code/run_final_experiments.py

"""Unified pipeline for the corrected thesis experiments.

Given a holdout split manifest (build_holdout_split.py), this script:
  1. Runs stratified 5-fold CV on the train+val pool for every selected model.
  2. After CV, retrains the model on the full train+val pool.
  3. Evaluates on the held-out independent test set and stores per-sample
     predictions for paired statistical testing.

Designed to run unattended on the Azure T4 VM. All artefacts go to
``output_dir`` so the user can inspect them after the long run finishes.
"""

from __future__ import annotations

import argparse
import copy
import json
import random
import time
from collections import defaultdict
from pathlib import Path
from typing import Sequence

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from sklearn.metrics import (
    accuracy_score,
    average_precision_score,
    confusion_matrix,
    f1_score,
    precision_recall_fscore_support,
    roc_auc_score,
)
from sklearn.model_selection import StratifiedKFold
from torch.utils.data import DataLoader, Dataset, WeightedRandomSampler
from torchvision import models, transforms


IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
INCEPTION_SIZE = 299


class ImageListDataset(Dataset):
    def __init__(self, root: Path, samples: Sequence[tuple[str, int]], transform):
        self.root = root
        self.samples = samples
        self.transform = transform

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        rel_path, label = self.samples[idx]
        with Image.open(self.root / rel_path) as img:
            img = img.convert("RGB")
            if self.transform is not None:
                img = self.transform(img)
        return img, label


class OpenCLIPClassifier(nn.Module):
    def __init__(self, num_classes, model_name="ViT-B-16", pretrained="openai"):
        super().__init__()
        import open_clip

        clip_model, _, _ = open_clip.create_model_and_transforms(model_name, pretrained=pretrained)
        self.backbone = clip_model
        with torch.no_grad():
            dummy = torch.zeros(1, 3, 224, 224)
            feat_dim = self.backbone.encode_image(dummy).shape[-1]
        self.head = nn.Linear(feat_dim, num_classes)

    def forward(self, x):
        feats = self.backbone.encode_image(x)
        return self.head(feats.float())


