scripts/05_evaluate.py

"""Official LUNA16 evaluation, uncertainty estimation and manuscript tables.

Produces every number the revision needs, from the cached full-volume candidate
tables. Nothing here touches the GPU, so the analysis can be re-run freely.

Outputs (under ``results/``)
---------------------------
``table_main.csv``            CPM + bootstrap CI for Exp1-Exp4, official protocol
``table_froc_points.csv``     sensitivity at the seven CPM operating points
``table_paired.csv``          paired bootstrap CIs for the key contrasts
``table_foldwise.csv``        per-fold CPM + paired Wilcoxon
``table_r_sweep.csv``         supervision-extent sweep
``table_wmin_sweep.csv``      minimum-box-size sweep
``table_negatives.csv``       negative-mining ablation
``table_seeds.csv``           multi-seed mean +/- SD
``table_yolo26.csv``          YOLO11n vs YOLO26n
``table_excluded_effect.csv`` effect of applying annotations_excluded.csv
``table_size_recall.csv``     size-stratified recall at a fixed operating point
``froc_curves.csv``           curve points for the figures

Usage
-----
    python scripts/05_evaluate.py [--bootstrap 1000]
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from luna_rev import config as cfg
from luna_rev import evaluate as ev
from luna_rev import splits, stats
from luna_rev.io_luna import all_uids, group_by_uid, load_annotations, load_excluded
from luna_rev.predict import candidates_path, load_candidates

R = cfg.RESULTS_DIR


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #

def available(exps):
    """Configurations whose candidate table exists on disk."""
    return [e for e in exps if candidates_path(e.name).exists()]


def cohort_for(exps, folds_by_index) -> list[str]:
    """Scans covered by *all* the given configurations (their common test folds)."""
    common = None
    for e in exps:
        cov = set()
        for k in cfg.folds_for(e):
            cov.update(folds_by_index[k].test)
        common = cov if common is None else (common & cov)
    return sorted(common or [])


def evaluate_set(exps, uids, included, excluded, n_boot, tag) -> tuple[pd.DataFrame, dict, dict]:
    """Point estimate + shared-resample bootstrap for a set of configurations."""
    vectors, points = {}, {}
    for e in exps:
        vec = ev.match(load_candidates(e.name), included, excluded, uids)
        vectors[e.name] = vec
        points[e.name] = ev.evaluate(vec)

    boots = ev.bootstrap_cpm(vectors, uids, n_iter=n_boot) if n_boot else {}

    rows = []
    for e in exps:
        p = points[e.name]
        row = {
            "cohort": tag, "n_scans": len(uids),
            "experiment": e.name, "label": e.label, "group": e.group,
            "model": e.model, "representation": e.representation,
            "r_sample": e.r_sample, "w_min_px": e.w_min_px,
            "negatives": e.negatives, "seed": e.seed,
            "cpm": round(p["cpm"], 4),
        }
        if n_boot:
            s = ev.summarise_bootstrap(boots[e.name], p["cpm"])
            row.update({"ci_low": round(s["ci_low"], 4), "ci_high": round(s["ci_high"], 4),
                        "ci_width": round(s["ci_width"], 4)})
        row.update({f"sens@{k}": round(v, 4) for k, v in p["sensitivity"].items()})
        row.update({
            "n_candidates": p["total_candidates"],
            "candidates_per_scan": round(p["candidates_per_scan"], 1),
            "true_positives": p["true_positives"],
            "false_positives": p["false_positives"],
            "ignored_on_excluded": p["ignored_on_excluded"],
            "max_recall": round(p["sensitivity_at_saturation"], 4),
        })
        rows.append(row)
    return pd.DataFrame(rows), boots, points


def paired_table(exps, boots, points, pairs, tag) -> pd.DataFrame:
    rows = []
    for a, b in pairs:
        if a not in boots or b not in boots:
            continue
        d = ev.paired_difference(boots[a], boots[b], points[a]["cpm"] - points[b]["cpm"])
        rows.append({
            "cohort": tag, "contrast": f"{a} - {b}",
            "cpm_a": round(points[a]["cpm"], 4), "cpm_b": round(points[b]["cpm"], 4),
            "delta_cpm": round(d["delta"], 4),
            "ci_low": round(d["ci_low"], 4), "ci_high": round(d["ci_high"], 4),
            "p_bootstrap": round(d["p_bootstrap_two_sided"], 4),
            "significant": d["excludes_zero"],
        })
    return pd.DataFrame(rows)


def training_set_sizes() -> pd.DataFrame:
    """Mean training/validation image counts per configuration, from the run log.

    Used to make the negative-mining comparison honest: the arm without
    nodule-free scans is trained on fewer images, not merely on less diverse
    ones, and a reader is entitled to see both numbers.
    """
    log = R / "training_log.jsonl"
    if not log.exists():
        return pd.DataFrame(columns=["experiment", "mean_train_images"])
    rows = []
    for line in log.read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        r = json.loads(line)
        if r.get("status") == "trained" and "n_train_images" in r:
            rows.append({"experiment": r["exp"], "n_train_images": r["n_train_images"]})
    if not rows:
        return pd.DataFrame(columns=["experiment", "mean_train_images"])
    g = pd.DataFrame(rows).groupby("experiment", as_index=False)["n_train_images"].mean()
    return g.rename(columns={"n_train_images": "mean_train_images"}).round(0)


def clamping_analysis(folds_by_index, included, excluded) -> pd.DataFrame | None:
    """Relate CPM to label geometry across every ``(r, w_min)`` cell.

    The supervision-extent sweep and the minimum-box sweep are two lines through
    the same two-dimensional design space, and neither factor alone explains the
    outcome: at fixed ``r`` the ``w_min`` variants span nearly the full range of
    the ``r`` sweep. What both manipulations have in common is the fraction of
    training boxes whose side is set by the floor rather than by the annotated
    diameter (Eq. 7). This table reports that fraction against CPM, together
    with the candidate counts, so the comparison can be read off directly rather
    than inferred.
    """
    stats_path = R / "dataset_stats.json"
    if not stats_path.exists():
        return None
    clamp = {(v["r_sample"], v["w_min_px"]): v["frac_clamped"]
             for v in json.loads(stats_path.read_text(encoding="utf-8"))["variants"]}

    grid = [e for e in cfg.ALL_EXPERIMENTS
            if e.representation == "naive" and e.negatives == "all888"
            and e.model == "yolo11n.pt" and e.seed == 42]
    grid = available(grid)
    if len(grid) < 3:
        return None

    uids = cohort_for(grid, folds_by_index)
    rows = []
    for e in grid:
        p = ev.evaluate(ev.match(load_candidates(e.name), included, excluded, uids))
        rows.append({
            "experiment": e.name, "label": e.label,
            "r_sample": e.r_sample, "w_min_px": e.w_min_px,
            "frac_clamped": round(clamp.get((e.r_sample, e.w_min_px), float("nan")), 4),
            "cpm": round(p["cpm"], 4),
            "false_positives": p["false_positives"],
            "candidates_per_scan": round(p["candidates_per_scan"], 1),
            "max_recall": round(p["sensitivity_at_saturation"], 4),
            "n_scans": len(uids),
        })
    df = pd.DataFrame(rows).sort_values("frac_clamped")

    # Which design variable actually tracks CPM? Written to its own file: folding
    # correlations into the grid table forces them to reuse unrelated columns,
    # which is unreadable for anyone but its author.
    from scipy import stats as sps
    corr = []
    for col in ("r_sample", "w_min_px", "frac_clamped", "false_positives",
                "candidates_per_scan", "max_recall"):
        x, y = df[col].to_numpy(float), df["cpm"].to_numpy(float)
        if np.std(x) == 0:
            continue
        sp, pe = sps.spearmanr(x, y), sps.pearsonr(x, y)
        corr.append({
            "predictor": col, "n_configurations": len(df),
            "pearson_r": round(float(pe.statistic), 4),
            "pearson_p": round(float(pe.pvalue), 4),
            "spearman_rho": round(float(sp.statistic), 4),
            "spearman_p": round(float(sp.pvalue), 4),
        })
    write(pd.DataFrame(corr).sort_values("pearson_p"), "table_clamping_correlations.csv")
    return df


def write(df: pd.DataFrame, name: str) -> None:
    if df is None or df.empty:
        print(f"  (skipped {name}: no data)")
        return
    path = R / name
    df.to_csv(path, index=False)
    print(f"  wrote {path.name:28s} ({len(df)} rows)")


# --------------------------------------------------------------------------- #
# Main
# --------------------------------------------------------------------------- #

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--bootstrap", type=int, default=cfg.N_BOOTSTRAP)
    args = ap.parse_args()

    annotations = load_annotations()
    included = group_by_uid(annotations)
    excluded = group_by_uid(load_excluded())
    folds = splits.get_folds("official")
    folds_by_index = {f.index: f for f in folds}
    uids888 = list(all_uids())

    main_exps = available(cfg.MAIN_EXPERIMENTS)
    if not main_exps:
        raise SystemExit("No candidate tables found - run scripts/04_predict.py first.")
    print(f"configurations with candidates: {len(available(cfg.ALL_EXPERIMENTS))}"
          f"/{len(cfg.ALL_EXPERIMENTS)}")

    # ---------------- 1) Main table, full 888-scan cohort -------------------
    print("\n[1] main table (official protocol, 888 scans)")
    main_df, main_boot, main_pts = evaluate_set(
        main_exps, uids888, included, excluded, args.bootstrap, "official_888")
    write(main_df, "table_main.csv")

    froc_rows = []
    for e in main_exps:
        p = main_pts[e.name]
        for fp, s in zip(cfg.FROC_POINTS, p["sensitivity"].values()):
            froc_rows.append({"experiment": e.name, "label": e.label,
                              "fp_per_scan": fp, "sensitivity": round(s, 4)})
    write(pd.DataFrame(froc_rows), "table_froc_points.csv")

    curve_rows = []
    for e in main_exps:
        p = main_pts[e.name]
        fps, sens = p["curve_fps"], p["curve_sens"]
        # Keep one point past the last operating point so the plotted curve
        # spans 8 FP/scan. Clipping at exactly 8.0 leaves the last stored point
        # just below it, which would make the figure mark a measured operating
        # point as extrapolated.
        keep = fps <= 8.0
        beyond = np.flatnonzero(fps > 8.0)
        if beyond.size:
            keep = keep.copy()
            keep[beyond[0]] = True
        curve_rows.append(pd.DataFrame({"experiment": e.name, "label": e.label,
                                        "fp_per_scan": fps[keep], "sensitivity": sens[keep]}))
    write(pd.concat(curve_rows, ignore_index=True) if curve_rows else None, "froc_curves.csv")

    # ---------------- 2) Paired bootstrap contrasts -------------------------
    print("\n[2] paired bootstrap contrasts")
    names = {e.name for e in main_exps}
    pairs = [(a, b) for a, b in [
        ("Exp4_2p5D_Strict", "Exp3_2p5D_Loose"),
        ("Exp4_2p5D_Strict", "Exp1_2D_Loose"),
        ("Exp3_2p5D_Loose", "Exp1_2D_Loose"),
        ("Exp3_2p5D_Loose", "Exp2_MIP_Loose"),
        ("Exp1_2D_Loose", "Exp2_MIP_Loose"),
    ] if a in names and b in names]
    write(paired_table(main_exps, main_boot, main_pts, pairs, "official_888"), "table_paired.csv")

    # ---------------- 3) Fold-wise analysis ---------------------------------
    print("\n[3] fold-wise CPM")
    fold_rows, fold_cpm = [], {}
    for e in main_exps:
        fc = stats.fold_wise_cpm(load_candidates(e.name), included, excluded, folds)
        fold_cpm[e.name] = fc
        for k, v in fc.items():
            fold_rows.append({"experiment": e.name, "label": e.label,
                              "fold": k, "subset": f"subset{k}", "cpm": round(v, 4)})
    write(pd.DataFrame(fold_rows), "table_foldwise.csv")

    ft_rows = []
    for a, b in pairs:
        if a in fold_cpm and b in fold_cpm:
            t = stats.paired_fold_test(fold_cpm[a], fold_cpm[b])
            t.update({"contrast": f"{a} - {b}"})
            ft_rows.append(t)
    if ft_rows:
        ft = pd.DataFrame(ft_rows)
        adj = stats.holm_bonferroni(dict(zip(ft["contrast"], ft["wilcoxon_p"])))
        ft["wilcoxon_p_holm"] = ft["contrast"].map(adj)
        write(ft, "table_foldwise_tests.csv")

    # ---------------- 4) Effect of annotations_excluded.csv -----------------
    print("\n[4] effect of applying annotations_excluded.csv")
    eff_rows = []
    for e in main_exps:
        cands = load_candidates(e.name)
        with_excl = ev.evaluate(ev.match(cands, included, excluded, uids888))
        without = ev.evaluate(ev.match(cands, included, None, uids888))
        legacy = ev.evaluate(ev.match(cands, included, excluded, uids888,
                                      excluded_policy="legacy_abs"))
        eff_rows.append({
            "experiment": e.name, "label": e.label,
            "cpm_without_excluded": round(without["cpm"], 4),
            "cpm_with_excluded_official": round(with_excl["cpm"], 4),
            "cpm_with_excluded_legacy_radius": round(legacy["cpm"], 4),
            "delta_cpm": round(with_excl["cpm"] - without["cpm"], 4),
            "fp_without_excluded": without["false_positives"],
            "fp_with_excluded": with_excl["false_positives"],
            "candidates_ignored": with_excl["ignored_on_excluded"],
            "pct_fp_removed": round(100 * (without["false_positives"] -
                                           with_excl["false_positives"]) /
                                    max(without["false_positives"], 1), 2),
        })
    write(pd.DataFrame(eff_rows), "table_excluded_effect.csv")

    # ---------------- 5) Size-stratified recall -----------------------------
    print("\n[5] size-stratified recall")
    size_rows = []
    for e in main_exps:
        cands = load_candidates(e.name)
        for op in (1.0, 4.0, None):
            df = stats.size_stratified_recall(cands, annotations, excluded, uids888,
                                              fp_per_scan=op)
            df.insert(0, "experiment", e.name)
            df.insert(1, "label", e.label)
            size_rows.append(df)
    write(pd.concat(size_rows, ignore_index=True), "table_size_recall.csv")

    # ---------------- 6) Sweeps (restricted common cohort) ------------------
    def sweep_table(sweep_exps, anchors, name, tag):
        exps = available(sweep_exps) + [e for e in main_exps if e.name in anchors]
        if len(exps) < 2:
            print(f"  (skipped {name}: not enough configurations)")
            return
        uids = cohort_for(exps, folds_by_index)
        df, boots, pts = evaluate_set(exps, uids, included, excluded, args.bootstrap, tag)
        write(df.sort_values("r_sample" if "r_sweep" in tag else "w_min_px"), name)
        anchor = "Exp4_2p5D_Strict"
        if anchor in boots:
            pr = [(e.name, anchor) for e in exps if e.name != anchor]
            write(paired_table(exps, boots, pts, pr, tag), name.replace(".csv", "_paired.csv"))

    print("\n[6] supervision-extent sweep")
    sweep_table(cfg.R_SWEEP_EXPERIMENTS, {"Exp4_2p5D_Strict", "Exp3_2p5D_Loose"},
                "table_r_sweep.csv", "r_sweep_cohort")

    print("\n[7] minimum-box-size sweep")
    sweep_table(cfg.WMIN_SWEEP_EXPERIMENTS, {"Exp4_2p5D_Strict"},
                "table_wmin_sweep.csv", "w_sweep_cohort")

    # ---------------- 6c) What actually predicts CPM across the design grid --
    print("\n[7b] label-geometry analysis across both sweeps")
    write(clamping_analysis(folds_by_index, included, excluded), "table_clamping.csv")

    # ---------------- 7) Negative mining ------------------------------------
    print("\n[8] negative-mining ablation")
    neg_exps = available(cfg.NEGATIVE_ABLATION_EXPERIMENTS + cfg.NEGATIVE_MATCHED_EXPERIMENTS)
    if neg_exps:
        exps = neg_exps + [e for e in main_exps
                           if e.name in {"Exp3_2p5D_Loose", "Exp4_2p5D_Strict"}]
        uids = cohort_for(exps, folds_by_index)
        df, boots, pts = evaluate_set(exps, uids, included, excluded, args.bootstrap, "neg_cohort")
        # Dropping the nodule-free scans removes their background slices too, so
        # the two arms differ in training-set *size* as well as in scan
        # diversity. Report the sizes rather than leaving the confound implicit.
        df = df.merge(training_set_sizes(), on="experiment", how="left")
        write(df, "table_negatives.csv")
        # all888 vs count-matched isolates scan provenance; count-matched vs
        # posonly isolates training-set size.
        pr = [("Exp3_2p5D_Loose", "NegMatch_Exp3_posonly12"),
              ("NegMatch_Exp3_posonly12", "NegAbl_Exp3_posonly"),
              ("Exp3_2p5D_Loose", "NegAbl_Exp3_posonly"),
              ("Exp4_2p5D_Strict", "NegAbl_Exp4_posonly")]
        write(paired_table(exps, boots, pts, [p for p in pr if p[0] in boots and p[1] in boots],
                           "neg_cohort"), "table_negatives_paired.csv")

    # ---------------- 8) Multi-seed -----------------------------------------
    print("\n[9] multi-seed repetition")
    seed_exps = available(cfg.SEED_EXPERIMENTS)
    if seed_exps:
        base = {"Exp3_2p5D_Loose": "Exp3 (2.5D loose)", "Exp4_2p5D_Strict": "Exp4 (2.5D strict)"}
        exps = seed_exps + [e for e in main_exps if e.name in base]
        uids = cohort_for(exps, folds_by_index)
        df, boots, pts = evaluate_set(exps, uids, included, excluded, 0, "seed_cohort")
        write(df, "table_seeds_raw.csv")

        summary = []
        for anchor, pretty in base.items():
            members = {e.name: e for e in exps
                       if (e.name == anchor) or (e.name.endswith(anchor.split("_", 1)[1])
                                                 and e.group == "seeds")}
            by_seed = {members[n].seed: float(df.loc[df.experiment == n, "cpm"].iloc[0])
                       for n in members if (df.experiment == n).any()}
            if len(by_seed) < 2:
                continue
            s = stats.seed_summary(by_seed)
            s.update({"configuration": pretty, "n_scans": len(uids),
                      "cpm_by_seed": json.dumps({str(k): round(v, 4) for k, v in by_seed.items()})})
            summary.append(s)
        write(pd.DataFrame(summary), "table_seeds.csv")

    # ---------------- 9) YOLO26 ---------------------------------------------
    print("\n[10] YOLO26n comparison")
    y26 = available(cfg.YOLO26_EXPERIMENTS)
    if y26:
        exps = y26 + [e for e in main_exps
                      if e.name in {"Exp3_2p5D_Loose", "Exp4_2p5D_Strict"}]
        uids = cohort_for(exps, folds_by_index)
        df, boots, pts = evaluate_set(exps, uids, included, excluded, args.bootstrap, "yolo26_cohort")
        write(df, "table_yolo26.csv")
        pr = [("Y26_Exp4_2p5D_Strict", "Exp4_2p5D_Strict"),
              ("Y26_Exp3_2p5D_Loose", "Exp3_2p5D_Loose"),
              ("Y26_Exp4_2p5D_Strict", "Y26_Exp3_2p5D_Loose")]
        write(paired_table(exps, boots, pts, [p for p in pr if p[0] in boots and p[1] in boots],
                           "yolo26_cohort"), "table_yolo26_paired.csv")

        # Curves for Figure 8 must come from the cohort its CPM values come
        # from. Scoring them on all 888 scans would divide by nodules the
        # five-fold runs never see, capping sensitivity near one half.
        rows = []
        for e in exps:
            fps, sens = pts[e.name]["curve_fps"], pts[e.name]["curve_sens"]
            keep = fps <= 8.0
            beyond = np.flatnonzero(fps > 8.0)
            if beyond.size:
                keep = keep.copy()
                keep[beyond[0]] = True
            rows.append(pd.DataFrame({"experiment": e.name, "label": e.label,
                                      "fp_per_scan": fps[keep], "sensitivity": sens[keep]}))
        write(pd.concat(rows, ignore_index=True), "froc_curves_yolo26.csv")

    print(f"\nAll tables written to {R}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())