"""CharacterEval: probe set for checking a BaichuanCharRM deployment (matrix V06/V07; charactereval_rm_server_check.py).

The probes are reward-model inputs built by the official logic (charactereval_reference.py: transform_format.py and
run_char_rm.py's format_input, pasted verbatim) from one of the repository's frozen responses,
results/generation_trans_chatglm3_6b.jsonl at commit c3d44a6, so they are texts the official pipeline itself scores:

- for each of the 12 metrics, the two distinct inputs with the smallest sample ids (24 probes);
- the longest of all 8,032 inputs of that file;
- one input made longer than the model's 4,096-token window (the longest input's context repeated until the text
  exceeds 20,000 characters), which the official code truncates from the left.

    python3 -I tools/experiments/charactereval_rm_probes.py <CharacterEval repo at c3d44a6> <out.json>

Prints the SHA-256 of the written file; charactereval_rm_server_check.py refuses a probe file with another hash.
"""
import copy
import hashlib
import json
import os
import sys

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

FROZEN = "results/generation_trans_chatglm3_6b.jsonl"


def main():
    repo, out = sys.argv[1:3]
    for f, h in ref.SHA256.items():
        if ref.sha256_file(os.path.join(repo, f)) != h:
            sys.exit(f"{f}: SHA-256 differs from commit {ref.COMMIT}")
    datas = json.load(open(os.path.join(repo, FROZEN), encoding="utf-8"))
    id_metric = json.load(open(os.path.join(repo, "data/id2metric.jsonl"), encoding="utf-8"))  # one JSON document, as transform_format.py reads it
    profiles = ref.load_profiles(repo)
    records = ref.transform(datas, id_metric)
    texts = ref.rm_inputs(records, profiles)
    rows = [{"id": r["id"], "role": r["role"], "metric_en": r["metric_en"], "text": t} for r, t in zip(records, texts)]

    probes = []
    for metric in sorted({r["metric_en"] for r in rows}):
        # The frozen file repeats some samples (e.g. id 37); the two probes of a metric are distinct texts.
        distinct = {}
        for r in sorted((r for r in rows if r["metric_en"] == metric), key=lambda r: r["id"]):
            distinct.setdefault(r["text"], r)
        for r in list(distinct.values())[:2]:
            probes.append({**r, "kind": "metric"})
    longest = max(rows, key=lambda r: (len(r["text"]), r["id"], r["metric_en"]))
    probes.append({**longest, "kind": "longest"})
    rec = copy.deepcopy(next(x for x in records if x["id"] == longest["id"] and x["metric_en"] == longest["metric_en"]))
    context = rec["context"]
    while len(ref.rm_inputs([rec], profiles)[0]) <= 20000:
        rec["context"] = rec["context"] + "\n" + context
    probes.append({"id": rec["id"], "role": rec["role"], "metric_en": rec["metric_en"], "text": ref.rm_inputs([rec], profiles)[0], "kind": "over-window"})

    doc = {
        "source": f"{ref.REPO_URL} @ {ref.COMMIT}, {FROZEN}",
        "builder": "tools/experiments/charactereval_rm_probes.py (official transform_format.py + run_char_rm.py format_input)",
        "probes": [{"n": i, **p} for i, p in enumerate(probes)],
    }
    data = json.dumps(doc, ensure_ascii=False, indent=1).encode("utf-8")
    open(out, "wb").write(data)
    print(len(probes), "probes", hashlib.sha256(data).hexdigest())


if __name__ == "__main__":
    main()