#!/usr/bin/env python3
"""
figure6.py — test accuracy boxplots across 20 seeds (Figure 6).

Usage:
    conda activate fundus_imaging
    python scripts/figure6.py
"""

import os, sys, json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from tqdm import tqdm

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from classes import PatientLeakageClassifier

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODELS = ["VGG16", "DenseNet121", "EfficientNetB1", "MobileNetV2", "ResNet50"]
SEEDS = list(range(1, 21))

clf = PatientLeakageClassifier(
    os.path.join(ROOT, "results", "simple_patient_manifest.csv"),
    os.path.join(ROOT, "features"),
    n_jobs=8)

# ---------------------------------------------------------------------------
# Run all seeds
# ---------------------------------------------------------------------------

all_data = []  # flat list of per-run dicts for raw data file

for name in tqdm(MODELS, desc="Models"):
    for seed in tqdm(SEEDS, desc=f"  {name} seeds", leave=False):
        for stype in ["image", "patient"]:
            r = clf.run(name, seed, stype)
            all_data.append(r)

# Save raw data
with open(os.path.join(ROOT, "results", "figure6_data.json"), "w") as f:
    json.dump(all_data, f, indent=2)

# ---------------------------------------------------------------------------
# Build per-model accuracy lists
# ---------------------------------------------------------------------------

accs = {}  # model -> {image: [accs], patient: [accs]}
for name in MODELS:
    accs[name] = {"image": [], "patient": []}
for r in all_data:
    accs[r["model"]][r["split_type"]].append(r["test"])

# ---------------------------------------------------------------------------
# Plot
# ---------------------------------------------------------------------------

fig, ax = plt.subplots(1, 1, figsize=(10, 6))

for i, name in enumerate(MODELS):
    pos_img = i * 2 + 0.7
    pos_pat = i * 2 + 1.3

    for pos, stype, color in [(pos_img, "image", '#4C9BD4'),
                               (pos_pat, "patient", '#6DBF6D')]:
        data = accs[name][stype]
        bp = ax.boxplot(data, positions=[pos], widths=0.5,
                        patch_artist=True, showfliers=True,
                        flierprops=dict(marker='o', markersize=3))
        bp['boxes'][0].set_facecolor(color)
        med = np.median(data)
        ax.annotate(f"{med:.3f}", (pos, med), fontsize=6,
                   ha='center', va='bottom')

ax.legend(handles=[
    Patch(facecolor='#4C9BD4', label='Image-level split'),
    Patch(facecolor='#6DBF6D', label='Patient-level split'),
], loc='lower right')

ax.set_xticks([p + 1 for p in range(0, len(MODELS) * 2, 2)])
ax.set_xticklabels(MODELS)
ax.set_ylabel("Test accuracy")
ax.set_title("Figure 6 — Image-level vs Patient-level test accuracy (20 seeds)")
ax.set_ylim(0.7, 1.02)
ax.grid(axis='y', alpha=0.3)

plt.tight_layout()
out = os.path.join(ROOT, "plots", "figure6.png")
plt.savefig(out, dpi=150)
plt.close()

print(f"\nSaved → {out}")
print(f"Saved → {ROOT}/results/figure6_data.json")
print("DONE")