"""Generate per_class_metrics_pooled.csv from the latest five-fold predictions.

This is the provenance script for research_v2_latest/analysis/per_class_metrics_pooled.csv,
the per-class precision/recall/F1/support table that drives Appendix B and Figure 4.2.

Aggregation matches Table 4.1 exactly under the uniform pooled-ensemble protocol:
  * the five convolutional networks and CLIP use their pooled five-fold test
    predictions (kfold/cnn_clip/<model>_test_preds.json), i.e. the probabilities
    averaged across the five fold checkpoints, scored once on the fixed test set;
  * the three foundation models are pooled the same way, averaging the probability
    vectors of their five fold checkpoints (kfold/foundation_fold<k>_<model>_preds.json)
    into one aligned per-image prediction before the per-class metrics are computed.

Run with the project .venv interpreter (needs numpy + scikit-learn):
    .venv/bin/python thesis_build/make_per_class_pooled_csv.py
"""
import csv
import json
import numpy as np
from pathlib import Path
from sklearn.metrics import precision_recall_fscore_support

ROOT = Path(__file__).resolve().parent.parent
KF = ROOT / "kfold"
OUT = ROOT / "analysis" / "per_class_metrics_pooled.csv"

# Class index -> display name (the dataset's alphabetical label order).
CLASSES = [
    "Central Serous Chorioretinopathy",  # 0
    "Diabetic Retinopathy",              # 1
    "Disc Edema",                        # 2
    "Glaucoma",                          # 3
    "Healthy",                           # 4
    "Macular Scar",                      # 5
    "Myopia",                            # 6
    "Pterygium",                         # 7
    "Retinal Detachment",                # 8
    "Retinitis Pigmentosa",              # 9
]
LABELS = list(range(len(CLASSES)))

# Model write order = accuracy-descending order of Table 4.1 (DINOv2-L is 5th).
CNN_CLIP = ["inception_v3", "clip_openai", "vgg19", "resnet101", "densenet121", "resnet50"]
FOUNDATION = ["dinov2_l", "swin_b", "retfound"]
MODEL_ORDER = ["inception_v3", "clip_openai", "vgg19", "resnet101", "dinov2_l",
               "densenet121", "resnet50", "swin_b", "retfound"]


def pooled_cnn_clip(model):
    """Per-class metrics from the single pooled five-fold prediction file."""
    d = json.load(open(KF / "cnn_clip" / f"{model}_test_preds.json"))
    y, p = np.array(d["labels"]), np.array(d["preds"])
    P, R, F, S = precision_recall_fscore_support(y, p, labels=LABELS, zero_division=0)
    return P, R, F, S


def pooled_foundation(model):
    """Per-class metrics from the pooled five-fold ensemble (probabilities averaged
    across the five fold checkpoints, then argmax), matching Table 4.1."""
    files = sorted(KF.glob(f"foundation_fold*_{model}_preds.json"))
    assert len(files) == 5, f"expected 5 folds for {model}, found {len(files)}"
    probs, y = [], None
    for f in files:
        d = json.load(open(f))
        probs.append(np.array(d["probs"])); y = np.array(d["labels"])
    p = np.mean(probs, axis=0).argmax(1)
    P, R, F, S = precision_recall_fscore_support(y, p, labels=LABELS, zero_division=0)
    return P, R, F, S


def main():
    rows = []
    for model in MODEL_ORDER:
        if model in CNN_CLIP:
            P, R, F, S = pooled_cnn_clip(model)
        else:
            P, R, F, S = pooled_foundation(model)
        for i in LABELS:
            rows.append([model, CLASSES[i], P[i], R[i], F[i], int(S[i])])

    with open(OUT, "w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["model", "class", "precision", "recall", "f1", "support"])
        w.writerows(rows)
    print(f"wrote {OUT.relative_to(ROOT)} ({len(rows)} rows)")


if __name__ == "__main__":
    main()