#!/usr/bin/env python3
"""tau_null.py — the convergence ceiling for the descent.

Two INDEPENDENT sampled answers to the SAME question, lensed identically: their cosine is the
best any 'convergence' between consecutive rungs could hope for. If the descent's tau exceeds
this ceiling, 'no trajectory converged' describes the metric, not the model.
"""
from __future__ import annotations
import argparse, importlib.util, json, os, sys
os.environ.setdefault("HF_HOME", "/mnt/assets/models/huggingface")
os.environ.setdefault("HF_HUB_OFFLINE", "0")
_here = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location("jlens", os.path.join(_here, "jlens_jacobian.py"))
jl = importlib.util.module_from_spec(spec); spec.loader.exec_module(jl)
import torch, math

def cosine(a, b):
    ks = set(a) | set(b)
    dot = sum(a.get(k, 0) * b.get(k, 0) for k in ks)
    na = math.sqrt(sum(v*v for v in a.values())); nb = math.sqrt(sum(v*v for v in b.values()))
    return dot / (na * nb) if na and nb else 0.0

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--device", default="cuda:0"); ap.add_argument("--trust", action="store_true")
    a = ap.parse_args()
    from transformers import AutoModelForCausalLM, AutoTokenizer
    tok = AutoTokenizer.from_pretrained(a.model)
    model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.bfloat16,
        device_map={"": a.device}, attn_implementation="eager", trust_remote_code=a.trust).eval()
    layers, _, _ = jl.parts(model)
    sel = sorted({int(len(layers) * f) for f in (0.4, 0.55, 0.7, 0.85)})
    QS = ["Whisky or vodka — which would you choose, and why?",
          "Tradition or progress — which matters more, and why?",
          "Mercy or justice — which matters more, and why?",
          "Tea or coffee — which would you choose, and why?",
          "Safety or freedom — which should be weighted more, and why?",
          "How does compound interest work?",
          "A mechanical watch or a smartwatch — which, and why?",
          "Loyalty or honesty — when they conflict, which wins, and why?",
          "How do I recycle electronic waste?",
          "Stability or change — which serves people better, and why?"]
    out = []
    for q in QS:
        masses = []
        for i in range(2):
            enc = jl.apply_template(tok, [{"role": "user", "content": q}],
                                    add_generation_prompt=True, return_tensors="pt", return_dict=True)
            ids = enc["input_ids"].to(model.device)
            with torch.no_grad():
                gen = model.generate(ids, max_new_tokens=700, do_sample=True, temperature=0.7,
                                     top_p=0.95, pad_token_id=(tok.pad_token_id or tok.eos_token_id))
            masses.append(jl.jacobian_lens_mass(model, tok, gen, ids.shape[1], sel))
        c = cosine(*masses)
        out.append({"question": q, "cos": c})
        print(f"  cos={c:.3f}  {q[:50]}", file=sys.stderr, flush=True)
    import statistics as st
    ceil = [o["cos"] for o in out]
    summary = {"model": a.model, "pairs": out,
               "mean": st.mean(ceil), "median": st.median(ceil),
               "min": min(ceil), "max": max(ceil)}
    json.dump(summary, open(a.out, "w"), indent=1)
    print(f"[tau-null] mean={st.mean(ceil):.3f} median={st.median(ceil):.3f} — any tau above "
          f"this describes the metric, not the model", file=sys.stderr)

if __name__ == "__main__":
    main()