#!/usr/bin/env python3
"""Combined CV-accuracy-vs-features curves for Figures 3, S1, and S4.

These three figures all compute the same per-model curve — RF importance
ranking + a per-nfeatures gamma grid at seed 20 — differing only in which
models and which split they show:

    Figure 3  : VGG16, image-level      (1 panel)   == panel (a) of S1
    Figure S1 : all 5 models, image     (5 panels)
    Figure S4 : all 5 models, patient   (5 panels)

Computing them together runs each (model, split) curve once (5 image + 5
patient = 10 curves) and caches them, instead of recomputing VGG16's image
curve for both Figure 3 and S1. The cache is incremental: dropping a model's
entries (e.g. after a feature change) recomputes only that model.

Usage:
    conda activate fundus_imaging
    python scripts/visualizations/figure3_s1_s4.py                # all three
    python scripts/visualizations/figure3_s1_s4.py --only s4      # just Fig S4
    python scripts/visualizations/figure3_s1_s4.py --from-cache   # re-render only
    python scripts/visualizations/figure3_s1_s4.py --force        # recompute all
"""

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

os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("MKL_NUM_THREADS", "1")

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__))))
# Shared results-cache tier (alongside classification_runs.json). These curves
# are model-selection output (best nfeat/gamma/cv), not plotting scaffolding, so
# they live in scripts/cache/ rather than under visualizations/.
CACHE_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "cache")
PLOTS_DIR = os.path.join(ROOT, "plots")

from classes import PatientLeakageClassifier

# Panel order matches the published supplementary figures.
MODELS = ["VGG16", "MobileNetV2", "DenseNet121", "ResNet50", "EfficientNetB1"]
SEED = 20
DATA_PATH = os.path.join(CACHE_DIR, "cv_curves.json")
MANIFEST = os.path.join(ROOT, "results", "simple_patient_manifest.csv")

# 16-point log grid of feature counts (each model caps at its own max dim).
NFEATS = np.unique(np.round(np.logspace(1.5, 5, 16)).astype(int)).tolist()
GAMMA_LOGSPACE = (-1.5, 1, 7)
PANELS = ["(a)", "(b)", "(c)", "(d)", "(e)"]


def key(model, split):
    return f"{model}::{split}"


def load_cache():
    if os.path.exists(DATA_PATH):
        with open(DATA_PATH) as f:
            return json.load(f)
    return {}


def compute(cache, splits, n_jobs, manifest):
    """Fill any missing (model, split) curves in the cache; return updated cache."""
    todo = [(m, s) for s in splits for m in MODELS if key(m, s) not in cache]
    if not todo:
        print(f"Cache complete for requested figures ({len(cache)} curves).")
        return cache
    clf = PatientLeakageClassifier(manifest, os.path.join(ROOT, "features"),
                                   n_jobs=n_jobs)
    for m, s in tqdm(todo, desc="Computing missing curves"):
        cache[key(m, s)] = clf.cv_curve(m, SEED, s, nfeatures_list=NFEATS,
                                        gamma_logspace=GAMMA_LOGSPACE)
    os.makedirs(CACHE_DIR, exist_ok=True)
    with open(DATA_PATH, "w") as f:
        json.dump(cache, f, indent=2)
    print(f"Computed {len(todo)} curves; cache now {len(cache)}.")
    return cache


def _draw_curve(ax, curve, color, ylim, title=None):
    # Keys are ints on a fresh compute but strings once round-tripped through
    # JSON; iterate items so either works.
    items = sorted((int(n), v) for n, v in curve["curves"].items())
    nfeats = [n for n, _ in items]
    accs = [v * 100 for _, v in items]
    ax.plot(nfeats, accs, "o-", color=color, markersize=4, linewidth=1.2)
    ax.set_xscale("log")
    ax.set_xlabel("Number of selected features")
    ax.set_ylabel("CV accuracy (%)")
    ax.set_ylim(*ylim)
    ax.grid(alpha=0.3)
    bn, ba = curve["best_nfeat"], curve["best_cv"] * 100
    ax.axvline(bn, color="red", linestyle="--", linewidth=0.8, alpha=0.6)
    ax.plot(bn, ba, "r*", markersize=11)
    ax.annotate(f"n={bn}\n{ba:.2f}%", (bn, ba), fontsize=8,
                xytext=(10, -10), textcoords="offset points", color="red")
    if title:
        ax.set_title(title, fontsize=10)


def render_fig3(cache, tag):
    fig, ax = plt.subplots(figsize=(8, 5))
    _draw_curve(ax, cache[key("VGG16", "image")], "#D62728", (92, 100))
    ax.set_title("Figure 3 — VGG16 Image-level CV accuracy vs number of "
                 "selected features", fontsize=12)
    _save(fig, "figure3", f"figure3{tag}.png")


def render_grid(cache, split, color, ylim, suptitle, subdir, fname):
    fig, axes = plt.subplots(2, 3, figsize=(14, 9))
    axes = axes.flatten()
    for i, model in enumerate(MODELS):
        _draw_curve(axes[i], cache[key(model, split)], color, ylim,
                    title=f"{PANELS[i]} {model}")
    axes[5].set_visible(False)
    fig.suptitle(suptitle, fontsize=13, y=1.01)
    _save(fig, subdir, fname, tight=True)


def _save(fig, subdir, fname, tight=False):
    plt.tight_layout()
    out = os.path.join(PLOTS_DIR, subdir, fname)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    fig.savefig(out, dpi=150, bbox_inches="tight" if tight else None)
    plt.close(fig)
    print(f"Saved → {out}")


def main():
    ap = argparse.ArgumentParser(description="Combined Figures 3, S1 & S4.")
    ap.add_argument("--only", choices=["3", "s1", "s4"],
                    help="Render only one figure (still computes its curves).")
    ap.add_argument("--from-cache", action="store_true",
                    help="Re-render from cache without computing.")
    ap.add_argument("--force", action="store_true",
                    help="Recompute all curves, ignoring the cache.")
    ap.add_argument("--n-jobs", type=int, default=6)
    ap.add_argument("--manifest", default=MANIFEST)
    ap.add_argument("--tag", default="")
    args = ap.parse_args()
    tag = f"_{args.tag}" if args.tag else ""

    # Which splits are needed for the requested figure(s)?
    need_image = args.only in (None, "3", "s1")
    need_patient = args.only in (None, "s4")
    splits = (["image"] if need_image else []) + (["patient"] if need_patient else [])

    cache = {} if args.force else load_cache()
    if args.from_cache:
        print(f"Loaded {len(cache)} cached curves ← {DATA_PATH}")
    else:
        cache = compute(cache, splits, args.n_jobs, args.manifest)

    if args.only in (None, "3"):
        render_fig3(cache, tag)
    if args.only in (None, "s1"):
        render_grid(cache, "image", "#D62728", (90, 100),
                    "Figure S1 — Image-level CV accuracy vs number of selected "
                    "features", "figure_s1", f"figure_s1{tag}.png")
    if args.only in (None, "s4"):
        render_grid(cache, "patient", "#2C7BB6", (82, 95),
                    "Figure S4 — Patient-level CV accuracy vs number of selected "
                    "features (C=10)", "figure_s4", f"figure_s4{tag}.png")
    print("DONE")


if __name__ == "__main__":
    main()