"""CharacterEval: mechanical checks of the paper's own numbers (reproduction report, Sections 1.5 and 2).

Reads the LaTeX source of arXiv:2401.01275v2 and checks, without any model:
  T1  the example counts: Table 1 (training + test) against the total stated in the text;
  T2  the correlation table: whether "Overall" equals the mean of the 12 metrics, per column, and whether the GPT-4
      columns increase with the number of demonstrations, as the text says;
  T4  the dimension averages of Table 4 and of the appendix GPT-4 table against the unweighted mean of their printed
      metrics (rounding of the printed metrics allows up to 0.001), and the bold (best) and underlined (second best)
      marks of Table 4, column by column;
  MB  whether every personality back-testing value is k/216 (54 characters x 4 MBTI letters) within rounding;
  NB  the number of baselines stated in the text against the rows of Table 4.

    python3 -I tools/experiments/charactereval_table_checks.py <e-print dir with acl_latex.tex> [out.json]

Prints one line per finding and writes them as JSON. Exit status 0 (findings are facts about the paper, not failures).
"""
import hashlib
import json
import re
import sys

SOURCE_SHA256 = "f98299737e75bfb7002e47a76a67161246fbf1286c7be43f8ccd5800b50d6c17"
CC = ["ke", "ka", "kh", "pb", "pu", "cc_avg", "mbti"]
CA_RA = ["flu", "coh", "cons", "ca_avg", "hl", "cs", "ed", "emp", "ra_avg"]
AVGS = {"cc_avg": ["ke", "ka", "kh", "pb", "pu"], "ca_avg": ["flu", "coh", "cons"], "ra_avg": ["hl", "cs", "ed", "emp"]}


def live_lines(tex):
    return [(i + 1, l) for i, l in enumerate(tex.split("\n")) if not l.lstrip().startswith("%")]


def table_rows(lines, label):
    """Rows (line number, raw cells) of the table carrying `label`, from its \\begin{table} on."""
    end = next(i for i, (_, l) in enumerate(lines) if f"\\label{{{label}}}" in l)
    start = max(i for i, (_, l) in enumerate(lines[:end]) if "\\begin{table" in l)
    return [(n, [c.strip() for c in l.split("\\\\")[0].split("&")]) for n, l in lines[start:end] if "&" in l]


def value(cell):
    s = re.sub(r"\\multicolumn\{\d+\}\{[^}]*\}\{(.*)\}", r"\1", cell)
    s = re.sub(r"\\(textbf|underline)\{([^}]*)\}", r"\2", s).strip()
    try:
        return float(s)
    except ValueError:
        return None


def model_rows(rows, keys):
    out = {}
    for _, cells in rows:
        vals = [value(c) for c in cells[1:]]
        if sum(v is not None for v in vals) >= 5 and len(vals) == len(keys):
            out[value(cells[0]) or re.sub(r"\\[a-z]+\{|\}", "", cells[0]).strip()] = (dict(zip(keys, vals)), dict(zip(keys, cells[1:])))
    return out


def tables(lines, label):
    rows = table_rows(lines, label)
    # Each results table holds two tabulars: character consistency (with back-testing), then the other two dimensions.
    split = next(i for i, (_, c) in enumerate(rows) if any("Flu." in x for x in c))
    a, b = model_rows(rows[:split], CC), model_rows(rows[split:], CA_RA)
    return {m: ({**a[m][0], **b[m][0]}, {**a[m][1], **b[m][1]}) for m in a}


