"""Confirmatory held-out tables for one cell: per-run and pooled (r2+r3) paired, template-clustered CIs.

Pooling resamples whole templates jointly across runs, so within-template correlation across runs and
seeds is kept. Primary preregistered comparison: C6 vs C4b (target rate).

    python scripts/confirmatory_table.py runs/remote-all/heldout/tiny-vangogh results/confirmatory-tiny-vangogh
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

import pandas as pd

from uvw.stats import cluster_bootstrap_ci, paired_cluster_bootstrap_diff

RUNS = {
    "C6": ["tiny/vangogh/C6/c6-tiny-vangogh-r2/vg_preserve_regularized",
           "tiny/vangogh/C6/c6-tiny-vangogh-r3/vg_value_rewrite_r10"],
    "C5": ["tiny/vangogh/C5/c5-tiny-vangogh-r2/vg_value_remap_bounded",
           "tiny/vangogh/C5/c5-tiny-vangogh-r3/vg_key_value_regularized"],
    "C4b": ["tiny/vangogh/C4b/c4b-tiny-vangogh-r2/c4b-00",
            "tiny/vangogh/C4b/c4b-tiny-vangogh-r3/c4b-08"],
}
REFERENCES = {"C1": "original",
              "C2": "tiny/vangogh/C2/c2-uce-tiny-vangogh/uce-preserve-x2",
              "C3": "tiny/vangogh/C3/c3-esd-tiny-vangogh/esd-esd-x"}
METRICS = {  # label: (suite, role, column)
    "target rate": ("vangogh-test", "target", "target_hit"),
    "target rate (CLIP-H)": ("vangogh-test", "target", "clip_h14_hit"),
    "indirect probes": ("vangogh-probes", "target", "target_hit"),
    "neighbour correct": ("vangogh-test", "neighbour", "neighbour_correct"),
    "neighbour CLIP": ("vangogh-test", "neighbour", "clip_score"),
    "COCO CLIP": ("coco-test", None, "clip_score"),
}
N_BOOT = 10_000


def load(root: Path, ckpt: str, label: str) -> pd.DataFrame:
    suite, role, col = METRICS[label]
    df = pd.read_json(root / ckpt / suite / "scores.jsonl", lines=True)
    if role:
        df = df[df.role == role]
    df = df[df[col].notna()]
    return df[["prompt_id", "seed", "group_id", col]].rename(columns={col: "y"})


def paired_frame(root: Path, a: str, b: str, label: str) -> pd.DataFrame:
    m = load(root, a, label).merge(load(root, b, label), on=["prompt_id", "seed"], suffixes=("_a", "_b"),
                                   validate="one_to_one")
    assert (m.group_id_a == m.group_id_b).all()
    return m


def diff(frames: list[pd.DataFrame]) -> dict:
    m = pd.concat(frames, ignore_index=True)
    d, lo, hi, p = paired_cluster_bootstrap_diff(m.y_a.astype(float), m.y_b.astype(float), m.group_id_a, n_boot=N_BOOT)
    return {"diff": d, "lo": lo, "hi": hi, "p": p, "n": len(m)}


def absolute(root: Path, ckpts: list[str], label: str) -> dict:
    df = pd.concat([load(root, c, label) for c in ckpts], ignore_index=True)
    m, lo, hi = cluster_bootstrap_ci(df.y.astype(float), df.group_id, n_boot=N_BOOT)
    return {"mean": m, "lo": lo, "hi": hi, "n": len(df)}


def main(root: Path, out: Path) -> None:
    res: dict = {"absolute": {}, "differences": {}}
    for cond, ck in REFERENCES.items():
        res["absolute"][cond] = {l: absolute(root, [ck], l) for l in METRICS}
    for cond, runs in RUNS.items():
        for i, ck in enumerate(runs, start=2):
            res["absolute"][f"{cond} r{i}"] = {l: absolute(root, [ck], l) for l in METRICS}
        res["absolute"][f"{cond} pooled r2+r3"] = {l: absolute(root, runs, l) for l in METRICS}
    comparisons = [(c, ref) for c in RUNS for ref in ("C1", "C2", "C3")] + [("C6", "C4b"), ("C6", "C5")]
    for a, b in comparisons:
        key = f"{a} − {b}"
        res["differences"][key] = {}
        for label in METRICS:
            if b in RUNS:  # run-paired: r2 vs r2, r3 vs r3
                frames = [paired_frame(root, ra, rb, label) for ra, rb in zip(RUNS[a], RUNS[b])]
            else:
                frames = [paired_frame(root, ra, REFERENCES[b], label) for ra in RUNS[a]]
            res["differences"][key][label] = diff(frames)
    out.with_suffix(".json").write_text(json.dumps(res, indent=2))
    lines = ["| condition | " + " | ".join(METRICS) + " |", "|---|" + "---|" * len(METRICS)]
    for cond, v in res["absolute"].items():
        lines.append(f"| {cond} | " + " | ".join(f"{v[l]['mean']:.3f} [{v[l]['lo']:.3f}, {v[l]['hi']:.3f}]" for l in METRICS) + " |")
    lines += ["", "| pooled r2+r3 difference | " + " | ".join(METRICS) + " |", "|---|" + "---|" * len(METRICS)]
    for key, v in res["differences"].items():
        lines.append(f"| {key} | " + " | ".join(f"{v[l]['diff']:+.3f} [{v[l]['lo']:+.3f}, {v[l]['hi']:+.3f}]" for l in METRICS) + " |")
    out.with_suffix(".md").write_text("\n".join(lines) + "\n")
    print("\n".join(lines))


if __name__ == "__main__":
    main(Path(sys.argv[1]), Path(sys.argv[2]))