"""Deeper check: torch full-precision vs each ONNX build, same prompts.
Runs 4 states / 10 questions: self-test (2), outlook (1), support (4), moderation (3).
Reports per-question top answer + worst per-option gap vs torch, plus size.
Usage: python compare_all.py [--models dir1 dir2 ...] [--cuda-torch]
"""
import argparse, json, sys, time
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

REQUESTS = [
    ("selftest",
     {"state": "Customer: my order arrived broken, I want my money back.",
      "questions": {"refund": {"type": "noul", "instructions": "Should the agent offer a refund?"},
                    "tone": {"type": "choice", "instructions": "Customer tone?",
                             "criteria": {"angry": "Angry", "neutral": "Neutral", "happy": "Happy"}}}}),
    ("outlook",
     {"state": "The quarterly revenue grew 12% year over year with margins stable.",
      "questions": {"q": {"type": "score", "instructions": "Rate the business outlook.",
                           "criteria": ["bad", "ok", "good", "great"]}}}),
    ("support", json.loads((HERE / "examples/support.json").read_text())),
    ("moderation", json.loads((HERE / "examples/moderation.json").read_text())),
]

def probs_of(answers):
    out = {}
    for qid, a in answers.items():
        if a["type"] == "noul":
            out[qid] = {"yes": a["noul"], "no": 1 - a["noul"]}
        elif a["type"] == "score":
            out[qid] = dict(a["probabilities"])
        else:
            out[qid] = dict(a["probabilities"])
    return out

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", nargs="+", default=["model-cuda-int4", "model-cpu"])
    ap.add_argument("--hf-model", default="/home/anon/LLM/KnowLine-4B-Gen5.5")
    ap.add_argument("--cuda-torch", action="store_true", help="load torch on CUDA (uses your GPU)")
    ap.add_argument("--out", default=None)
    a = ap.parse_args()

    from transformers import AutoTokenizer
    from knowline_layout import Layout
    import knowline_server as ks

    tok = AutoTokenizer.from_pretrained(str(HERE / "orig_tokenizer"), local_files_only=True)
    layout = Layout(read_twice=True, twice_max_tokens=3000)

    print("Loading torch full-precision ...", flush=True)
    t0 = time.time()
    device = "cuda" if a.cuda_torch else "cpu"
    hf = ks.HF(a.hf_model, device=device)
    eng_hf = ks.KnowLine(tok, hf, 1.0, 4, {}, layout)
    print(f"torch loaded in {time.time()-t0:.1f}s on {device}", flush=True)

    # torch answers once (reference) + timed second pass (first pass is warmup)
    ref = {}
    torch_times = {}
    for name, body in REQUESTS:
        ans, usage = eng_hf.run(body["state"], body["questions"])
        ref[name] = (ans, usage)
        print(f"torch {name}: usage={usage}", flush=True)
    for name, body in REQUESTS:
        t1 = time.time()
        ans, _ = eng_hf.run(body["state"], body["questions"])
        torch_times[name] = round(time.time() - t1, 2)
        ref[name] = (ans, ref[name][1])
    print(f"torch times: {torch_times}", flush=True)

    from onnx_raw import OnnxLM
    from score_knowline import OnnxScorer

    report = {"torch_device": device, "models": {}, "torch_times": torch_times}
    for mdir in a.models:
        mp = HERE / mdir
        size_gb = (mp / "model.onnx.data").stat().st_size / 1e9 if (mp / "model.onnx.data").exists() else -1
        is_cuda = "cpu" not in mdir  # model-cpu* -> CPU, rest CUDA
        providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if is_cuda else ["CPUExecutionProvider"]
        print(f"\n=== {mdir} ({size_gb:.2f} GB, {'CUDA' if is_cuda else 'CPU'}) ===")
        lm = OnnxLM(str(mp / "model.onnx"), providers=providers)
        eng = ks.KnowLine(tok, OnnxScorer(lm, tok), 1.0, 4, {}, layout)
        for name, body in REQUESTS:  # warmup, untimed
            eng.run(body["state"], body["questions"])
        mrep = {"size_gb": round(size_gb, 2), "questions": {}, "times": {}}
        worst = 0.0
        agree = 0
        total = 0
        t_req = 0.0
        for name, body in REQUESTS:
            t1 = time.time()
            ans_ox, _ = eng.run(body["state"], body["questions"])
            dt = time.time() - t1
            t_req += dt
            mrep["times"][name] = round(dt, 2)
            ans_hf, _ = ref[name]
            ph, po = probs_of(ans_hf), probs_of(ans_ox)
            for qid in ph:
                total += 1
                top_h = max(ph[qid], key=ph[qid].__getitem__)
                top_o = max(po[qid], key=po[qid].__getitem__)
                ok = top_h == top_o
                agree += ok
                dd = {k: abs(ph[qid][k] - po[qid][k]) for k in ph[qid]}
                w = max(dd.values())
                worst = max(worst, w)
                mrep["questions"][f"{name}/{qid}"] = {
                    "torch_top": top_h, "onnx_top": top_o, "agree": ok,
                    "worst_gap": round(w, 4),
                    "torch": {k: round(v, 4) for k, v in ph[qid].items()},
                    "onnx": {k: round(v, 4) for k, v in po[qid].items()},
                }
                mark = "" if ok else "  <-- FLIP"
                print(f"  {name}/{qid}: torch={top_h} onnx={top_o}{mark} worst_gap={w:.4f}")
        mrep["worst_gap"] = round(worst, 4)
        mrep["top1"] = f"{agree}/{total}"
        mrep["seconds_4_requests"] = round(t_req, 1)
        print(f"  WORST gap={worst:.4f} top-1 {agree}/{total} time={t_req:.1f}s")
        report["models"][mdir] = mrep
        del lm, eng

    if a.out:
        Path(a.out).write_text(json.dumps(report, indent=2))
        print(f"\nwrote {a.out}")
    else:
        print(json.dumps(report, indent=2))

if __name__ == "__main__":
    main()