scripts/bootstrap_clinical_all.py

"""Bootstrap 95% CI for clinical probing (AUROC) across ALL 6 encoders.

Runs clinical-only bootstrap (no Stage-1) for each encoder:
  1. Load checkpoint, extract embeddings from train/val/test clinical splits
  2. Train logistic probe on train/val (with C sweep)
  3. Bootstrap 1000 resamples over test set for 3-way and AD-vs-CN AUROC

Usage:
    cd /data/Albus/Brain
    CUDA_VISIBLE_DEVICES=4 python scripts/bootstrap_clinical_all.py
"""

from __future__ import annotations
import gc
import sys
import time
from pathlib import Path

import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent))

from bootstrap_ci import (
    extract_embeddings,
    _subset_cls,
    train_probe,
    bootstrap_clinical_auroc,
)
from train_pet_foundation import PETSUVRFoundationModel, build_encoder
from pet_vlm_dataset import PETSUVRDataset


ENCODERS = [
    ("ReMAP-PET",           "runs/foundation/medicalnet_layer4_regalign_best.pt"),
    ("MedicalNet frozen",   "runs/foundation/medicalnet_frozen_mlp.pt"),
    ("BrainIAC frozen",     "runs/foundation/brainiac_frozen_mlp.pt"),
    ("BrainFM frozen",      "runs/foundation/brainfm_frozen_mlp_b4_best.pt"),
    ("SAM-Med3D frozen",    "runs/foundation/sam_med3d_frozen_mlp_best.pt"),
    ("SwinUNETR frozen",    "runs/foundation/swinunetr_frozen_mlp_best.pt"),
]

TRAIN_CSV = Path("data/metadata/splits/train_clinical_server.csv")
VAL_CSV   = Path("data/metadata/splits/val_clinical_server.csv")
TEST_CSV  = Path("data/metadata/splits/test_clinical_server.csv")
TEST_MANIFEST = Path("metadata/splits/test.csv")

B = 1000
SEED = 42
BATCH_SIZE = 4
NUM_WORKERS = 2


def load_model(ckpt_path: Path, device: torch.device):
    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
    saved = ckpt.get("args", {})

    class _A:
        pass
    a = _A()
    a.backbone = saved.get("backbone", "medicalnet")
    a.medicalnet_weights = Path(saved.get("medicalnet_weights",
                                          "pretrained/medicalnet/resnet_50_23dataset.pth"))
    a.brainiac_weights = Path(saved.get("brainiac_weights",
                                        "pretrained/brainiac/backbone.safetensors"))
    a.brainfm_weights = Path("pretrained/brainfm/assets/brainfm_pretrained.pth")
    a.brainfm_code_root = Path("pretrained/brainfm")
    a.swinunetr_weights = Path("pretrained/swinunetr/model_swinvit.pt")
    a.sam_med3d_weights = Path("pretrained/sam-med3d/sam_med3d_turbo.pth")
    a.output_size = tuple(saved.get("output_size", (96, 96, 96)))
    embed_dim = saved.get("embed_dim", 256)
    freeze = bool(saved.get("freeze_encoder", False))

    ds_tmp = PETSUVRDataset(TEST_MANIFEST, output_size=a.output_size)
    n_regions = int(ds_tmp[0]["suvr"].numel())
    encoder = build_encoder(a)
    model = PETSUVRFoundationModel(encoder, n_regions, embed_dim, freeze).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()
    return model, a.output_size


def run_one(name: str, ckpt_path: str, device: torch.device):
    print(f"\n{'='*60}", flush=True)
    print(f"  {name}  ({ckpt_path})", flush=True)
    print(f"{'='*60}", flush=True)
    t0 = time.time()

    model, output_size = load_model(Path(ckpt_path), device)

    # extract embeddings
    train_df, x_train = extract_embeddings(
        model, TRAIN_CSV, output_size, BATCH_SIZE, NUM_WORKERS, device)
    val_df, x_val = extract_embeddings(
        model, VAL_CSV, output_size, BATCH_SIZE, NUM_WORKERS, device)
    test_df, x_test = extract_embeddings(
        model, TEST_CSV, output_size, BATCH_SIZE, NUM_WORKERS, device)

    results = {}

    # --- 3-way ---
    x_tr3, y_tr3 = _subset_cls(train_df, x_train, "clinical_label", ["CN", "MCI", "AD"])
    x_v3,  y_v3  = _subset_cls(val_df,   x_val,   "clinical_label", ["CN", "MCI", "AD"])
    x_te3, y_te3 = _subset_cls(test_df,  x_test,  "clinical_label", ["CN", "MCI", "AD"])
    print(f"  3-way:  train={len(y_tr3)}  val={len(y_v3)}  test={len(y_te3)}", flush=True)
    probe3, enc3 = train_probe(x_tr3, y_tr3, x_v3, y_v3)
    pt3, lo3, hi3 = bootstrap_clinical_auroc(
        probe3, enc3, x_tr3, y_tr3, x_v3, y_v3, x_te3, y_te3, B=B, seed=SEED)
    results["3way"] = (pt3, lo3, hi3)
    print(f"  3-way AUROC:   {pt3:.4f}   95% CI [{lo3:.4f}, {hi3:.4f}]", flush=True)

    # --- AD vs CN ---
    x_tr2, y_tr2 = _subset_cls(train_df, x_train, "clinical_label", ["CN", "AD"])
    x_v2,  y_v2  = _subset_cls(val_df,   x_val,   "clinical_label", ["CN", "AD"])
    x_te2, y_te2 = _subset_cls(test_df,  x_test,  "clinical_label", ["CN", "AD"])
    print(f"  AD/CN:  train={len(y_tr2)}  val={len(y_v2)}  test={len(y_te2)}", flush=True)
    probe2, enc2 = train_probe(x_tr2, y_tr2, x_v2, y_v2)
    pt2, lo2, hi2 = bootstrap_clinical_auroc(
        probe2, enc2, x_tr2, y_tr2, x_v2, y_v2, x_te2, y_te2, B=B, seed=SEED)
    results["adcn"] = (pt2, lo2, hi2)
    print(f"  AD/CN AUROC:   {pt2:.4f}   95% CI [{lo2:.4f}, {hi2:.4f}]", flush=True)

    elapsed = time.time() - t0
    print(f"  (elapsed {elapsed:.0f}s)", flush=True)

    # free GPU memory
    del model
    gc.collect()
    torch.cuda.empty_cache()

    return results


def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}", flush=True)

    all_results = {}
    for name, ckpt in ENCODERS:
        all_results[name] = run_one(name, ckpt, device)

    # --- summary table ---
    print(f"\n\n{'='*80}", flush=True)
    print("  SUMMARY: Clinical Probing Bootstrap 95% CI (B=1000)", flush=True)
    print(f"{'='*80}", flush=True)
    print(f"{'Model':<22s}  {'3-way AUROC':>22s}  {'AD/CN AUROC':>22s}", flush=True)
    print("-" * 72, flush=True)
    for name, res in all_results.items():
        pt3, lo3, hi3 = res["3way"]
        pt2, lo2, hi2 = res["adcn"]
        hw3 = (hi3 - lo3) / 2
        hw2 = (hi2 - lo2) / 2
        print(f"{name:<22s}  {pt3:.4f} +/- {hw3:.4f}  {pt2:.4f} +/- {hw2:.4f}", flush=True)
    print(f"{'='*80}", flush=True)
    print("Done.", flush=True)


if __name__ == "__main__":
    main()