"""Compare a candidate run against a reference golden (both in the reference_golden.py .pt format).

    python -I examples/neuron/compare_golden.py ~/di/golden/medium.pt <candidate>.pt [--json out.json]

The candidate must have been run TEACHER-FORCED on the golden's tokens (prompt_ids, then golden `tokens[:-1]` as
decode inputs); otherwise decode-step logits are compared across different contexts. Required keys in the
candidate: prompt_ids, logits [steps+1, vocab]; optional: fed_tokens, hidden (list per forward call of
[n_layers, seqlen, hc_mult, dim]), indexer_topk.

Metrics per forward call (step 0 = prefill, which predicts the first token):
  argmax match, rel = ||a-b|| / ||a|| of the fp32 logits, cosine, max |a-b|, top-5 overlap, KL(p_ref || p_cand),
  and the golden top1-top2 margin (an argmax mismatch where the margin is below the noise floor is not a bug).
Per layer (if both have `hidden`): relative L2 error of the block output over all tokens, and of the worst token.

The reference is not bit-stable even against itself: FP8 activation quantization turns sub-ulp differences
(e.g. a different CPU thread count) into ~1e-2 relative per-layer differences, and top-k routing / index
selection compound them. golden/noise_floor_*.json records the reference-vs-reference spread; judge a
candidate against that floor, not against zero.
"""

from __future__ import annotations

import argparse
import json

import torch


def rel(a, b):
    a, b = a.float(), b.float()
    return ((a - b).norm() / a.norm().clamp_min(1e-30)).item()


def compare(ref: dict, cand: dict) -> dict:
    assert list(ref["prompt_ids"]) == list(cand["prompt_ids"]), "different prompts"
    if "fed_tokens" in cand:
        exp = list(ref["tokens"][: len(cand["fed_tokens"])])
        if list(cand["fed_tokens"]) != exp:
            print("WARNING: candidate was not teacher-forced on the golden tokens; decode steps are not comparable")
    la, lb = ref["logits"].float(), cand["logits"].float()
    n = min(la.size(0), lb.size(0))
    steps = []
    for i in range(n):
        a, b = la[i], lb[i]
        top2 = a.topk(2).values
        pa, lpb = a.log_softmax(-1), b.log_softmax(-1)
        steps.append(
            dict(
                step=i,
                argmax_ref=int(a.argmax()),
                argmax_cand=int(b.argmax()),
                match=bool(a.argmax() == b.argmax()),
                margin_ref=round((top2[0] - top2[1]).item(), 5),
                rel=rel(a, b),
                cos=torch.nn.functional.cosine_similarity(a, b, dim=0).item(),
                maxabs=(a - b).abs().max().item(),
                top5_overlap=len(set(a.topk(5).indices.tolist()) & set(b.topk(5).indices.tolist())),
                kl=(pa.exp() * (pa - lpb)).sum().item(),
            )
        )
    layers = []
    if ref.get("hidden") and cand.get("hidden"):
        for i in range(min(len(ref["hidden"]), len(cand["hidden"]))):
            ha, hb = ref["hidden"][i].float(), cand["hidden"][i].float()
            if ha.shape != hb.shape:
                layers.append(dict(step=i, error=f"shape {tuple(ha.shape)} vs {tuple(hb.shape)}"))
                continue
            dims = tuple(range(2, ha.dim()))
            per_tok = (ha - hb).pow(2).sum(dims).sqrt() / ha.pow(2).sum(dims).sqrt().clamp_min(1e-30)  # [L, s]
            layers.append(
                dict(
                    step=i,
                    rel=[rel(ha[l], hb[l]) for l in range(ha.size(0))],
                    worst_token_rel=per_tok.max(-1).values.tolist(),
                )
            )
    return dict(steps=steps, layers=layers)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("ref")
    ap.add_argument("cand")
    ap.add_argument("--json", default=None)
    args = ap.parse_args()
    ref = torch.load(args.ref, weights_only=False)
    cand = torch.load(args.cand, weights_only=False)
    res = compare(ref, cand)
    print(f"{'step':>4} {'match':>5} {'margin':>7} {'rel':>9} {'cos':>9} {'maxabs':>8} {'top5':>4} {'KL':>9}")
    for s in res["steps"]:
        print(
            f"{s['step']:>4} {str(s['match']):>5} {s['margin_ref']:>7.3f} {s['rel']:>9.2e} {s['cos']:>9.6f} "
            f"{s['maxabs']:>8.3f} {s['top5_overlap']:>4} {s['kl']:>9.2e}"
        )
    for L in res["layers"]:
        if "error" in L:
            print(f"hidden step {L['step']}: {L['error']}")
        else:
            print(f"hidden step {L['step']}: layer rel " + " ".join(f"{v:.1e}" for v in L["rel"]))
    if args.json:
        with open(args.json, "w") as fh:
            json.dump(res, fh, indent=1)


if __name__ == "__main__":
    main()