#!/usr/bin/env python3
"""Combined 20-seed classification for Figures 4 and 6 — one compute pass.

Usage:
    python scripts/visualizations/figure4_6.py              # both figures
    python scripts/visualizations/figure4_6.py --only 4     # Fig 4 only
    python scripts/visualizations/figure4_6.py --only 6     # Fig 6 only
    python scripts/visualizations/figure4_6.py --from-cache # use saved data
"""

import os, sys, json, argparse
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.dirname(os.path.abspath(__file__)))))
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from classes import PatientLeakageClassifier

PLOTS_DIR = os.path.join(ROOT, "plots")
RESULTS_DIR = os.path.join(ROOT, "results")
# Shared classification-runs cache, owned by scripts/classification.py. Both
# scripts fill it incrementally (same model/seed/split keys), merge-safe.
CACHE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "cache")
MODELS = ["VGG16", "DenseNet121", "EfficientNetB1", "MobileNetV2", "ResNet50"]
SEEDS = list(range(1, 21))
DATA_PATH = os.path.join(CACHE_DIR, "classification_runs.json")
MANIFEST = os.path.join(RESULTS_DIR, "simple_patient_manifest.csv")

ap = argparse.ArgumentParser()
ap.add_argument("--only", default=None, choices=["4", "6"])
ap.add_argument("--from-cache", action="store_true")
ap.add_argument("--force", action="store_true", help="Recompute all runs, ignoring cache.")
ap.add_argument("--manifest", default=MANIFEST)
ap.add_argument("--tag", default="")
args = ap.parse_args()
tag = f"_{args.tag}" if args.tag else ""

# Run or load. The cache is a flat list of per-run dicts; compute only the
# (model, seed, split) combos it's missing, so dropping a stale model's rows
# (e.g. ResNet50 after a feature change) recomputes just that model.
all_data = []
if os.path.exists(DATA_PATH) and not args.force:
    with open(DATA_PATH) as f: all_data = json.load(f)

if args.from_cache:
    print(f"Loaded {len(all_data)} cached runs ← {DATA_PATH}")
else:
    have = {(r["model"], r["seed"], r["split_type"]) for r in all_data}
    todo = [(m, s, st) for m in MODELS for s in SEEDS
            for st in ("image", "patient") if (m, s, st) not in have]
    if todo:
        clf = PatientLeakageClassifier(args.manifest, os.path.join(ROOT, "features"), n_jobs=6)
        for m, s, st in tqdm(todo, desc="Computing missing runs"):
            all_data.append(clf.run(m, s, st))
        os.makedirs(CACHE_DIR, exist_ok=True)
        with open(DATA_PATH, "w") as f: json.dump(all_data, f, indent=2)
        print(f"Computed {len(todo)} missing runs; cache now {len(all_data)}")
    else:
        print(f"Cache complete ({len(all_data)} runs); nothing to compute.")

# Build accs. Only plot models that actually have data, so a partially filled
# cache (e.g. ResNet50 dropped pending recompute) still renders without error.
accs = {m: {"image": [], "patient": []} for m in MODELS}
for r in all_data:
    if r["model"] in accs: accs[r["model"]][r["split_type"]].append(r["test"])
PLOT_MODELS = [m for m in MODELS if accs[m]["image"] or accs[m]["patient"]]
missing = [m for m in MODELS if m not in PLOT_MODELS]
if missing:
    print(f"WARNING: no cached runs for {missing}; run without --from-cache "
          "to compute them. Plotting remaining models only.")

# Figure 4: image-level only
if args.only is None or args.only == "4":
    fig, ax = plt.subplots(figsize=(8, 5))
    pos = list(range(1, len(PLOT_MODELS) + 1))
    bp = ax.boxplot([accs[m]["image"] for m in PLOT_MODELS], positions=pos,
                    widths=0.5, patch_artist=True, showfliers=True,
                    flierprops=dict(marker='o', markersize=3))
    for i, b in enumerate(bp['boxes']):
        b.set_facecolor('#4C9BD4')
        ax.annotate(f"{np.median(accs[PLOT_MODELS[i]]['image']):.3f}",
                    (pos[i], np.median(accs[PLOT_MODELS[i]]['image'])),
                    fontsize=6, ha='center', va='bottom')
    ax.set_xticks(pos); ax.set_xticklabels(PLOT_MODELS)
    ax.set_ylabel("Test accuracy"); ax.set_ylim(0.96, 1.00); ax.grid(axis='y', alpha=0.3)
    ax.set_title("Figure 4 — Image-level test accuracy across 20 seeds", fontsize=12)
    plt.tight_layout()
    out = os.path.join(PLOTS_DIR, "figure4", f"figure4{tag}.png")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    plt.savefig(out, dpi=150); plt.close()
    print(f"Saved → {out}")

# Figure 6: image vs patient
if args.only is None or args.only == "6":
    fig, ax = plt.subplots(figsize=(10, 6))
    for i, name in enumerate(PLOT_MODELS):
        for pos, stype, color in [(i*2+0.7, "image", '#4C9BD4'),
                                   (i*2+1.3, "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)
            ax.annotate(f"{np.median(data):.3f}", (pos, np.median(data)),
                        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(PLOT_MODELS)*2, 2)])
    ax.set_xticklabels(PLOT_MODELS)
    ax.set_ylabel("Test accuracy"); ax.set_ylim(0.70, 1.00); ax.grid(axis='y', alpha=0.3)
    ax.set_title("Figure 6 — Image vs Patient-level test accuracy (20 seeds)", fontsize=13)
    plt.tight_layout()
    out = os.path.join(PLOTS_DIR, "figure6", f"figure6{tag}.png")
    os.makedirs(os.path.dirname(out), exist_ok=True)
    plt.savefig(out, dpi=150); plt.close()
    print(f"Saved → {out}")

print("DONE")