"""Regenerate the three foundation rows of research_v2_latest/kfold/kfold_v2_summary.csv
under the uniform pooled-ensemble protocol (Table 4.1 source).

The six convolutional/CLIP rows were produced on the training VM as a pooled
five-fold ensemble (probabilities averaged across the five fold checkpoints,
scored once on the fixed 3,208-image test set). The three foundation models were
originally summarised by the mean of their per-fold scores instead, a different,
non-comparable aggregation that also denied them the same ensemble treatment.
This script brings the foundation models into the identical pooled-ensemble
protocol so all nine rows are directly comparable and every model carries a
single aligned per-image prediction (which also lets the McNemar test in
make_mcnemar9.py cover all nine).

The six CNN/CLIP rows are kept verbatim from the authoritative VM summary; only
the three foundation rows are recomputed from
kfold/foundation_fold<k>_<model>_preds.json. Confidence intervals use the
bootstrap documented in research_v2_latest/code/ensemble_and_stats.py: 2000
resamples, seed 42, 2.5/97.5 percentile. Brier uses the element-wise mean
convention of the original summary.

Run with the project .venv interpreter:
    .venv/bin/python thesis_build/make_kfold_summary.py
"""
import csv
import json
import numpy as np
from pathlib import Path
from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support, roc_auc_score,
    average_precision_score, cohen_kappa_score,
)

ROOT = Path(__file__).resolve().parent.parent
KF = ROOT / "kfold"
CSV = KF / "kfold_v2_summary.csv"
FOUNDATION = ["dinov2_l", "swin_b", "retfound"]
NUM_CLASSES = 10
LABELS = list(range(NUM_CLASSES))


def pooled_foundation(model):
    probs, labels = [], None
    for k in range(5):
        d = json.load(open(KF / f"foundation_fold{k}_{model}_preds.json"))
        probs.append(np.array(d["probs"])); labels = np.array(d["labels"])
    P = np.mean(probs, axis=0)
    return labels, P.argmax(1), P


def ece(probs, labels, n_bins=15):
    conf = probs.max(1); pred = probs.argmax(1); correct = (pred == labels).astype(float)
    bins = np.linspace(0, 1, n_bins + 1); e = 0.0
    for i in range(n_bins):
        m = (conf > bins[i]) & (conf <= bins[i + 1])
        if m.sum():
            e += m.mean() * abs(correct[m].mean() - conf[m].mean())
    return float(e)


def boot_ci(fn, labels, preds, n=2000, seed=42):
    rng = np.random.default_rng(seed); N = len(labels); vals = []
    for _ in range(n):
        idx = rng.integers(0, N, N)
        vals.append(fn(labels[idx], preds[idx]))
    return float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5))


def foundation_row(model):
    y, p, pr = pooled_foundation(model)
    acc = accuracy_score(y, p) * 100
    prec, rec, f1, _ = precision_recall_fscore_support(y, p, average="macro", zero_division=0)
    acc_lo, acc_hi = boot_ci(lambda a, b: accuracy_score(a, b) * 100, y, p)
    f1_lo, f1_hi = boot_ci(
        lambda a, b: precision_recall_fscore_support(a, b, average="macro", zero_division=0)[2] * 100, y, p)
    roc = roc_auc_score(y, pr, multi_class="ovr", average="macro", labels=LABELS)
    oh = np.eye(NUM_CLASSES)[y]
    prauc = average_precision_score(oh, pr, average="macro")
    brier = float(((pr - oh) ** 2).mean())
    return [model, "Foundation", round(acc, 2), round(acc_lo, 2), round(acc_hi, 2),
            round(f1 * 100, 2), round(f1_lo, 2), round(f1_hi, 2),
            round(prec * 100, 2), round(rec * 100, 2), round(roc, 4), round(prauc, 4),
            round(ece(pr, y), 4), round(cohen_kappa_score(y, p), 3), round(brier, 3),
            5, "five-fold pooled preds"]


def main():
    with open(CSV) as fh:
        rdr = csv.reader(fh); header = next(rdr); rows = list(rdr)
    kept = [r for r in rows if r[1] != "Foundation"]
    for r in kept:
        r[16] = "five-fold pooled preds"
    new_found = [foundation_row(m) for m in FOUNDATION]
    allrows = kept + new_found
    allrows.sort(key=lambda r: -float(r[2]))
    with open(CSV, "w", newline="") as fh:
        w = csv.writer(fh); w.writerow(header); w.writerows(allrows)
    with open("/tmp/summary_check.txt", "w") as fh:
        for r in allrows:
            fh.write(f"{r[0]:14s} acc {r[2]:6} ({r[3]}-{r[4]})  f1 {r[5]:6}  "
                     f"roc {r[10]}  ece {r[12]}  kappa {r[13]}  brier {r[14]}  {r[16]}\n")
    print("wrote", CSV.relative_to(ROOT))


if __name__ == "__main__":
    main()