"""CoSER: the numbers of the paper checked against each other, against the platform's data module and against the released files.

Reads the LaTeX tables of the arXiv e-prints (never a PDF) and the test set and examples of the official repository at the
pinned commit (`git show`, nothing imported). Prints one line per finding and writes them to a JSON file. Findings about
the paper are facts to report, not pass/fail, except T01-T03, which guard the platform's data and fail with exit 1.
Prints counts only, never test-set text.

    python3 -I tools/experiments/coser_table_checks.py <eprint root holding v1 v2 v3> <CoSER repo> [--json out.json] [--main-table out.json]

--main-table writes the 21 rows of the main table as read from the LaTeX source; the report's Appendix B reads that file
(tools/experiments/coser_main_table.result.json), and T02 checks the same rows against published.ts.
"""
import argparse
import hashlib
import json
import math
import os
import re
import statistics
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
PINNED = "624efc50db7be3a4e37cec6a1c58ed434dc9f2ce"
COLS = ["sc", "an", "cf", "sq", "avg", "bleu", "rouge"]
SECTIONS = {"Close-source Models": "closed", "Open-source Models": "open"}
RELEVANT = [
    "Tables/main_exp.tex", "Tables/exp_idood.tex", "Tables/main_exp_cf3.tex", "Tables/main_exp_cot.tex", "Tables/exp_main_cot_full.tex", "Tables/exp_diff_judges.tex",
    "Tables/judge_alignment.tex", "Tables/human_eval.tex", "Tables/exp_wr.tex", "Tables/exp_rag.tex", "Tables/exp_otherbmks.tex",
    "Tables/prompts_agent.tex", "Tables/prompts_agent_2.tex", "Tables/prompts_eval.tex", "Tables/prompts_eval_2.tex",
    "Sections/4_Methodology.tex", "Sections/5_Experiments.tex", "Sections/Appendix.tex",
]

out = {}
failed = []


def say(cid, text, ok=None, **data):
    tag = "INFO" if ok is None else ("PASS" if ok else "FAIL")
    print(f"{tag} {cid} {text}")
    out[cid] = {"status": tag, "text": text, **data}
    if ok is False:
        failed.append(cid)


def read(eprint, v, path):
    return open(os.path.join(eprint, v, path), encoding="utf-8").read()


def cell(s):
    bold, under = "\\textbf{" in s, "\\underline{" in s
    m = re.search(r"([\d.]+)\s*(?:\\scriptsize|\\tiny)?\{?\$?\\?\{?\$?\\pm\s*([\d.]+)\$", s) or re.search(r"([\d.]+)\s*\\(?:scriptsize|tiny)\{\$\\pm([\d.]+)\$\}", s)
    if m:
        return {"v": float(m.group(1)), "pm": float(m.group(2)), "bold": bold, "under": under, "text": m.group(1)}
    m = re.search(r"(\d+\.\d+)", s)
    if m:
        return {"v": float(m.group(1)), "pm": None, "bold": bold, "under": under, "text": m.group(1)}
    return None


def rows_of(tex):
    rows, section = {}, None
    for line in tex.split("\n"):
        line = line.strip()
        m = re.search(r"\\textit\{([^}]*)\}", line)
        if line.startswith("\\multicolumn") and m and m.group(1) in SECTIONS:
            section = SECTIONS[m.group(1)]
            continue
        if not line.endswith("\\\\") or "&" not in line or "multirow" in line:
            continue
        parts = [p.strip() for p in line[:-2].split("&")]
        cells = [cell(p) for p in parts[1:]]
        if cells and all(c is not None for c in cells):
            rows[parts[0]] = {"section": section, "cells": cells}
    return rows


def decimals(c):
    return len(c["text"].split(".")[1]) if "." in c["text"] else 0


def git_show(repo, path):
    return subprocess.run(["git", "-C", repo, "show", f"{PINNED}:{path}"], capture_output=True, check=True).stdout


