"""CharacterEval anchor (matrix V07; reproduction report, Section 5.3), through the reward-model endpoint the platform calls.

1. Probes: the 26 probes of charactereval_rm_probes.json are scored through the platform's endpoint and compared with the
   reference scores the operator's run of the official run_char_rm.py produced (charactereval_rm_server_check.result.json):
   every |difference| must be <= 1e-6. This checks, from the platform's side, the endpoint the platform itself uses.
2. Anchor: the five frozen result files of the repository (results/generation_trans_*.jsonl, 8,032 rows each, already in the
   form run_char_rm.py reads) are turned into reward-model inputs by the official format_input (charactereval_reference.py),
   scored through the same endpoint, aggregated with the official compute_score and the platform's dimension averages
   (matrix E17, E18), and compared cell by cell with the paper's Table 4 (arXiv v2 LaTeX).

Criteria, fixed on 2026-10-28 before any measurement (matrix V07):
  reproduced       every one of the 75 cells (5 models x (12 metrics + 3 dimension averages)) within 0.005
  within one step  every cell within 0.02, some beyond 0.005
  not reproduced   any cell beyond 0.02
Also reported: the distribution of the differences, their mean sign, and whether the five models keep Table 4's order on
character consistency.

    RPB_SESSION_FILE=<file holding an admin session token> python3 -I tools/experiments/charactereval_anchor.py \\
        <CharacterEval repo at c3d44a6> <e-print dir> <work dir> [--api https://sociabench.com/api]

Scores are cached in the work dir (one JSON line per input), so an interrupted run resumes. The session token is read from
the file and never printed. Writes <work dir>/anchor.json and prints a summary.
"""
import argparse
import json
import os
import statistics as st
import sys
import time
import urllib.error
import urllib.request

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

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"}
AVG = {"cc": "cc_avg", "ca": "ca_avg", "ra": "ra_avg"}
BATCH = 64
EXACT, REPRODUCED, ONE_STEP = 1e-6, 0.005, 0.02


def score(api, token, texts):
    body = json.dumps({"texts": texts}).encode()
    for attempt in range(6):
        req = urllib.request.Request(f"{api}/settings/judge-endpoints/charrm/score", data=body, method="POST",
                                     headers={"content-type": "application/json", "cookie": f"rpb_session={token}", "origin": "https://sociabench.com", "user-agent": "charactereval-anchor"})
        try:
            with urllib.request.urlopen(req, timeout=600) as r:
                out = json.load(r)["scores"]
            assert len(out) == len(texts)
            return out
        except (urllib.error.URLError, TimeoutError, AssertionError) as e:
            if attempt == 5:
                raise
            print(f"  retry {attempt + 1}: {str(e)[:120]}", flush=True)
            time.sleep(5 * 2 ** attempt)


