"""CharacterBench audit, paper anchor: platform runs of a model the paper reports, against the paper's row.

    python3 -I tools/experiments/characterbench_anchor.py EPRINT OUT.json [--judge JUDGE_AUDIT.json] RUN_1.json [RUN_2.json ...]

EPRINT: the paper's e-print directory (the row is read from its LaTeX source by characterbench_paper.py's parser);
RUN_*.json: D1 exports of the runs (runs table row, with result_json). Prints, per language and column, the paper's
value, each run's value and their mean, the difference, and the spread between repeated runs. With --judge (the output of
characterbench_judge_audit.py), also how far the column differences follow the deployed judge's own systematic
differences from the authors' released predictions on the annotated responses (Pearson over the 13 columns).
"""
import json
import os
import re
import statistics
import sys

COLS = ["mc", "fa", "bck", "acb", "ach", "bcpb", "bcph", "es", "er", "ms", "mr", "hl", "eg"]
FILES = ["memory_consistency", "fact_accuracy", "boundary_consistency", "attribute_consistency_bot", "attribute_consistency_human",
         "behavior_consistency_bot", "behavior_consistency_human", "emotion_self_regulation", "empathetic_responsiveness",
         "morality_stability", "morality_robustness", "human_likeness", "engagement"]


def pearson(x, y):
    mx, my = statistics.fmean(x), statistics.fmean(y)
    return sum((a - mx) * (b - my) for a, b in zip(x, y)) / (sum((a - mx) ** 2 for a in x) * sum((b - my) ** 2 for b in y)) ** 0.5


ROW = "GPT-4o"


def paper_row(eprint):
    """The appendix table's GPT-4o row: AVG then the 13 columns, each "zh/en"."""
    tex = open(os.path.join(eprint, "appendix.tex"), encoding="utf-8").read()
    m = re.search(r"^\s*" + re.escape(ROW) + r"\s*&(.*?)\\\\", tex, re.M)
    cells = [c.strip() for c in m.group(1).split("&")]
    vals = [tuple(float(re.sub(r"\\\w+\{|\}", "", x)) for x in c.split("/")) for c in cells]
    return {lang: {"avg": vals[0][i], **{k: vals[j + 1][i] for j, k in enumerate(COLS)}} for i, lang in enumerate(("zh", "en"))}


def main(eprint, out, runs):
    judge = None
    if runs and runs[0] == "--judge":
        judge, runs = json.load(open(runs[1], encoding="utf-8")), runs[2:]
    paper = paper_row(eprint)
    results = []
    for path in runs:
        run = json.load(open(path, encoding="utf-8"))
        res = json.loads(run["result_json"])
        results.append({"id": run["id"], "model": run["target_model"], "config": json.loads(run["config_json"]).get("targetParams"),
                        "perLang": {l: {"avg": res["perLang"][l]["overall"], **res["perLang"][l]["dims"]} for l in ("zh", "en")},
                        "excluded": res.get("excludedSessions")})
    report = {"paper_row": ROW, "paper": paper, "runs": results, "compare": {}}
    for l in ("zh", "en"):
        rows = {}
        for k in ["avg"] + COLS:
            vs = [r["perLang"][l][k] for r in results]
            m = statistics.fmean(vs)
            rows[k] = {"paper": paper[l][k], "runs": vs, "mean": round(m, 3), "diff": round(m - paper[l][k], 3),
                       "repeat_range": round(max(vs) - min(vs), 3)}
        report["compare"][l] = rows
        if judge:
            bias = [judge[l]["vs_released_sc"]["mean_diff_5pt_per_column"][f] for f in FILES]
            report.setdefault("judge_bias_correlation", {})[l] = round(pearson([rows[k]["diff"] for k in COLS], bias), 2)
        d = [abs(rows[k]["diff"]) for k in COLS]
        print(f"{l}: AVG paper {paper[l]['avg']} platform {rows['avg']['runs']} (diff {rows['avg']['diff']:+.2f}); "
              f"columns |diff| mean {statistics.fmean(d):.3f} max {max(d):.2f}; within 0.10: {sum(x <= 0.10 for x in d)}/13; "
              f"repeat range AVG {rows['avg']['repeat_range']}; follows judge bias r = {report.get('judge_bias_correlation', {}).get(l)}")
    json.dump(report, open(out, "w"), indent=1, ensure_ascii=False)


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], sys.argv[3:])