"""CharacterEval V07a, checked on our side: the operator's numerical-sensitivity result against the repository and the anchor.

    python3 -I tools/experiments/charactereval_rm_numerics_verify.py <CharacterEval repo at c3d44a6> <result.json> <anchor scores dir> [out.json]

  N1  the 1,200 inputs are the ones charactereval_rm_numerics_check.py selects from the repository (first 20 per metric per file)
  N2  the root-mean-square of (variant - bf16), recomputed from the scores, and the verdict under the criteria fixed in advance
  N3  the operator's bf16 scores against the platform's anchor scores of the same inputs (scores_<name>.jsonl): equal within 1e-6?
  N4  bf16-strict - bf16 by metric, and the size of the changes (median and 90th percentile of |difference|)
Reads only; exit status 0 when N1 holds and N3 is equal for every input, 1 otherwise.
"""
import collections
import json
import math
import os
import statistics as st
import sys

FROZEN = ["chatglm3_6b", "baichuan_7b", "xverse_7b", "internlm_7b", "qwen_7b"]
PER_METRIC, SUFFICIENT, RULED_OUT, EXACT = 20, 0.06, 0.02, 1e-6


def rms(x):
    return math.sqrt(sum(t * t for t in x) / len(x))


def main():
    repo, result, anchor = sys.argv[1:4]
    out = sys.argv[4] if len(sys.argv) > 4 else None
    d = json.load(open(result))
    expected = []
    for name in FROZEN:
        records = json.load(open(os.path.join(repo, "results", f"generation_trans_{name}.jsonl"), encoding="utf-8"))
        taken = collections.Counter()
        for i, r in enumerate(records):
            if taken[r["metric_en"]] < PER_METRIC:
                taken[r["metric_en"]] += 1
                expected.append({"file": name, "row": i, "metric_en": r["metric_en"]})
    r = {"N1": {"inputs": len(d["inputs"]), "equal_to_selection": expected == d["inputs"]}, "N2": {}}
    base = d["variants"]["bf16"]["scores"]
    for v in ["bf16-strict", "fp32"]:
        diff = [b - a for a, b in zip(base, d["variants"][v]["scores"])]
        x = rms(diff)
        r["N2"][v] = {"rms": round(x, 4), "reported_rms": round(d["comparison"][v]["rms"], 4), "mean_signed": round(st.mean(diff), 4),
                      "changed": sum(t != 0 for t in diff), "median_abs": round(st.median(map(abs, diff)), 4),
                      "p90_abs": round(sorted(map(abs, diff))[int(0.9 * len(diff))], 4), "max_abs": round(max(map(abs, diff)), 4),
                      "verdict": "sufficient" if x >= SUFFICIENT else "ruled out" if x < RULED_OUT else "partial"}
    platform = {}
    for name in FROZEN:
        platform[name] = {}
        for line in open(os.path.join(anchor, f"scores_{name}.jsonl")):
            x = json.loads(line)
            platform[name][x["i"]] = x["score"]
    gap = [abs(platform[it["file"]][it["row"]] - s) for it, s in zip(d["inputs"], base)]
    r["N3"] = {"n": len(gap), "equal": sum(g <= EXACT for g in gap), "max_abs_diff": max(gap)}
    by = collections.defaultdict(list)
    for it, a, b in zip(d["inputs"], base, d["variants"]["bf16-strict"]["scores"]):
        by[it["metric_en"]].append(b - a)
    r["N4"] = {"bf16-strict_rms_by_metric": {m: round(rms(x), 3) for m, x in sorted(by.items())}}
    r["environment"] = d["environment"]
    print(f"N1 inputs {r['N1']['inputs']}, equal to the selection: {r['N1']['equal_to_selection']}")
    for v, c in r["N2"].items():
        print(f"N2 {v} - bf16: rms {c['rms']} (reported {c['reported_rms']}), mean {c['mean_signed']:+}, median |d| {c['median_abs']}, "
              f"p90 |d| {c['p90_abs']}, max |d| {c['max_abs']}, changed {c['changed']}/{len(base)} -> {c['verdict']}")
    print(f"N3 operator bf16 = platform: {r['N3']['equal']}/{r['N3']['n']}, max |diff| {r['N3']['max_abs_diff']}")
    print(f"N4 bf16-strict rms by metric: {r['N4']['bf16-strict_rms_by_metric']}")
    if out:
        json.dump(r, open(out, "w"), indent=1)
    sys.exit(0 if r["N1"]["equal_to_selection"] and r["N3"]["equal"] == r["N3"]["n"] else 1)


if __name__ == "__main__":
    main()