"""CharacterEval anchor diagnostics (matrix V07; reproduction report, Section 5.3).

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

Reads the anchor result (<work>/anchor.json) and the cached scores (<work>/scores_<name>.jsonl) written by
charactereval_anchor.py, and reports:
  D1  rows per frozen file against id2metric: duplicates, missing and extra (id, metric) pairs
  D2  per metric cell: records n, standard deviation of the scores, difference to Table 4, z = diff / (sd / sqrt(n))
  D3  the per-score noise the differences imply: root-mean-square over the 60 metric cells of diff * sqrt(n)
  D4  pairwise orders: for each of the 15 columns and each pair of the five models, whether the platform keeps Table 4's order
  D5  the cells beyond 0.005, largest first
Reads only; exit status 0.
"""
import collections
import itertools
import json
import math
import os
import statistics as st
import sys

FROZEN = {"chatglm3_6b": "ChatGLM3-6B", "baichuan_7b": "Baichuan2-7B", "xverse_7b": "XVERSE-7B", "internlm_7b": "InternLM-7B", "qwen_7b": "Qwen-7B"}
KEY = {"Exposure": "ke", "Accuracy": "ka", "Hallucination": "kh", "Behavior": "pb", "Utterance": "pu", "Fluency": "flu", "Coherence": "coh",
       "Consistency": "cons", "Humanlikeness": "hl", "Communication_skills": "cs", "Diversity": "ed", "Empathy": "emp"}


def main():
    repo, work = sys.argv[1:3]
    out = sys.argv[3] if len(sys.argv) > 3 else None
    rows = json.load(open(os.path.join(work, "anchor.json")))["anchor"]["rows"]
    id2m = json.load(open(os.path.join(repo, "data", "id2metric.jsonl"), encoding="utf-8"))
    pairs = {(str(i), m[0]) for i, ms in id2m.items() for m in ms}
    r = {"D1": {}, "D2": {}}
    zs, noise = [], []
    for name, row in FROZEN.items():
        records = json.load(open(os.path.join(repo, "results", f"generation_trans_{name}.jsonl"), encoding="utf-8"))
        keys = collections.Counter((str(x["id"]), x["metric_en"]) for x in records)
        r["D1"][row] = {"rows": len(records), "duplicates": sum(c - 1 for c in keys.values() if c > 1),
                        "missing": len(pairs - set(keys)), "extra": len(set(keys) - pairs)}
        scores = {}
        for line in open(os.path.join(work, f"scores_{name}.jsonl")):
            x = json.loads(line)
            scores[x["i"]] = x["score"]
        by = collections.defaultdict(list)
        for i, x in enumerate(records):
            by[x["metric_en"]].append(scores[i])
        r["D2"][row] = {}
        for m, v in by.items():
            d, sd = rows[row]["cells"][KEY[m]]["diff"], st.stdev(v)
            z = d / (sd / math.sqrt(len(v)))
            zs.append(z)
            noise.append(d * d * len(v))
            r["D2"][row][KEY[m]] = {"n": len(v), "sd": round(sd, 4), "diff": round(d, 5), "z": round(z, 3)}
    r["D3"] = {"implied_per_score_rms": round(math.sqrt(sum(noise) / len(noise)), 4), "z_mean": round(st.mean(zs), 3),
               "z_sd": round(st.stdev(zs), 3), "z_max_abs": round(max(map(abs, zs)), 3)}
    flips, n = [], 0
    for c in next(iter(rows.values()))["cells"]:
        for a, b in itertools.combinations(rows, 2):
            n += 1
            pa, pb = rows[a]["cells"][c]["paper"], rows[b]["cells"][c]["paper"]
            qa, qb = rows[a]["cells"][c]["platform"], rows[b]["cells"][c]["platform"]
            if (pa - pb) * (qa - qb) <= 0:
                flips.append([c, a, b])
    r["D4"] = {"pairs": n, "not_kept": flips}
    r["D5"] = sorted(([row, k, round(c["diff"], 4)] for row, x in rows.items() for k, c in x["cells"].items() if abs(c["diff"]) > 0.005),
                     key=lambda x: -abs(x[2]))
    print(f"D1 rows: {r['D1']}")
    for row, cells in r["D2"].items():
        print(f"D2 {row}: " + " ".join(f"{k}:n{c['n']} d{c['diff']:+.4f} z{c['z']:+.2f}" for k, c in cells.items()))
    print(f"D3 implied per-score noise (rms) {r['D3']['implied_per_score_rms']}; z mean {r['D3']['z_mean']:+}, sd {r['D3']['z_sd']}, max |z| {r['D3']['z_max_abs']}")
    print(f"D4 pairwise orders kept: {n - len(flips)}/{n}")
    print(f"D5 beyond 0.005 ({len(r['D5'])}): {r['D5']}")
    if out:
        json.dump(r, open(out, "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()