#!/usr/bin/env python3
"""
Generate one merged white-box LaTeX table
(AUROC → Balanced Accuracy → Recall@1% FPR) and write to white_box_results.tex
"""

import re
from datetime import datetime
from pathlib import Path
from math import isnan
import pandas as pd

CSV_FILE = "multi_experiment_dashboard.csv"
OUTFILE  = "white_box_results.tex"
DECIMALS = 2
MISS     = "—"

# ---------- dataset → column label ----------
DATASET_MAP = {
    "harm_refusal_choice__plain"          : "HP-C",
    "harm_refusal_knowledge_report__plain": "HP-KR",
    "mask__plain"                         : "MASK",
    "insider_trading_XYZ__plain"      : "IT",
    "instruct_dishonesty__plain"          : "ID",
    "convince_me__plain"                  : "CM",
}
SLEEPER_VARIANTS = [
    "sleeper_agents__greeting",
    "sleeper_agents__gender",
    "sleeper_agents__time",
]
ST_LABEL = "ST (avg)"

# ---------- model-id → display ----------
MODEL_DISPLAY = {
    "llama-70b-3.3"         : "Llama 3.3 70B",
    "qwen-72b"              : "Qwen 2.5 72B",
    "mistral-small-3.1-24b" : "Mistral 3.1 24B",
    "llama-70b-3.1"         : "Llama 3.1 70B",
}
MODEL_ORDER = [
    "Llama 3.3 70B",
    "Qwen 2.5 72B",
    "Mistral 3.1 24B",
    "Llama 3.1 70B",
]

# ---------- helpers ----------
TS_RE = re.compile(r"__(\d{8}_\d{6})")
def ts(s):
    m = TS_RE.search(str(s))
    return datetime.strptime(m.group(1), "%Y%m%d_%H%M%S") if m else None
def num(x):
    try:
        f = float(x);  return None if isnan(f) else f
    except Exception:
        return None
def f(x):     # format
    return MISS if x is None else f"{x:.{DECIMALS}f}"

# ---------- load CSV & keep latest row per (model,dataset) ----------
if not Path(CSV_FILE).exists():
    raise SystemExit(f"{CSV_FILE} not found")

df = pd.read_csv(CSV_FILE)
df["__ts"] = df.apply(lambda r: ts(r.get("exp_path")) or ts(r.get("exp_name")), axis=1)
df = (df.sort_values("__ts")
        .groupby(["model", "dataset"], as_index=False)
        .tail(1)
        .reset_index(drop=True))

# ---------- build record dict ----------
records = {}
for _, r in df.iterrows():
    m = MODEL_DISPLAY.get(r["model"])
    if m is None:          # skip models not in display map
        continue
    rec = records.setdefault(m, {"ds": {}, "fprs": []})

    d = r["dataset"]
    au   = num(r.get("auroc"))
    ba   = num(r.get("balanced_accuracy"))
    rc   = num(r.get("recall 1.0%"))
    fp   = num(r.get("fpr 1.0%"))
    if fp is not None:
        rec["fprs"].append(fp)

    if d in DATASET_MAP:
        lbl = DATASET_MAP[d]
        rec["ds"][lbl] = {"auroc": au, "bal": ba, "rec": rc}
    elif d in SLEEPER_VARIANTS:
        rec["ds"].setdefault("__ST", []).append({"auroc": au, "bal": ba, "rec": rc})

# aggregate ST
for rec in records.values():
    variants = rec["ds"].pop("__ST", [])
    if variants:
        rec["ds"][ST_LABEL] = {
            "auroc":         sum(v["auroc"] for v in variants if v["auroc"] is not None) / len(variants) if variants else None,
            "bal":           sum(v["bal"]   for v in variants if v["bal"]   is not None) / len(variants) if variants else None,
            "rec":           sum(v["rec"]   for v in variants if v["rec"]   is not None) / len(variants) if variants else None,
        }

# ---------- LaTeX ----------
COLS = list(DATASET_MAP.values()) + [ST_LABEL]
tabcols = "l c " + "c " * len(COLS)

lines = [
    r"\begin{table*}[t]",
    r"\centering",
    fr"\begin{{tabular}}{{{tabcols.strip()}}}",
    r"\toprule",
    fr"Model & & \multicolumn{{{len(COLS)}}}{{c}}{{\textbf{{Model Organism}}}} \\",
    r"\textit{(all instruction-tuned)} & & " + " & ".join(COLS) + r" \\",
    r"\hline",
    r"& & \multicolumn{" + str(len(COLS)) + r"}{c}{\textsc{AUROC}} \\"
]
# AUROC rows
for m in MODEL_ORDER:
    ds = records.get(m, {}).get("ds", {})
    vals = [f(ds.get(c, {}).get("auroc")) for c in COLS]
    lines.append(rf"\textit{{{m}}} & & " + " & ".join(vals) + r" \\")

# balanced accuracy block
lines += [
    r"\hline",
    r"& & \multicolumn{" + str(len(COLS)) + r"}{c}{\textsc{Balanced Accuracy}} \\"
]
for m in MODEL_ORDER:
    ds = records.get(m, {}).get("ds", {})
    vals = [f(ds.get(c, {}).get("bal")) for c in COLS]
    lines.append(rf"\textit{{{m}}} & & " + " & ".join(vals) + r" \\")

# recall block
lines += [
    r"\hline",
    r"& FPR (1\%) & \multicolumn{" + str(len(COLS)) + r"}{c}{\textsc{Recall @ FPR}} \\"
]
for m in MODEL_ORDER:
    rec = records.get(m, {})
    fpr = f(max(rec.get("fprs", [])) if rec.get("fprs") else None)
    ds  = rec.get("ds", {})
    vals = [f(ds.get(c, {}).get("rec")) for c in COLS]
    lines.append(rf"\textit{{{m}}} & {fpr} & " + " & ".join(vals) + r" \\")

lines += [
    r"\end{tabular}",
    r"\vspace{0.75em}",
    r"\caption{White-box probe (AUROC, Balanced Accuracy, Recall). ST column averages sleeper variants; CM is Convince-Me. Per-model FPR (1\%) = max realised dataset FPR.}",
    r"\label{table:white-box-combined}",
    r"\end{table*}",
]

Path(OUTFILE).write_text("\n".join(lines) + "\n", encoding="utf-8")
print(f"Wrote {OUTFILE}")