def platform_rows(path):
    text = open(os.path.join(ROOT, path), encoding="utf-8").read()
    data = text[text.index("PAPER_ROWS: PaperRow[] = ") + len("PAPER_ROWS: PaperRow[] = "):].rstrip().rstrip(";")
    return json.loads(data)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("eprint")
    ap.add_argument("repo")
    ap.add_argument("--json")
    ap.add_argument("--main-table")
    args = ap.parse_args()

    # Versions ---------------------------------------------------------------------------------------------------------
    same, diff = {}, {}
    for path in RELEVANT:
        texts = {v: read(args.eprint, v, path) if os.path.exists(os.path.join(args.eprint, v, path)) else None for v in ("v1", "v2", "v3")}
        (same if len(set(texts.values())) == 1 else diff)[path] = {v: hashlib.sha256(t.encode()).hexdigest()[:12] if t else None for v, t in texts.items()}
    say("VERSIONS", f"{len(same)} of {len(RELEVANT)} source files identical in v1, v2, v3; differing: {sorted(diff)}", same=sorted(same), differing=diff)

    main3 = rows_of(read(args.eprint, "v3", "Tables/main_exp.tex"))
    tex3 = read(args.eprint, "v3", "Tables/main_exp.tex")
    say("T01", f"main table: {len(main3)} rows, identical in v1, v2, v3",
        len(main3) == 21 and all(read(args.eprint, v, "Tables/main_exp.tex") == tex3 for v in ("v1", "v2")))
    bad = []
    for path in ("worker/src/benchmarks/coser/published.ts", "server/src/benchmarks/coser/published.ts"):
        rows = platform_rows(path)
        if [r["name"] for r in rows] != list(main3):
            bad.append((path, "names or order"))
        for r in rows:
            cs = main3[r["name"]]["cells"]
            for k, c in zip(COLS, cs):
                if r["scores"][k] != c["v"] or r["spread"][k] != c["pm"]:
                    bad.append((path, r["name"], k))
    say("T02", f"worker and server published.ts equal the LaTeX table: {21 * 7 * 2} means and {21 * 7 * 2} ± values, {len(bad)} differ", not bad, differing=bad)
    wsrc = open(os.path.join(ROOT, "worker/src/benchmarks/coser/published.ts"), encoding="utf-8").read()
    ssrc = open(os.path.join(ROOT, "server/src/benchmarks/coser/published.ts"), encoding="utf-8").read()
    sha = hashlib.sha256(tex3.encode()).hexdigest()
    say("T03", f"published.ts records the SHA-256 of Tables/main_exp.tex ({sha[:12]}…) and the two copies are the same file", sha in wsrc and wsrc == ssrc)

    # Main table: its own arithmetic ---------------------------------------------------------------------------------------
    gaps = {n: abs(sum(c["v"] for c in r["cells"][:4]) / 4 - r["cells"][4]["v"]) for n, r in main3.items()}
    worst = max(gaps, key=gaps.get)
    say("AVG", f"Average vs the mean of the four printed dimensions: largest gap {gaps[worst]:.4f} ({worst}); within 0.005 (rounding of one value) {sum(g <= 0.005 for g in gaps.values())}/21, within 0.01 {sum(g <= 0.01 for g in gaps.values())}/21",
        gaps={k: round(v, 4) for k, v in gaps.items()})
    marks = []
    for j, k in enumerate(COLS):
        best = max(r["cells"][j]["v"] for r in main3.values())
        best_open = max(r["cells"][j]["v"] for r in main3.values() if r["section"] == "open")
        for n, r in main3.items():
            c = r["cells"][j]
            if c["bold"] != (c["v"] == best):
                marks.append((k, n, "bold", c["v"], best))
            if r["section"] == "open" and c["under"] != (c["v"] == best_open):
                marks.append((k, n, "underline", c["v"], best_open))
    say("MARKS", f"bold (best overall) and underlined (best open) marks against the values: {len(marks)} mismatches", mismatches=marks)
    spreads = {k: [r["cells"][j]["pm"] for r in main3.values()] for j, k in enumerate(COLS)}
    say("SPREAD", "printed ± on the 21 rows: " + ", ".join(f"{k} {min(v):.2f}-{max(v):.2f} (median {statistics.median(v):.2f})" for k, v in spreads.items()) + "; the paper states 'averaged across three runs' and does not define ±",
        spreads={k: statistics.median(v) for k, v in spreads.items()})
    avg_pm_gt = [n for n, r in main3.items() if r["cells"][4]["pm"] > max(c["pm"] for c in r["cells"][:4])]
    say("SPREAD2", f"rows whose ± on Average exceeds the ± of all four dimensions: {len(avg_pm_gt)}", rows=avg_pm_gt)

    # In-domain / out-of-domain table against the main table -----------------------------------------------------------------
    idood = rows_of(read(args.eprint, "v3", "Tables/exp_idood.tex"))
    cmp, cols = [], ["sc", "an", "cf", "sq", "avg"]
    for n, r in idood.items():
        if n not in main3:
            continue
        for j, k in enumerate(cols):
            m = (r["cells"][j]["v"] + r["cells"][j + 5]["v"]) / 2
            cmp.append({"model": n, "col": k, "id": r["cells"][j]["v"], "ood": r["cells"][j + 5]["v"], "mean": round(m, 3), "main": main3[n]["cells"][j]["v"], "diff": round(m - main3[n]["cells"][j]["v"], 3)})
    avg_cmp = [c for c in cmp if c["col"] == "avg"]
    worst_c = max(cmp, key=lambda c: abs(c["diff"]))
    say("IDOOD", f"table exp_idood: {len(idood)} rows; mean of in-domain and out-of-domain against the main table (equal sizes, so it must equal it if it is the same runs): "
        f"Average differs by up to {max(abs(c['diff']) for c in avg_cmp):.3f}, mean signed {sum(c['diff'] for c in avg_cmp) / len(avg_cmp):+.3f}; any dimension up to {abs(worst_c['diff']):.3f} ({worst_c['model']} {worst_c['col']}); "
        f"cells within 0.01: {sum(abs(c['diff']) <= 0.01 for c in cmp)}/{len(cmp)}; beyond: {[(c['model'], c['col'], c['diff']) for c in cmp if abs(c['diff']) > 0.01]}", cells=cmp)

    # Different judges: the "GPT-4o" column ----------------------------------------------------------------------------------
    dj = {}
    for line in read(args.eprint, "v3", "Tables/exp_diff_judges.tex").split("\n"):
        parts = [p.strip() for p in line.strip().rstrip("\\").split("&")]
        if len(parts) == 4 and re.search(r"\d", parts[1]) and "Model" not in parts[0]:
            dj[parts[0].replace("\\method", "CoSER")] = [float(re.sub(r"[^\d.]", "", p)) for p in parts[1:]]
    agree = {}
    for n, v in dj.items():
        agree[n] = {"judges_table": v[0], "main_avg": main3[n]["cells"][4]["v"], "in_domain_avg": idood[n]["cells"][4]["v"]}
    say("JUDGES", f"table exp_diff_judges: the GPT-4o column equals the in-domain Average of exp_idood for {sum(abs(a['in_domain_avg'] - a['judges_table']) <= 0.051 for a in agree.values())}/{len(agree)} models, the full-test Average for {sum(abs(a['main_avg'] - a['judges_table']) <= 0.051 for a in agree.values())}/{len(agree)}; the caption does not say so",
        rows=agree)

    # Ablations -------------------------------------------------------------------------------------------------------------
    cot, full = rows_of(read(args.eprint, "v3", "Tables/main_exp_cot.tex")), rows_of(read(args.eprint, "v3", "Tables/exp_main_cot_full.tex"))
    std_ok = {n: (r["cells"][0]["v"], main3[n]["cells"][4]["v"]) for n, r in cot.items() if n in main3}
    say("ABL", f"ablation table: Standard equals the main table's Average for {sum(a == b for a, b in std_ok.values())}/{len(std_ok)} models; removing motivations raises the score for "
        f"{[n for n, r in cot.items() if n in main3 and r['cells'][2]['v'] > r['cells'][0]['v']]} although the text says improvements are 'consistent across all models'", rows=std_ok)

    # Claims in the text recomputed -----------------------------------------------------------------------------------------
    above57 = sorted(n for n, r in main3.items() if r["cells"][4]["v"] > 57)
    bleu = sorted(((r["cells"][5]["v"], n) for n, r in main3.items()), reverse=True)
    non_coser = next(b for b in bleu if not b[1].startswith("CoSER"))
    say("CLAIMS", f"Average above 57%: {above57} (the text lists GPT-4o, Gemini Pro, Claude-3.5-Sonnet, Doubao-pro, Step-2, Qwen-2-72B, CoSER-70B); GPT-4o Mini at {main3['GPT-4o Mini']['cells'][4]['v']} also qualifies; "
        f"Claim '58% above the second-best on BLEU': {bleu[0][1]} {bleu[0][0]} / best other {non_coser[1]} {non_coser[0]} = {bleu[0][0] / non_coser[0] - 1:.3f}; "
        f"Qwen-2 72B-7B gap {main3['Qwen-2-72B']['cells'][4]['v'] - main3['Qwen-2-7B']['cells'][4]['v']:.2f}")
    app = read(args.eprint, "v3", "Sections/Appendix.tex")
    meth = read(args.eprint, "v3", "Sections/4_Methodology.tex")
    say("FORMULA", f"initial score in Methodology v3: {'100-5*' in meth}; in Appendix v3 (length correction): {'100-5' in app} (still 100 − Σv); regression quoted: {'-1.5909' in app}; in v1 Methodology: {'100-5*' in read(args.eprint, 'v1', 'Sections/4_Methodology.tex')}")

    # Released examples --------------------------------------------------------------------------------------------------------
    ev = json.loads(git_show(args.repo, "exp/evaluation/test_set_gpt-4o.json"))
    per = []
    for key, c in ev["cases"].items():
        R = len([m for m in c["simulation"] if m["role"] != "Environment"])
        for dim, cr in c["critique"].items():
            if dim == "bleu" or dim == "rouge" or not isinstance(cr, dict):
                continue
            S = sum(f["severity"] for f in cr["flaws"] if isinstance(f["severity"], int))
            for name, value in (("100-5S+1.5R", 100 - 5 * S + 1.5 * R), ("100-S+1.5R", 100 - S + 1.5 * R), ("100-5S", 100 - 5 * S)):
                per.append({"case": hashlib.sha256(key.encode()).hexdigest()[:8], "dim": dim, "S": S, "R": R, "form": name, "value": value, "stored": cr["score"]})
    for form in ("100-5S+1.5R", "100-S+1.5R", "100-5S"):
        rows = [p for p in per if p["form"] == form]
        hit = sum(abs(min(max(p["value"], 0), 100) - p["stored"]) < 1e-9 for p in rows)
        say(f"EX-{form}", f"released evaluation of 3 scenes × 4 dimensions: the stored score equals clip({form}) for {hit}/{len(rows)}")

    # Test set -------------------------------------------------------------------------------------------------------------
    blob = git_show(args.repo, "data/test/test_set.json")
    test = json.loads(blob)
    speakers = [[x for x in s["speaking_characters_w_env"] if x != "Environment"] for s in test]
    thoughts = sum(1 for s, sp in zip(test, speakers) for n in sp if any(k.get("name") == n and k.get("thought") for k in s["key_characters"]))
    lens = [len(s["dialogues"]) for s in test]
    no_profile = sum(1 for s, sp in zip(test, speakers) for n in sp if not (s["character_profiles"].get(n) or any(k.get("name") == n and k.get("description") for k in s["plot"].get("key_characters", []))))
    majors = [len(s["major_characters"]) for s in test]
    env_listed = sum(1 for s in test if "Environment" in s["speaking_characters_w_env"])
    say("TEST", f"test set: SHA-256 {hashlib.sha256(blob).hexdigest()}; {len(test)} scenes, {sum(s.get('tag') == 'id' for s in test)} id, {sum(s.get('tag') == 'ood' for s in test)} ood; {len({s['book'] for s in test})} books "
        f"({len({s['book'] for s in test if s.get('tag') == 'id'})} id, {len({s['book'] for s in test if s.get('tag') == 'ood'})} ood); speaking characters per scene {min(map(len, speakers))}-{max(map(len, speakers))} (mean {sum(map(len, speakers)) / 200:.2f}, {sum(map(len, speakers))} in all), "
        f"{thoughts} with a motivation; original lines per scene {min(lens)}-{max(lens)} (median {statistics.median(lens)}), {sum(l > 20 for l in lens)} scenes longer than the 20-line limit; major characters per scene {min(majors)}-{max(majors)}; Environment listed as a speaker in {env_listed}; speakers without any profile text {no_profile}")
    other_sets = sorted(n for n in subprocess.run(["git", "-C", args.repo, "ls-tree", "-r", "--name-only", PINNED, "exp/"], capture_output=True, text=True).stdout.split("\n") if n)
    say("EXP-FILES", f"released result files under exp/: {other_sets}")

    if args.json:
        with open(args.json, "w", encoding="utf-8") as f:
            json.dump(out, f, ensure_ascii=False, indent=1)
    if args.main_table:
        with open(args.main_table, "w", encoding="utf-8") as f:
            json.dump({"source": "Tables/main_exp.tex", "sha256": sha, "rows": [{"name": n, "section": r["section"], "scores": {k: c["v"] for k, c in zip(COLS, r["cells"])}} for n, r in main3.items()]}, f, ensure_ascii=False, indent=1)
    print("failures:" if failed else "no failures", failed if failed else "")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())