def build_transforms(image_size, use_clip_norm=False):
    mean = CLIP_MEAN if use_clip_norm else IMAGENET_MEAN
    std = CLIP_STD if use_clip_norm else IMAGENET_STD
    train_tf = transforms.Compose([
        transforms.RandomRotation(30),
        transforms.RandomHorizontalFlip(),
        transforms.RandomResizedCrop((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])
    eval_tf = transforms.Compose([
        transforms.Resize((image_size, image_size)),
        transforms.ToTensor(),
        transforms.Normalize(mean, std),
    ])
    return train_tf, eval_tf


def build_model(name, num_classes):
    if name == "vgg19":
        m = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1)
        m.classifier[-1] = nn.Linear(m.classifier[-1].in_features, num_classes)
        return m, 224, False
    if name == "resnet50":
        m = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        m.fc = nn.Linear(m.fc.in_features, num_classes)
        return m, 224, False
    if name == "resnet101":
        m = models.resnet101(weights=models.ResNet101_Weights.IMAGENET1K_V2)
        m.fc = nn.Linear(m.fc.in_features, num_classes)
        return m, 224, False
    if name == "densenet121":
        m = models.densenet121(weights=models.DenseNet121_Weights.IMAGENET1K_V1)
        m.classifier = nn.Linear(m.classifier.in_features, num_classes)
        return m, 224, False
    if name == "inception_v3":
        m = models.inception_v3(weights=models.Inception_V3_Weights.IMAGENET1K_V1, aux_logits=True)
        m.fc = nn.Linear(m.fc.in_features, num_classes)
        m.AuxLogits.fc = nn.Linear(m.AuxLogits.fc.in_features, num_classes)
        return m, INCEPTION_SIZE, False
    if name == "clip_openai":
        return OpenCLIPClassifier(num_classes), 224, True
    raise ValueError(name)


def extract_logits_loss(out, labels, criterion):
    if isinstance(out, tuple):
        logits = out[0]
        loss = criterion(logits, labels) + 0.4 * criterion(out[1], labels)
        return logits, loss
    return out, criterion(out, labels)


def train_epoch(model, loader, criterion, optim, scaler, device):
    model.train()
    total, count = 0.0, 0
    for x, y in loader:
        x = x.to(device, non_blocking=True)
        y = y.to(device, non_blocking=True)
        optim.zero_grad(set_to_none=True)
        with torch.amp.autocast("cuda", enabled=scaler.is_enabled()):
            out = model(x)
            _, loss = extract_logits_loss(out, y, criterion)
        scaler.scale(loss).backward()
        scaler.step(optim)
        scaler.update()
        total += loss.item() * x.size(0)
        count += x.size(0)
    return total / max(count, 1)


@torch.no_grad()
def evaluate(model, loader, device, return_preds=False, num_classes=None):
    model.eval()
    preds, labels, probs = [], [], []
    for x, y in loader:
        x = x.to(device, non_blocking=True)
        out = model(x)
        logits = out[0] if isinstance(out, tuple) else out
        p = F.softmax(logits, dim=1)
        preds.extend(logits.argmax(1).cpu().tolist())
        probs.extend(p.cpu().numpy().tolist())
        labels.extend(y.tolist())
    p_m, r_m, f_m, _ = precision_recall_fscore_support(labels, preds, average="macro", zero_division=0)
    acc = accuracy_score(labels, preds)
    res = {"accuracy": acc, "precision_macro": p_m, "recall_macro": r_m, "f1_macro": f_m}
    # ROC-AUC and PR-AUC (one-vs-rest, macro)
    try:
        n = num_classes or (max(labels) + 1)
        y_onehot = np.eye(n)[np.array(labels)]
        probs_arr = np.array(probs)
        res["roc_auc_macro"] = float(roc_auc_score(y_onehot, probs_arr, average="macro", multi_class="ovr"))
        res["pr_auc_macro"] = float(average_precision_score(y_onehot, probs_arr, average="macro"))
    except Exception as exc:
        res["roc_auc_macro"] = None
        res["pr_auc_macro"] = None
    # Expected Calibration Error (15 bins)
    res["ece"] = float(expected_calibration_error(np.array(probs), np.array(labels)))
    if return_preds:
        res["preds"] = preds
        res["labels"] = labels
        res["probs"] = probs
    return res


def expected_calibration_error(probs, labels, n_bins=15):
    confidences = probs.max(axis=1)
    predictions = probs.argmax(axis=1)
    accuracies = (predictions == labels).astype(float)
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for lo, hi in zip(bin_boundaries[:-1], bin_boundaries[1:]):
        in_bin = (confidences > lo) & (confidences <= hi)
        if in_bin.sum() > 0:
            avg_conf = confidences[in_bin].mean()
            avg_acc = accuracies[in_bin].mean()
            ece += (in_bin.sum() / len(probs)) * abs(avg_conf - avg_acc)
    return ece


def bootstrap_ci(labels, preds, metric_fn, n_resamples=1000, alpha=0.05, seed=0):
    rng = np.random.RandomState(seed)
    labels = np.array(labels)
    preds = np.array(preds)
    stats = []
    n = len(labels)
    for _ in range(n_resamples):
        idx = rng.randint(0, n, size=n)
        stats.append(metric_fn(labels[idx], preds[idx]))
    stats = np.array(stats)
    return {"mean": float(stats.mean()), "lo": float(np.quantile(stats, alpha / 2)),
            "hi": float(np.quantile(stats, 1 - alpha / 2))}


def compute_class_weights(labels, num_classes, smoothing="sqrt"):
    counts = np.bincount(labels, minlength=num_classes).astype(float)
    counts[counts == 0] = 1.0  # avoid div0
    if smoothing == "sqrt":
        weights = 1.0 / np.sqrt(counts)
    else:
        weights = 1.0 / counts
    weights = weights * num_classes / weights.sum()
    return torch.tensor(weights, dtype=torch.float32)


def class_names_to_int(samples, classes):
    cls_to_idx = {c: i for i, c in enumerate(classes)}
    return [(p, cls_to_idx[c]) for p, c in samples]


def train_model(model_name, train_samples, val_samples, classes, root, args, device, log):
    model_tuple = build_model(model_name, len(classes))
    model, image_size, use_clip_norm = model_tuple
    model = model.to(device)
    train_tf, eval_tf = build_transforms(image_size, use_clip_norm)
    train_set = ImageListDataset(root, train_samples, train_tf)
    val_set = ImageListDataset(root, val_samples, eval_tf)
    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True,
                              num_workers=args.workers, pin_memory=device.type == "cuda")
    val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False,
                            num_workers=args.workers, pin_memory=device.type == "cuda")
    # Class-weighted cross-entropy to mitigate severe imbalance
    train_labels = [s[1] for s in train_samples]
    cls_weights = compute_class_weights(train_labels, len(classes)).to(device) if args.class_weighted else None
    criterion = nn.CrossEntropyLoss(weight=cls_weights)
    optim = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-4)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=args.epochs)
    scaler = torch.amp.GradScaler("cuda", enabled=device.type == "cuda")

    best_state, best_acc, best_epoch = None, 0.0, 0
    bad_epochs = 0
    for ep in range(1, args.epochs + 1):
        tloss = train_epoch(model, train_loader, criterion, optim, scaler, device)
        scheduler.step()
        v = evaluate(model, val_loader, device, num_classes=len(classes))
        line = (f"{model_name} ep {ep:03d} train_loss={tloss:.4f} "
                f"val_acc={v['accuracy']:.4f} val_f1={v['f1_macro']:.4f}")
        log(line)
        if v["accuracy"] > best_acc:
            best_acc = v["accuracy"]
            best_epoch = ep
            best_state = copy.deepcopy(model.state_dict())
            bad_epochs = 0
        else:
            bad_epochs += 1
            if bad_epochs >= args.patience:
                log(f"  early stop at ep {ep} (best ep {best_epoch} acc {best_acc:.4f})")
                break
    if best_state is not None:
        model.load_state_dict(best_state)
    return model, {"best_val_acc": best_acc, "best_epoch": best_epoch, "image_size": image_size,
                   "use_clip_norm": use_clip_norm}


