"""CharacterEval: mechanical checks of the released files (reproduction report, Sections 2 and 5.1).

    python3 -I tools/experiments/charactereval_release_checks.py <CharacterEval repo at c3d44a6> <e-print dir> [out.json]

R1  file formats: each data and result file is one JSON document (an object for id2metric, arrays for the others), not JSON lines
R2  the reward-model training data (rm_train_data.json): rows, distinct ids, labels; whether every (id, role, context) is a test
    example; whether its (id, metric) pairs are a subset of id2metric
R3  the five frozen result files: rows, ids, failed or empty outputs
R4  dimension averages of the paper's Table 4 weighted by the number of examples per metric, against the printed averages
    (the unweighted reading is checked by charactereval_table_checks.py)
Prints one line per check and writes them as JSON. Reads only; exit status 0.
"""
import collections
import json
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import charactereval_table_checks as tc  # noqa: E402

AVGS = {"cc_avg": ["Exposure", "Accuracy", "Hallucination", "Behavior", "Utterance"], "ca_avg": ["Fluency", "Coherence", "Consistency"],
        "ra_avg": ["Humanlikeness", "Communication_skills", "Diversity", "Empathy"]}
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"}
FROZEN = ["chatglm3_6b", "baichuan_7b", "xverse_7b", "internlm_7b", "qwen_7b"]


def load(repo, rel):
    return json.load(open(os.path.join(repo, rel), encoding="utf-8"))


def main():
    repo, eprint = sys.argv[1:3]
    out = sys.argv[3] if len(sys.argv) > 3 else None
    r = {}
    files = ["data/test_data.jsonl", "data/id2metric.jsonl", "data/character_profiles.json", "rm_train_data.json"] + [f"results/generation_trans_{m}.jsonl" for m in FROZEN]
    r["R1"] = {f: type(load(repo, f)).__name__ for f in files}

    test = {str(x["id"]): x for x in load(repo, "data/test_data.jsonl")}
    id2m = load(repo, "data/id2metric.jsonl")
    pairs = {(i, m[0]) for i, ms in id2m.items() for m in ms}
    rm = load(repo, "rm_train_data.json")
    ids = {str(x["id"]) for x in rm}
    same = sum(1 for x in rm if str(x["id"]) in test and test[str(x["id"])]["role"] == x["role"] and test[str(x["id"])]["context"] == x["context"])
    rm_pairs = {(str(x["id"]), x["metric_en"]) for x in rm}
    labels = collections.Counter(str(x[x["metric_en"]]) for x in rm)
    r["R2"] = {"rows": len(rm), "distinct_ids": len(ids), "ids_in_test_set": len(ids & set(test)), "rows_with_test_role_and_context": same,
               "pairs_subset_of_id2metric": rm_pairs <= pairs, "pairs": len(rm_pairs), "id2metric_pairs": len(pairs), "labels": dict(sorted(labels.items()))}

    frozen = {}
    for m in FROZEN:
        d = load(repo, f"results/generation_trans_{m}.jsonl")
        frozen[m] = {"rows": len(d), "ids": len({x["id"] for x in d}), "none_or_error": sum(1 for x in d if x["model_output"] in (None, "ERROR")),
                     "empty": sum(1 for x in d if x["model_output"] == "")}
    r["R3"] = frozen

    counts = collections.Counter(m[0] for ms in id2m.items() for m in ms[1])
    raw = open(os.path.join(eprint, "acl_latex.tex"), encoding="utf-8").read()
    t4 = tc.tables(tc.live_lines(raw), "tab:detail_result")
    worst = 0.0
    for v, _ in t4.values():
        for avg, ms in AVGS.items():
            w = sum(v[KEY[m]] * counts[m] for m in ms) / sum(counts[m] for m in ms)
            worst = max(worst, abs(w - v[avg]))
    r["R4"] = {"examples_per_metric": dict(counts), "max_abs_diff_weighted_vs_printed": round(worst, 4)}

    print(f"R1 formats: {r['R1']}")
    print(f"R2 rm_train_data: {r['R2']['rows']} rows, {r['R2']['distinct_ids']} ids, {r['R2']['ids_in_test_set']} of them test ids, "
          f"{r['R2']['rows_with_test_role_and_context']} rows with a test example's role and context; (id, metric) pairs {r['R2']['pairs']} "
          f"of id2metric's {r['R2']['id2metric_pairs']}, subset: {r['R2']['pairs_subset_of_id2metric']}; labels {r['R2']['labels']}")
    print(f"R3 frozen outputs: {frozen}")
    print(f"R4 example-weighted dimension averages vs printed: max |diff| {r['R4']['max_abs_diff_weighted_vs_printed']}")
    if out:
        json.dump(r, open(out, "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()