"""Diff per-layer hidden states between DeepInfer and the reference implementation.

Both dumps keep only the final position of the prompt, reduced over the `hc_mult` copies, so
the comparison is one vector per layer. Reports cosine similarity, relative L2 error and the
ratio of norms, which together distinguish "slightly off" from "different direction".

    python scripts/compare_layers.py --ref ref_layers.npz --mine deepinfer_layers.npz
"""

import argparse

import numpy as np


def rel_err(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.linalg.norm(a - b) / (np.linalg.norm(b) + 1e-12))


def cosine(a: np.ndarray, b: np.ndarray) -> float:
    na, nb = np.linalg.norm(a), np.linalg.norm(b)
    if na == 0 or nb == 0:
        return float("nan")
    return float(np.dot(a, b) / (na * nb))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--ref", required=True)
    parser.add_argument("--mine", required=True)
    parser.add_argument("--tol", type=float, default=1e-2, help="first layer above this rel_err is the culprit")
    args = parser.parse_args()

    ref = np.load(args.ref)
    mine = np.load(args.mine)

    layers = sorted(int(k[5:]) for k in ref.files if k.startswith("layer") and k[5:].isdigit())
    mine_layers = {int(k[5:]) for k in mine.files if k.startswith("layer") and k[5:].isdigit()}
    missing = [i for i in layers if i not in mine_layers]
    if missing:
        print(f"WARNING: no dump from DeepInfer for layers {missing}")

    for key in ("embed", "layer0_attn", "layer1_attn"):
        if key in ref.files and key in mine.files:
            r, m = ref[key].astype(np.float64), mine[key].astype(np.float64)
            if r.shape == m.shape:
                print(
                    f"{key:>12}  cos={cosine(m, r):9.5f}  rel_err={rel_err(m, r):9.2e}"
                    f"  |mine|/|ref|={np.linalg.norm(m) / (np.linalg.norm(r) + 1e-12):8.4f}"
                )
    print()
    print(f"{'layer':>5}  {'cos':>9}  {'rel_err':>9}  {'|mine|/|ref|':>12}")
    first_bad = None
    for i in layers:
        if i not in mine_layers:
            continue
        r, m = ref[f"layer{i}"].astype(np.float64), mine[f"layer{i}"].astype(np.float64)
        if r.shape != m.shape:
            print(f"{i:>5}  shape mismatch ref={r.shape} mine={m.shape}")
            continue
        err, cos = rel_err(m, r), cosine(m, r)
        ratio = float(np.linalg.norm(m) / (np.linalg.norm(r) + 1e-12))
        flag = ""
        if first_bad is None and err > args.tol:
            first_bad = i
            flag = "   <-- first > tol"
        print(f"{i:>5}  {cos:9.5f}  {err:9.2e}  {ratio:12.4f}{flag}")

    ref_tok = int(ref["next_token"][0]) if "next_token" in ref.files else None
    print(f"\nreference next token: {ref_tok}")

    if "final_row_logits" in ref.files:
        rl = ref["final_row_logits"].astype(np.float64)
        top = np.argsort(-rl)[:5]
        print("reference top-5 tokens:", [int(x) for x in top], "logits:", [round(float(rl[x]), 3) for x in top])
        # DeepInfer's logits are not dumped; derive its choice from the harness output instead.

    print()
    if first_bad is None:
        print(f"RESULT: every dumped layer is within rel_err <= {args.tol}; divergence is downstream")
    else:
        print(f"RESULT: first layer above rel_err {args.tol} is layer {first_bad}")


if __name__ == "__main__":
    main()