"""Compute the nine-model soft-vote ensemble on the pooled five-fold predictions.

Every other result in the thesis reduces each model to a single pooled prediction
per image (the probabilities of its five fold checkpoints averaged on the fixed
3,208-image test set). The original ensemble report shipped from the training VM
was instead computed on the single-checkpoint predictions, a different and
non-comparable aggregation. This script recomputes the ensemble in the same
pooled paradigm so Section 4.8 is consistent with Table 4.1 and the McNemar
analysis: each model is first pooled across its folds, the nine pooled
probability vectors are then combined by a soft vote weighted by each model's
macro-F1, and the result is compared against the best single pooled model with
the same exact-binomial McNemar test used in Section 4.5.

Writes research_v2_latest/analysis/ensemble_pooled.json.

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

ROOT = Path(__file__).resolve().parent.parent
KF = ROOT / "kfold"
OUT = ROOT / "analysis" / "ensemble_pooled.json"

CNN_CLIP = {"inception_v3", "clip_openai", "vgg19", "resnet101", "densenet121", "resnet50"}
ALL = ["inception_v3", "clip_openai", "vgg19", "resnet101", "dinov2_l",
       "densenet121", "resnet50", "swin_b", "retfound"]
NUM_CLASSES = 10
LABELS = list(range(NUM_CLASSES))


def pooled(model):
    if model in CNN_CLIP:
        d = json.load(open(KF / "cnn_clip" / f"{model}_test_preds.json"))
        return np.array(d["labels"]), np.array(d["probs"])
    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"])
    return labels, np.mean(probs, axis=0)


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 mcnemar(pred_a, pred_b, labels):
    ca = pred_a == labels; cb = pred_b == labels
    b = int(np.sum(ca & ~cb)); c = int(np.sum(~ca & cb)); n = b + c
    p = 1.0 if n == 0 else float(min(1.0, 2 * binom.cdf(min(b, c), n, 0.5)))
    return b, c, p


def main():
    P = {}; L = None
    for m in ALL:
        L, P[m] = pooled(m)
    f1 = {m: precision_recall_fscore_support(L, P[m].argmax(1), average="macro", zero_division=0)[2]
          for m in ALL}

    # soft vote weighted by each model's pooled macro-F1
    wsum = sum(f1.values())
    ens_probs = sum(f1[m] * P[m] for m in ALL) / wsum
    ens_pred = ens_probs.argmax(1)

    acc = accuracy_score(L, ens_pred) * 100
    prec, rec, fm, _ = precision_recall_fscore_support(L, ens_pred, average="macro", zero_division=0)
    roc = roc_auc_score(L, ens_probs, multi_class="ovr", average="macro", labels=LABELS)

    # best single pooled model
    best = max(ALL, key=lambda m: accuracy_score(L, P[m].argmax(1)))
    best_acc = accuracy_score(L, P[best].argmax(1)) * 100
    b, c, p = mcnemar(ens_pred, P[best].argmax(1), L)

    report = {
        "protocol": "pooled five-fold soft vote, weight = pooled macro-F1, all nine models",
        "members": ALL,
        "ensemble": {
            "acc": round(acc, 2), "f1": round(fm * 100, 2), "roc_auc": round(roc, 4),
            "ece": round(ece(ens_probs, L), 4), "kappa": round(cohen_kappa_score(L, ens_pred), 3),
        },
        "best_single": {"model": best, "acc": round(best_acc, 2), "f1": round(f1[best] * 100, 2)},
        "ensemble_vs_best_mcnemar": {"b": b, "c": c, "p": round(p, 3),
                                     "significant_005": bool(p < 0.05)},
        "n_test": int(len(L)),
    }
    with open(OUT, "w") as fh:
        json.dump(report, fh, indent=2)
    with open("/tmp/ensemble_pooled_check.txt", "w") as fh:
        fh.write(json.dumps(report, indent=2))
    print("wrote", OUT.relative_to(ROOT))


if __name__ == "__main__":
    main()