def main():
    src = sys.argv[1]
    out = sys.argv[2] if len(sys.argv) > 2 else None
    raw = open(f"{src}/acl_latex.tex", encoding="utf-8").read()
    digest = hashlib.sha256(raw.encode()).hexdigest()
    lines = live_lines(raw)
    findings = {"source_sha256": digest, "source_is_v2": digest == SOURCE_SHA256}

    # T1
    examples = next(c for _, c in table_rows(lines, "tab:statistic") if "Examples" in c[0])
    nums = [int(x.replace(",", "")) for x in examples[1:]]
    total_text = sorted({(n, m.group(0)) for n, l in lines for m in re.finditer(r"(\d{1,3},\d{3}) examples", l)})
    findings["T1"] = {"table1_numbers_with_thousands": nums, "text_totals": [{"line": n, "text": t} for n, t in total_text]}

    # T2
    corr = {}
    for _, cells in table_rows(lines, "table:coefficients"):
        vals = [value(c) for c in cells[1:]]
        if len(vals) == 4 and all(v is not None for v in vals):
            corr[re.sub(r"\\[a-z]+\{|\}", "", cells[0]).strip()] = vals
    metrics = [k for k in corr if k != "Overall"]
    cols = ["Char-RM", "1-shot", "2-shot", "3-shot"]
    findings["T2"] = {
        "metrics": len(metrics),
        "overall_printed": dict(zip(cols, corr["Overall"])),
        "mean_of_metrics": {c: round(sum(corr[m][i] for m in metrics) / len(metrics), 4) for i, c in enumerate(cols)},
        "gpt4_overall_increasing": corr["Overall"][1] < corr["Overall"][2] < corr["Overall"][3],
    }

    # T4
    t4, gpt4 = tables(lines, "tab:detail_result"), tables(lines, "tab:gpt4results")
    avgs, off = [], []
    for name, tab in (("Table 4", t4), ("GPT-4 table", gpt4)):
        for model, (v, _) in tab.items():
            for avg, parts in AVGS.items():
                mean = sum(v[p] for p in parts) / len(parts)
                avgs.append(abs(mean - v[avg]))
                if round(mean, 3) != v[avg]:
                    off.append({"table": name, "model": model, "column": avg, "printed": v[avg], "mean_of_printed_metrics": round(mean, 5)})
    marks = []
    for col in [k for k in CC + CA_RA if k != "mbti"]:
        vals = sorted({v[col] for v, _ in t4.values()}, reverse=True)
        for model, (v, c) in t4.items():
            bold, under = "\\textbf" in c[col], "\\underline" in c[col]
            if bold != (v[col] == vals[0]) or under != (len(vals) > 1 and v[col] == vals[1]):
                marks.append({"column": col, "model": model, "value": v[col], "bold": bold, "underline": under, "best": vals[0], "second": vals[1]})
    findings["T4"] = {"models": [len(t4), len(gpt4)], "averages": len(avgs), "max_abs_diff": round(max(avgs), 5),
                      "printed_not_equal_round3_of_mean": off, "bold_underline_mismatches": marks}

    # MB
    mbti = {m: v["mbti"] for m, (v, _) in t4.items() if v["mbti"] is not None}
    findings["MB"] = {"values": len(mbti), "all_k_over_216": all(any(abs(k / 216 - x) <= 0.0005 for k in range(217)) for x in mbti.values()),
                      "k": {m: min(range(217), key=lambda k: abs(k / 216 - x)) for m, x in mbti.items()}}

    # NB
    nb = [{"line": n, "text": m.group(0)} for n, l in lines for m in re.finditer(r"\b(\d+|seven|ten) baselines\b", l)]
    findings["NB"] = {"text": nb, "table4_rows": len(t4)}

    print(f"source sha256 {digest[:16]} ({'arXiv v2' if findings['source_is_v2'] else 'NOT v2'})")
    print(f"T1 Table 1 numbers {nums}; totals in text: {[t['text'] + ' (line ' + str(t['line']) + ')' for t in findings['T1']['text_totals']]}")
    print(f"T2 Overall printed {findings['T2']['overall_printed']}; mean of {len(metrics)} metrics {findings['T2']['mean_of_metrics']}; GPT-4 Overall increasing with shots: {findings['T2']['gpt4_overall_increasing']}")
    print(f"T4 {len(avgs)} dimension averages, max |printed - mean of printed metrics| {findings['T4']['max_abs_diff']}; {len(off)} not equal to the rounded mean")
    for m in marks:
        print(f"T4 mark: {m['column']} {m['model']} {m['value']} bold={m['bold']} underline={m['underline']} (best {m['best']}, second {m['second']})")
    print(f"MB {len(mbti)} back-testing values, all k/216: {findings['MB']['all_k_over_216']}")
    print(f"NB baselines in text: {[x['text'] + ' (line ' + str(x['line']) + ')' for x in nb]}; Table 4 rows: {len(t4)}")
    if out:
        json.dump(findings, open(out, "w"), ensure_ascii=False, indent=1)


if __name__ == "__main__":
    main()