"""Compare two validate.py result dumps (baseline vs converted).

Usage: python3 compare.py baseline.json converted.json
"""

import json
import sys

a = json.load(open(sys.argv[1]))
b = json.load(open(sys.argv[2]))

print(f"baseline={a['model']}  converted={b['model']}")
print(f"arith 17*23: base={a['arith_correct']} conv={b['arith_correct']}")
print(
    f"long gen: base {a['long_gen_tokens']} tok degenerate={a['long_gen_degenerate']} | "
    f"conv {b['long_gen_tokens']} tok degenerate={b['long_gen_degenerate']}"
)

# greedy 50-token comparison: positionwise top-1 agreement until first
# divergence (after divergence contexts differ), plus |dlogprob| on agreeing
# positions.
tot_pos = agree = 0
dl = []
full_match = 0
for ga, gb in zip(a["greedy"], b["greedy"]):
    ta, tb = ga["tokens"], gb["tokens"]
    la, lb = ga["token_logprobs"], gb["token_logprobs"]
    n = min(len(ta), len(tb))
    diverged = False
    prompt_match = 0
    for i in range(n):
        if diverged:
            break
        tot_pos += 1
        if ta[i] == tb[i]:
            agree += 1
            prompt_match += 1
            if la[i] is not None and lb[i] is not None:
                dl.append(abs(la[i] - lb[i]))
        else:
            diverged = True
    if not diverged and len(ta) == len(tb):
        full_match += 1
print(
    f"greedy: {full_match}/{len(a['greedy'])} prompts identical for all 50 tok; "
    f"agreement until divergence {agree}/{tot_pos} = {100 * agree / max(tot_pos, 1):.2f}%"
)
if dl:
    print(
        f"greedy |dlogprob| on agreeing tokens: mean {sum(dl) / len(dl):.5f} "
        f"max {max(dl):.5f} (n={len(dl)})"
    )

# teacher-forced prompt logprobs: identical context by construction.
tf_dl = []
tf_pos = tf_agree = 0
for sa, sb in zip(a["scored"], b["scored"]):
    for pa, pb in zip(sa["prompt_logprobs"], sb["prompt_logprobs"]):
        if pa is None or pb is None:
            continue
        # chosen prompt token = the key with rank field that matches; both dicts
        # contain the actual token (rank r) and the top-1. Find common token id
        # present in both (the prompt token id is the same on both sides).
        common = set(pa) & set(pb)
        if not common:
            continue
        # the prompt token appears in both dicts (it is always included)
        for tid in common:
            ra, rb = pa[tid].get("rank"), pb[tid].get("rank")
            tf_dl.append(abs(pa[tid]["logprob"] - pb[tid]["logprob"]))
            tf_pos += 1
            if (ra == 1) == (rb == 1):
                tf_agree += 1
            break
print(
    f"teacher-forced: n={tf_pos} mean |dlogprob| {sum(tf_dl) / max(len(tf_dl), 1):.5f} "
    f"max {max(tf_dl) if tf_dl else 0:.5f}; top-1-status agreement "
    f"{100 * tf_agree / max(tf_pos, 1):.2f}%"
)