def scored(api, token, texts, cache):
    done = {}
    if os.path.exists(cache):
        for line in open(cache):
            x = json.loads(line)
            done[x["i"]] = x["score"]
    todo = [i for i in range(len(texts)) if i not in done]
    with open(cache, "a") as f:
        for k in range(0, len(todo), BATCH):
            idx = todo[k:k + BATCH]
            for i, s in zip(idx, score(api, token, [texts[i] for i in idx])):
                done[i] = s
                f.write(json.dumps({"i": i, "score": s}) + "\n")
            f.flush()
            print(f"  {os.path.basename(cache)}: {len(done)}/{len(texts)}", flush=True)
    return [done[i] for i in range(len(texts))]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("repo")
    ap.add_argument("eprint")
    ap.add_argument("work")
    ap.add_argument("--api", default="https://sociabench.com/api")
    args = ap.parse_args()
    token = open(os.environ["RPB_SESSION_FILE"]).read().strip()
    os.makedirs(args.work, exist_ok=True)
    result = {"criteria": {"probes": f"every |platform endpoint - official reference| <= {EXACT}", "reproduced": f"all 75 cells within {REPRODUCED}",
                           "within_one_step": f"all within {ONE_STEP}", "not_reproduced": f"any cell beyond {ONE_STEP}"}}

    # 1. Probes against the operator's reference.
    probes = json.load(open(os.path.join(HERE, "charactereval_rm_probes.json"), encoding="utf-8"))["probes"]
    reference = [p["score"] for p in json.load(open(os.path.join(HERE, "charactereval_rm_server_check.result.json")))["B_reference"]["probes"]]
    got = score(args.api, token, [p["text"] for p in probes])
    diffs = [abs(a - b) for a, b in zip(got, reference)]
    result["probes"] = {"n": len(got), "max_abs_diff": max(diffs), "equal": sum(d <= EXACT for d in diffs), "pass": max(diffs) <= EXACT, "scores": got}
    print(f"probes: {result['probes']['equal']}/{len(got)} equal to the official reference, max |diff| {max(diffs)}", flush=True)
    if not result["probes"]["pass"]:
        json.dump(result, open(os.path.join(args.work, "anchor.json"), "w"), indent=1)
        sys.exit("the platform's endpoint does not reproduce the official reference; anchor not run")

    # 2. Anchor.
    profiles = ref.load_profiles(args.repo)
    t4 = tc.tables(tc.live_lines(open(os.path.join(args.eprint, "acl_latex.tex"), encoding="utf-8").read()), "tab:detail_result")
    cells, rows = [], {}
    for name, row in FROZEN.items():
        records = json.load(open(os.path.join(args.repo, "results", f"generation_trans_{name}.jsonl"), encoding="utf-8"))
        texts = ref.rm_inputs(records, profiles)
        scores = scored(args.api, token, texts, os.path.join(args.work, f"scores_{name}.jsonl"))
        for r, s in zip(records, scores):
            r[r["metric_en"]] = s
        means = ref.compute_score(records)
        dims = ref.dimension_averages(means)
        paper = t4[row][0]
        out = {}
        for m, v in means.items():
            out[KEY[m]] = {"platform": v, "paper": paper[KEY[m]], "diff": v - paper[KEY[m]]}
        for g, v in dims.items():
            out[AVG[g]] = {"platform": v, "paper": paper[AVG[g]], "diff": v - paper[AVG[g]]}
        rows[row] = {"scores": len(scores), "cells": out}
        cells += [c["diff"] for c in out.values()]
    worst = max(abs(d) for d in cells)
    verdict = "reproduced" if worst <= REPRODUCED else "within one step" if worst <= ONE_STEP else "not reproduced"
    order = lambda key: [m for m in sorted(rows, key=lambda m: -rows[m]["cells"]["cc_avg"][key])]
    result["anchor"] = {
        "rows": rows, "cells": len(cells), "verdict": verdict, "max_abs_diff": worst, "mean_signed_diff": st.mean(cells),
        "within_0.0005": sum(abs(d) <= 0.0005 for d in cells), "within_0.001": sum(abs(d) <= 0.001 for d in cells),
        "within_0.005": sum(abs(d) <= REPRODUCED for d in cells), "within_0.02": sum(abs(d) <= ONE_STEP for d in cells),
        "character_consistency_order": {"platform": order("platform"), "paper": order("paper")},
    }
    json.dump(result, open(os.path.join(args.work, "anchor.json"), "w"), ensure_ascii=False, indent=1)
    a = result["anchor"]
    print(f"anchor: {verdict}; {a['cells']} cells, max |diff| {worst:.5f}, mean signed {a['mean_signed_diff']:+.5f}; "
          f"within 0.0005: {a['within_0.0005']}, 0.001: {a['within_0.001']}, 0.005: {a['within_0.005']}, 0.02: {a['within_0.02']}")
    print(f"order on character consistency: platform {a['character_consistency_order']['platform']}; paper {a['character_consistency_order']['paper']}")
    for row, r in rows.items():
        print(row, " ".join(f"{k}:{c['platform']:.3f}({c['diff']:+.4f})" for k, c in r["cells"].items()))


if __name__ == "__main__":
    main()