def run_kfold(model_name, samples_pool, classes, root, args, device, log):
    paths = [s[0] for s in samples_pool]
    labels = [s[1] for s in samples_pool]
    skf = StratifiedKFold(n_splits=args.folds, shuffle=True, random_state=args.seed)
    fold_metrics = []
    for k, (tr_idx, vl_idx) in enumerate(skf.split(paths, labels), start=1):
        log(f"=== {model_name} fold {k}/{args.folds} ===")
        tr = [samples_pool[i] for i in tr_idx]
        vl = [samples_pool[i] for i in vl_idx]
        model, info = train_model(model_name, tr, vl, classes, root, args, device, log)
        # final fold metrics on val
        _, eval_tf = build_transforms(info["image_size"], info["use_clip_norm"])
        v_loader = DataLoader(
            ImageListDataset(root, vl, eval_tf), batch_size=args.batch_size,
            shuffle=False, num_workers=args.workers, pin_memory=device.type == "cuda",
        )
        m = evaluate(model, v_loader, device, return_preds=True, num_classes=len(classes))
        fold_metrics.append({"fold": k, **{k2: v for k2, v in m.items() if k2 not in ("preds", "labels", "probs")}})
        log(f"  fold {k} acc={m['accuracy']:.4f} f1={m['f1_macro']:.4f}")
        del model
        torch.cuda.empty_cache()
    accs = np.array([f["accuracy"] for f in fold_metrics])
    f1s = np.array([f["f1_macro"] for f in fold_metrics])
    summary = {
        "model": model_name,
        "folds": args.folds,
        "accuracy_mean": float(accs.mean()),
        "accuracy_std": float(accs.std(ddof=1)),
        "accuracy_ci95": [float(accs.mean() - 1.96 * accs.std(ddof=1) / np.sqrt(len(accs))),
                          float(accs.mean() + 1.96 * accs.std(ddof=1) / np.sqrt(len(accs)))],
        "f1_macro_mean": float(f1s.mean()),
        "f1_macro_std": float(f1s.std(ddof=1)),
        "fold_metrics": fold_metrics,
    }
    return summary


