"""Diff ONNX-bidi slot logits / answers vs the torch noncausal reference.

Usage:
    python diff_bidi.py --model ./model --cuda            # fp16 bidi
    python diff_bidi.py --model ./model-cuda-int4 --cuda  # int4 bidi

Compares per-branch slot LOGITS (torch-ref.json diag) and final answer
probabilities. Needs torch-ref.json (patched qev engine, debug=True).
"""

import argparse
import json
import sys
from pathlib import Path

import numpy as np

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

from onnx_raw import OnnxLM  # noqa: E402
from score_elephant import OnnxScorer, decide, plan, verify_slots  # noqa: E402


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--model", default=str(HERE / "model"))
    p.add_argument("--onnx", default="model.onnx")
    p.add_argument("--cuda", action="store_true")
    p.add_argument("--ref", default=str(HERE / "torch-ref.json"))
    a = p.parse_args()

    from transformers import AutoTokenizer

    tok = AutoTokenizer.from_pretrained(str(HERE / "orig_tokenizer"), local_files_only=True)
    slot_ids = verify_slots(tok)
    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if a.cuda else ["CPUExecutionProvider"]
    lm = OnnxLM(str(Path(a.model) / a.onnx), providers=providers)
    ref = json.load(open(a.ref))

    worst_logit, worst_prob, agree = 0.0, 0.0, True
    for name in ["support", "moderation"]:
        body = json.load(open(HERE / "examples" / f"{name}.json"))
        full, rendered, owners, total, plen = plan(body["state"], body["questions"], tok, 1)
        scorer = OnnxScorer(lm, slot_ids, plen)
        print(f"--- {name} ---")
        for (br_ids, br) in zip(full, rendered):
            ologits, _ = scorer.score(br_ids, br.n_slots)
            tlogits = ref[name]["diag"][br.qid]["slot_logits"][0]
            d = max(abs(x - y) for x, y in zip(ologits, tlogits))
            worst_logit = max(worst_logit, d)
            print(f"  {br.qid}: max|Δlogit|={d:.4f} (k={br.n_slots})")
        ans_ox, _ = decide(body["state"], body["questions"], tok,
                           lambda pl: OnnxScorer(lm, slot_ids, pl), 1.0, 1, 4)
        for qid, ta in ref[name]["answers"].items():
            oa = ans_ox[qid]
            tp = ta["noul"] if ta["type"] == "noul" else None
            if ta["type"] == "noul":
                d = abs(ta["noul"] - oa["noul"])
                worst_prob = max(worst_prob, d)
                if (ta["noul"] > 0.5) != (oa["noul"] > 0.5):
                    agree = False
                print(f"  {qid}: torch={ta['noul']:.4f} onnx={oa['noul']:.4f} |Δ|={d:.4f}")
            else:
                tk = {k: v for k, v in ta["probabilities"].items()}
                ok = oa["probabilities"]
                top_t, top_o = max(tk, key=tk.__getitem__), max(ok, key=ok.__getitem__)
                if top_t != top_o:
                    agree = False
                for opt in tk:
                    d = abs(tk[opt] - ok[opt])
                    worst_prob = max(worst_prob, d)
                print(f"  {qid}: top torch={top_t} onnx={top_o}; "
                      + " ".join(f"{o}:Δ={abs(tk[o]-ok[o]):.4f}" for o in tk))
    print(f"WORST |Δlogit|={worst_logit:.4f} WORST |Δprob|={worst_prob:.4f} "
          f"top-1: {'PASS' if agree else 'FAIL'}")


if __name__ == "__main__":
    main()