"""CharacterBench: recompute every audited run's scores from the published frozen judge outputs, with the authors' own
reading and aggregation, and compare them with the scores the platform published. Needs only the audit bundle and the
authors' repository; nothing of the platform is imported.

    git clone https://github.com/thu-coai/CharacterBench repo && git -C repo checkout 66c113c
    python3 -I characterbench_recompute.py repo <bundle>/frozen

Per run: each sample's ten outputs read by all_eval.py's check() (executed from the repository's source), averaged; per
language and test file, mean / max(human scores of the judged samples) × 5 (all_eval.py); AVG = mean of the 13 columns.
Compared with result.perLang in runs.json at two decimals (the platform publishes two). Exit status 1 on any difference.
"""
import ast
import collections
import json
import math
import os
import statistics
import sys

KEYS = {"memory_consistency": "mc", "fact_accuracy": "fa", "boundary_consistency": "bck", "attribute_consistency_bot": "acb",
        "attribute_consistency_human": "ach", "behavior_consistency_bot": "bcpb", "behavior_consistency_human": "bcph",
        "emotion_self_regulation": "es", "empathetic_responsiveness": "er", "morality_stability": "ms", "morality_robustness": "mr",
        "human_likeness": "hl", "engagement": "eg"}


def authors_check(repo):
    src = open(os.path.join(repo, "all_eval.py"), encoding="utf-8").read()
    node = next(n for n in ast.parse(src).body if isinstance(n, ast.FunctionDef) and n.name == "check")
    ns = {}
    exec(ast.get_source_segment(src, node), ns)
    return ns["check"]


def main(repo, frozen):
    check = authors_check(repo)
    raw = {f: json.load(open(os.path.join(repo, "eval_data/raw_data", f"{f}_test.json"), encoding="utf-8")) for f in KEYS}
    runs = json.load(open(os.path.join(frozen, "runs.json"), encoding="utf-8"))
    half_up = lambda x: math.floor(x * 100 + 0.5 + 1e-9) / 100
    bad = 0
    for run in runs:
        scores = {l: collections.defaultdict(list) for l in ("zh", "en")}
        golds = {l: collections.defaultdict(list) for l in ("zh", "en")}
        for line in open(os.path.join(frozen, f"judge_outputs_{run['id']}.jsonl"), encoding="utf-8"):
            r = json.loads(line)
            if "outputs" not in r:
                continue  # excluded sample: no score
            lang, file, index = r["key"].split("|")
            vals = [float(check(t)) for t in r["outputs"]]
            scores[lang][file].append(sum(vals) / len(vals))
            golds[lang][file].append(float(check(str(raw[file][int(index)]["annotation_score"]))))
        for lang in ("zh", "en"):
            if not scores[lang]:
                continue
            cols = {KEYS[f]: statistics.fmean(v) / max(golds[lang][f]) * 5 for f, v in scores[lang].items()}
            cols["avg"] = statistics.fmean(cols[k] for k in KEYS.values())
            published = {**run["result"]["perLang"][lang]["dims"], "avg": run["result"]["perLang"][lang]["overall"]}
            diff = [k for k in cols if abs(half_up(cols[k]) - published[k]) > 1e-9]
            bad += len(diff)
            print(f"{run['id']} {run['model']} {lang}: AVG {cols['avg']:.4f} (published {published['avg']}); "
                  f"{14 - len(diff)}/14 values equal at two decimals{'; differ: ' + ', '.join(diff) if diff else ''}")
    print("all equal" if not bad else f"{bad} values differ")
    sys.exit(1 if bad else 0)


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