def run_indep_test(model_name, samples_pool, test_samples, classes, root, args, device, log, weights_dir):
    log(f"=== {model_name} FINAL (train on pool, eval on indep test) ===")
    # 90/10 split within the pool to keep an internal val for early stopping
    pool_paths = [s[0] for s in samples_pool]
    pool_labels = [s[1] for s in samples_pool]
    rng = np.random.RandomState(args.seed)
    idx = np.arange(len(samples_pool))
    rng.shuffle(idx)
    cut = int(0.9 * len(idx))
    tr = [samples_pool[i] for i in idx[:cut]]
    vl = [samples_pool[i] for i in idx[cut:]]
    model, info = train_model(model_name, tr, vl, classes, root, args, device, log)
    weights_path = weights_dir / f"{model_name}_final.pth"
    torch.save(model.state_dict(), weights_path)
    log(f"  saved weights: {weights_path}")
    _, eval_tf = build_transforms(info["image_size"], info["use_clip_norm"])
    test_loader = DataLoader(
        ImageListDataset(root, test_samples, eval_tf), batch_size=args.batch_size,
        shuffle=False, num_workers=args.workers, pin_memory=device.type == "cuda",
    )
    m = evaluate(model, test_loader, device, return_preds=True, num_classes=len(classes))
    cm = confusion_matrix(m["labels"], m["preds"], labels=list(range(len(classes))))
    # Bootstrap 95% CIs on accuracy and macro-F1
    boot_acc = bootstrap_ci(m["labels"], m["preds"], lambda y, p: float((y == p).mean()), seed=args.seed)
    boot_f1 = bootstrap_ci(m["labels"], m["preds"], lambda y, p: f1_score(y, p, average="macro", zero_division=0), seed=args.seed + 1)
    log(f"  test acc={m['accuracy']:.4f} f1={m['f1_macro']:.4f} "
        f"roc_auc={m.get('roc_auc_macro')} ece={m['ece']:.4f}")
    del model
    torch.cuda.empty_cache()
    return {
        "model": model_name,
        "test_accuracy": m["accuracy"],
        "test_precision_macro": m["precision_macro"],
        "test_recall_macro": m["recall_macro"],
        "test_f1_macro": m["f1_macro"],
        "test_roc_auc_macro": m.get("roc_auc_macro"),
        "test_pr_auc_macro": m.get("pr_auc_macro"),
        "test_ece": m["ece"],
        "bootstrap_accuracy_ci95": boot_acc,
        "bootstrap_f1_macro_ci95": boot_f1,
        "preds": m["preds"],
        "labels": m["labels"],
        "probs": m["probs"],
        "confusion_matrix": cm.tolist(),
    }


