scripts/bootstrap_ci.py

"""Bootstrap 95 % confidence intervals for ReMAP-PET key metrics.

Stage-1 metrics (153 test subjects):
    - SUVR MAE
    - Pearson r (voxel-level across all subjects x regions)
    - PET->SUVR Recall@1  (retrieval)

Clinical probe metrics:
    - AD vs CN AUROC  (logistic regression on PET embeddings)
    - 3-way (CN/MCI/AD) AUROC

Usage (from /data/Albus/Brain):
    CUDA_VISIBLE_DEVICES=1 python scripts/bootstrap_ci.py
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import balanced_accuracy_score, roc_auc_score
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import LabelEncoder, StandardScaler, label_binarize

# -- project imports (scripts/ is the working dir's sibling) -----------------
sys.path.insert(0, str(Path(__file__).resolve().parent))
from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr
from train_pet_foundation import PETSUVRFoundationModel, build_encoder


# ---------------------------------------------------------------------------
# helpers copied from evaluate_pet_foundation.py
# ---------------------------------------------------------------------------

def _pearson_flat(pred: np.ndarray, target: np.ndarray) -> float:
    p = pred.reshape(-1)
    t = target.reshape(-1)
    if p.std() < 1e-8 or t.std() < 1e-8:
        return float("nan")
    return float(np.corrcoef(p, t)[0, 1])


def _retrieval_recall_at_1(logits: np.ndarray) -> float:
    ranks = []
    for i in range(logits.shape[0]):
        order = np.argsort(-logits[i])
        rank = int(np.where(order == i)[0][0]) + 1
        ranks.append(rank)
    return float(np.mean(np.asarray(ranks) <= 1))


# ---------------------------------------------------------------------------
# Stage-1: forward pass -> per-subject arrays
# ---------------------------------------------------------------------------

@torch.no_grad()
def collect_stage1(
    model: PETSUVRFoundationModel,
    loader: DataLoader,
    device: torch.device,
) -> dict[str, np.ndarray]:
    """Return pred_suvr, target_suvr, pet_z, suvr_z  (all numpy, N-first)."""
    model.eval()
    pred_chunks, target_chunks = [], []
    pet_z_chunks, suvr_z_chunks = [], []

    for batch in loader:
        image = batch["image"].to(device, non_blocking=True)
        suvr = batch["suvr"].to(device, non_blocking=True)
        outputs = model(image, suvr)
        pred_chunks.append(outputs["pred_suvr"].cpu().numpy())
        target_chunks.append(suvr.cpu().numpy())

        pet_feat = model.pet_encoder(image)
        pet_z = F.normalize(model.pet_projector(pet_feat), dim=-1)
        suvr_z = F.normalize(model.suvr_encoder(suvr), dim=-1)
        pet_z_chunks.append(pet_z.cpu().numpy())
        suvr_z_chunks.append(suvr_z.cpu().numpy())

    return {
        "pred":   np.concatenate(pred_chunks,   axis=0),
        "target": np.concatenate(target_chunks, axis=0),
        "pet_z":  np.concatenate(pet_z_chunks,  axis=0),
        "suvr_z": np.concatenate(suvr_z_chunks, axis=0),
    }


def stage1_metrics(d: dict[str, np.ndarray], idx: np.ndarray) -> dict[str, float]:
    """Compute stage-1 metrics on a subset given by *idx*.

    MAE and Pearson work fine with duplicate indices (bootstrap).
    For retrieval R@1 we need unique subjects (duplicates would make the
    diagonal ground-truth ambiguous), so we deduplicate *idx* first.
    """
    pred   = d["pred"][idx]
    target = d["target"][idx]

    # retrieval: use unique indices only
    uid = np.unique(idx)
    pet_z  = d["pet_z"][uid]
    suvr_z = d["suvr_z"][uid]
    logits = pet_z @ suvr_z.T

    return {
        "mae":     float(np.mean(np.abs(pred - target))),
        "pearson": _pearson_flat(pred, target),
        "pet_suvr_r1": _retrieval_recall_at_1(logits),
    }


# ---------------------------------------------------------------------------
# Clinical: extract embeddings, train probe, evaluate
# ---------------------------------------------------------------------------

@torch.no_grad()
def extract_embeddings(
    model: PETSUVRFoundationModel,
    manifest: Path,
    output_size: tuple[int, int, int],
    batch_size: int,
    num_workers: int,
    device: torch.device,
) -> tuple[pd.DataFrame, np.ndarray]:
    dataset = PETSUVRDataset(manifest, output_size=output_size)
    loader  = DataLoader(dataset, batch_size=batch_size, shuffle=False,
                         num_workers=num_workers, collate_fn=collate_pet_suvr)
    feats = []
    model.eval()
    for batch in loader:
        image = batch["image"].to(device, non_blocking=True)
        pet_feat = model.pet_encoder(image)
        pet_z = F.normalize(model.pet_projector(pet_feat), dim=-1)
        feats.append(pet_z.cpu().numpy())
    return pd.read_csv(manifest), np.concatenate(feats, axis=0)


def _subset_cls(df, x, column, labels):
    mask = df[column].isin(labels).to_numpy()
    return x[mask], df.loc[mask, column].astype(str).to_numpy()


def train_probe(x_train, y_train, x_val, y_val):
    """Train logistic probe with C sweep; return best model + encoder."""
    enc = LabelEncoder()
    enc.fit(np.concatenate([y_train, y_val]))
    y_tr = enc.transform(y_train)
    y_v  = enc.transform(y_val)
    best_m, best_s = None, -np.inf
    for c in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0]:
        m = make_pipeline(StandardScaler(),
                          LogisticRegression(C=c, max_iter=5000,
                                            class_weight="balanced"))
        m.fit(x_train, y_tr)
        s = balanced_accuracy_score(y_v, m.predict(x_val))
        if s > best_s:
            best_m, best_s = m, s
    return best_m, enc


def clinical_auroc(model_probe, encoder, x_test, y_test):
    """Return AUROC (binary or macro-OVR)."""
    y_int = encoder.transform(y_test)
    proba = model_probe.predict_proba(x_test)
    if len(encoder.classes_) == 2:
        return roc_auc_score(y_int, proba[:, 1])
    else:
        y_bin = label_binarize(y_int, classes=np.arange(len(encoder.classes_)))
        return roc_auc_score(y_bin, proba, average="macro", multi_class="ovr")


# ---------------------------------------------------------------------------
# Bootstrap
# ---------------------------------------------------------------------------

def bootstrap_ci(
    metric_fn,
    n: int,
    B: int = 1000,
    seed: int = 42,
    alpha: float = 0.05,
) -> tuple[float, float, float]:
    """
    metric_fn(idx) -> float   where idx is array of resampled indices.
    Returns (point_estimate, lo, hi) for the (1-alpha) CI.
    """
    rng = np.random.RandomState(seed)
    all_idx = np.arange(n)
    point = metric_fn(all_idx)
    boots = np.empty(B)
    for b in range(B):
        idx = rng.choice(n, size=n, replace=True)
        boots[b] = metric_fn(idx)
    lo = float(np.percentile(boots, 100 * alpha / 2))
    hi = float(np.percentile(boots, 100 * (1 - alpha / 2)))
    return point, lo, hi


def bootstrap_clinical_auroc(
    probe, encoder,
    x_train, y_train_raw,
    x_val, y_val_raw,
    x_test, y_test_raw,
    B: int = 1000,
    seed: int = 42,
    alpha: float = 0.05,
) -> tuple[float, float, float]:
    """
    Bootstrap over the *test* set only (probe is fixed).
    """
    rng = np.random.RandomState(seed)
    n = len(y_test_raw)
    all_idx = np.arange(n)
    point = clinical_auroc(probe, encoder, x_test, y_test_raw)

    boots = np.empty(B)
    for b in range(B):
        idx = rng.choice(n, size=n, replace=True)
        try:
            boots[b] = clinical_auroc(probe, encoder, x_test[idx], y_test_raw[idx])
        except ValueError:
            # can happen if a resample has only one class
            boots[b] = np.nan
    boots = boots[~np.isnan(boots)]
    lo = float(np.percentile(boots, 100 * alpha / 2))
    hi = float(np.percentile(boots, 100 * (1 - alpha / 2)))
    return point, lo, hi


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path,
                        default=Path("runs/foundation/medicalnet_layer4_regalign_best.pt"))
    parser.add_argument("--test-manifest", type=Path,
                        default=Path("metadata/splits/test.csv"))
    parser.add_argument("--train-clinical", type=Path,
                        default=Path("data/metadata/splits/train_clinical_server.csv"))
    parser.add_argument("--val-clinical", type=Path,
                        default=Path("data/metadata/splits/val_clinical_server.csv"))
    parser.add_argument("--test-clinical", type=Path,
                        default=Path("data/metadata/splits/test_clinical_server.csv"))
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--B", type=int, default=1000, help="bootstrap resamples")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    # ---- load model -------------------------------------------------------
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    saved = ckpt.get("args", {})

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

    output_size = margs.output_size

    # build model
    dataset_tmp = PETSUVRDataset(args.test_manifest, output_size=output_size)
    n_regions = int(dataset_tmp[0]["suvr"].numel())
    encoder = build_encoder(margs)
    model = PETSUVRFoundationModel(encoder, n_regions, embed_dim, freeze_encoder).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()
    print(f"Loaded checkpoint: {args.checkpoint}", flush=True)
    print(f"backbone={margs.backbone}  embed_dim={embed_dim}  "
          f"freeze={freeze_encoder}  output_size={output_size}", flush=True)

    # ===== STAGE 1 =========================================================
    print("\n===== Stage-1 evaluation (test set) =====", flush=True)
    test_ds = PETSUVRDataset(args.test_manifest, output_size=output_size)
    test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False,
                             num_workers=args.num_workers, collate_fn=collate_pet_suvr)
    d = collect_stage1(model, test_loader, device)
    N = d["pred"].shape[0]
    print(f"  N = {N}", flush=True)

    for name in ("mae", "pearson", "pet_suvr_r1"):
        fn = lambda idx, _n=name: stage1_metrics(d, idx)[_n]
        pt, lo, hi = bootstrap_ci(fn, N, B=args.B, seed=args.seed)
        print(f"  {name:20s}  {pt:.4f}   95% CI [{lo:.4f}, {hi:.4f}]", flush=True)

    # ===== CLINICAL ========================================================
    print("\n===== Clinical downstream probes =====", flush=True)
    train_df, x_train_all = extract_embeddings(
        model, args.train_clinical, output_size, args.batch_size, args.num_workers, device)
    val_df, x_val_all = extract_embeddings(
        model, args.val_clinical, output_size, args.batch_size, args.num_workers, device)
    test_df, x_test_all = extract_embeddings(
        model, args.test_clinical, output_size, args.batch_size, args.num_workers, device)

    # ---- AD vs CN ---------------------------------------------------------
    print("\n  -- AD vs CN --", flush=True)
    x_tr, y_tr = _subset_cls(train_df, x_train_all, "clinical_label", ["CN", "AD"])
    x_v,  y_v  = _subset_cls(val_df,   x_val_all,   "clinical_label", ["CN", "AD"])
    x_te, y_te = _subset_cls(test_df,  x_test_all,  "clinical_label", ["CN", "AD"])
    print(f"     train={len(y_tr)}  val={len(y_v)}  test={len(y_te)}", flush=True)
    probe_ad, enc_ad = train_probe(x_tr, y_tr, x_v, y_v)
    pt, lo, hi = bootstrap_clinical_auroc(
        probe_ad, enc_ad, x_tr, y_tr, x_v, y_v, x_te, y_te,
        B=args.B, seed=args.seed)
    print(f"     {'ad_vs_cn_auroc':20s}  {pt:.4f}   95% CI [{lo:.4f}, {hi:.4f}]", flush=True)

    # ---- 3-way CN / MCI / AD ---------------------------------------------
    print("\n  -- 3-way (CN / MCI / AD) --", flush=True)
    x_tr3, y_tr3 = _subset_cls(train_df, x_train_all, "clinical_label", ["CN", "MCI", "AD"])
    x_v3,  y_v3  = _subset_cls(val_df,   x_val_all,   "clinical_label", ["CN", "MCI", "AD"])
    x_te3, y_te3 = _subset_cls(test_df,  x_test_all,  "clinical_label", ["CN", "MCI", "AD"])
    print(f"     train={len(y_tr3)}  val={len(y_v3)}  test={len(y_te3)}", flush=True)
    probe_3w, enc_3w = train_probe(x_tr3, y_tr3, x_v3, y_v3)
    pt, lo, hi = bootstrap_clinical_auroc(
        probe_3w, enc_3w, x_tr3, y_tr3, x_v3, y_v3, x_te3, y_te3,
        B=args.B, seed=args.seed)
    print(f"     {'3way_auroc':20s}  {pt:.4f}   95% CI [{lo:.4f}, {hi:.4f}]", flush=True)

    print("\nDone.", flush=True)


if __name__ == "__main__":
    main()