def mcnemar_pairwise(test_results):
    """Compute pairwise McNemar p-values between models on the independent test set."""
    from scipy.stats import binom
    out = {}
    names = sorted(test_results.keys())
    for i in range(len(names)):
        for j in range(i + 1, len(names)):
            a, b = names[i], names[j]
            preds_a = np.array(test_results[a]["preds"])
            preds_b = np.array(test_results[b]["preds"])
            labels = np.array(test_results[a]["labels"])
            ca = preds_a == labels
            cb = preds_b == labels
            b_count = int(np.sum(ca & ~cb))   # a right, b wrong
            c_count = int(np.sum(~ca & cb))   # a wrong, b right
            n = b_count + c_count
            if n == 0:
                p = 1.0
            else:
                k = min(b_count, c_count)
                p = float(2 * binom.cdf(k, n, 0.5))
                p = min(p, 1.0)
            out[f"{a}_vs_{b}"] = {"b": b_count, "c": c_count, "p_value": p}
    return out


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--manifest", default="holdout_split.json")
    p.add_argument("--models", nargs="+", default=[
        "vgg19", "resnet50", "resnet101", "densenet121", "inception_v3", "clip_openai",
    ])
    p.add_argument("--epochs", type=int, default=60)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--workers", type=int, default=4)
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--folds", type=int, default=5)
    p.add_argument("--patience", type=int, default=10)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--output-dir", default="final_experiments")
    p.add_argument("--skip-kfold", action="store_true")
    p.add_argument("--skip-test", action="store_true")
    p.add_argument("--class-weighted", action="store_true", default=True,
                   help="Use sqrt-inverse-frequency class-weighted CE loss to mitigate imbalance.")
    p.add_argument("--no-class-weighted", dest="class_weighted", action="store_false")
    return p.parse_args()


def main():
    args = parse_args()
    manifest = json.loads(Path(args.manifest).read_text())
    root = Path(manifest["data_dir"])
    classes = manifest["classes"]

    pool_samples = class_names_to_int(
        [(p, c) for p, c in zip(manifest["kfold"]["pool_paths"], manifest["kfold"]["pool_labels"])],
        classes,
    )
    test_samples = class_names_to_int(manifest["splits"]["test"], classes)

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    weights_dir = out_dir / "weights"
    weights_dir.mkdir(exist_ok=True)
    log_path = out_dir / "run.log"

    def log(line):
        msg = f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {line}"
        print(msg, flush=True)
        with log_path.open("a") as fh:
            fh.write(msg + "\n")

    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    log(f"device={device} classes={len(classes)} pool={len(pool_samples)} test={len(test_samples)}")

    kfold_summaries = {}
    test_results = {}

    for model_name in args.models:
        log(f"########## {model_name} ##########")
        try:
            if not args.skip_kfold:
                s = run_kfold(model_name, pool_samples, classes, root, args, device, log)
                kfold_summaries[model_name] = s
                (out_dir / f"{model_name}_kfold.json").write_text(json.dumps(s, indent=2))
            if not args.skip_test:
                r = run_indep_test(model_name, pool_samples, test_samples, classes, root, args, device, log, weights_dir)
                test_results[model_name] = r
                # Save without huge preds/labels/probs arrays inline
                slim = {k: v for k, v in r.items() if k not in ("preds", "labels", "probs")}
                (out_dir / f"{model_name}_test.json").write_text(json.dumps(slim, indent=2))
                # Save preds + probs separately for stat tests, ROC, calibration analysis
                (out_dir / f"{model_name}_test_preds.json").write_text(json.dumps({
                    "preds": r["preds"], "labels": r["labels"], "probs": r["probs"]}))
        except Exception as exc:
            log(f"!! {model_name} FAILED: {exc!r}")

    if test_results:
        mc = mcnemar_pairwise(test_results)
        (out_dir / "mcnemar.json").write_text(json.dumps(mc, indent=2))
        log(f"McNemar pairwise saved -> {out_dir / 'mcnemar.json'}")

    if kfold_summaries:
        summary = {m: {k: v for k, v in s.items() if k != "fold_metrics"} for m, s in kfold_summaries.items()}
        (out_dir / "kfold_summary.json").write_text(json.dumps(summary, indent=2))
    log("ALL DONE")


if __name__ == "__main__":
    main()