scripts/run_revised_experiments.py

#!/usr/bin/env python3
"""Run revised RAVEL staged experiments with reusable outputs.

The script is resumable: a run is skipped when its metrics JSON and prediction
CSV already exist, unless `--overwrite` is passed.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import random
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score
from torch.optim import AdamW
from transformers import CLIPProcessor, DebertaV2Tokenizer
from transformers.utils import logging as hf_logging

hf_logging.set_verbosity_error()

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.revised_ravel_model import token_loss


@dataclass(frozen=True)
class MethodSpec:
    key: str
    display: str
    architecture: str
    train_target: str
    enable_lora: bool = True
    use_coattention: bool = True
    use_auxiliary: bool = True
    use_disagreement: bool = True
    parameter_matched: bool = False
    disagreement_formulation: str = "signed_difference"
    additional_loss: str = "none"
    lambda_primary: float = 0.5
    lambda_unimodal: float = 0.25
    consistency_beta: float = 0.25
    infonce_alpha: float = 0.10
    infonce_temperature: float = 0.07


METHODS: Dict[str, MethodSpec] = {
    "frozen_concat": MethodSpec(
        key="frozen_concat",
        display="Frozen concat",
        architecture="token",
        train_target="no_coattn",
        enable_lora=False,
        use_coattention=False,
        use_auxiliary=False,
        use_disagreement=False,
    ),
    "lora_concat": MethodSpec(
        key="lora_concat",
        display="LoRA concat",
        architecture="token",
        train_target="no_coattn",
        enable_lora=True,
        use_coattention=False,
        use_auxiliary=False,
        use_disagreement=False,
    ),
    "frozen_token_coattn": MethodSpec(
        key="frozen_token_coattn",
        display="Frozen token co-attention",
        architecture="token",
        train_target="primary",
        enable_lora=False,
        use_coattention=True,
        use_auxiliary=False,
        use_disagreement=False,
    ),
    "legacy_global": MethodSpec(
        key="legacy_global",
        display="RAVEL-Global",
        architecture="legacy_global",
        train_target="legacy_full",
        enable_lora=True,
        use_coattention=False,
        use_auxiliary=False,
        use_disagreement=False,
    ),
    "token_coattn": MethodSpec(
        key="token_coattn",
        display="Token co-attention",
        architecture="token",
        train_target="primary",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=False,
        use_disagreement=False,
    ),
    "token_aux": MethodSpec(
        key="token_aux",
        display="Token co-attention + unimodal heads",
        architecture="token",
        train_target="primary_aux",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=False,
    ),
    "param_mlp": MethodSpec(
        key="param_mlp",
        display="Parameter-matched MLP",
        architecture="token",
        train_target="param_mlp",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=False,
        parameter_matched=True,
    ),
    "full_revised": MethodSpec(
        key="full_revised",
        display="Full revised RAVEL",
        architecture="token",
        train_target="full",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
    ),
    "text_only": MethodSpec(
        key="text_only",
        display="Text-only",
        architecture="token",
        train_target="text_only",
        enable_lora=True,
        use_coattention=False,
        use_auxiliary=True,
        use_disagreement=False,
    ),
    "vision_only": MethodSpec(
        key="vision_only",
        display="Vision-only",
        architecture="token",
        train_target="vision_only",
        enable_lora=True,
        use_coattention=False,
        use_auxiliary=True,
        use_disagreement=False,
    ),
    "absolute_difference": MethodSpec(
        key="absolute_difference",
        display="Absolute posterior difference",
        architecture="token",
        train_target="disagreement_feature",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        disagreement_formulation="absolute_difference",
    ),
    "js_divergence": MethodSpec(
        key="js_divergence",
        display="Jensen-Shannon divergence",
        architecture="token",
        train_target="disagreement_feature",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        disagreement_formulation="js_divergence",
    ),
    "log_probability_ratio": MethodSpec(
        key="log_probability_ratio",
        display="Log-probability ratio",
        architecture="token",
        train_target="disagreement_feature",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        disagreement_formulation="log_probability_ratio",
    ),
    "attention_discrepancy": MethodSpec(
        key="attention_discrepancy",
        display="Attention discrepancy",
        architecture="token",
        train_target="disagreement_feature",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        disagreement_formulation="attention_discrepancy",
    ),
    "consistency_loss": MethodSpec(
        key="consistency_loss",
        display="JS consistency loss",
        architecture="token",
        train_target="consistency_loss",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=False,
        additional_loss="js_consistency",
    ),
    "infonce_alignment": MethodSpec(
        key="infonce_alignment",
        display="InfoNCE alignment",
        architecture="token",
        train_target="infonce_alignment",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=False,
        additional_loss="infonce_alignment",
    ),
    "lambda_u_0_0": MethodSpec(
        key="lambda_u_0_0",
        display="Full revised RAVEL lambda_u=0.0",
        architecture="token",
        train_target="full",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        lambda_unimodal=0.0,
    ),
    "lambda_u_0_25": MethodSpec(
        key="lambda_u_0_25",
        display="Full revised RAVEL lambda_u=0.25",
        architecture="token",
        train_target="full",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        lambda_unimodal=0.25,
    ),
    "lambda_u_0_5": MethodSpec(
        key="lambda_u_0_5",
        display="Full revised RAVEL lambda_u=0.5",
        architecture="token",
        train_target="full",
        enable_lora=True,
        use_coattention=True,
        use_auxiliary=True,
        use_disagreement=True,
        lambda_unimodal=0.5,
    ),
}


STAGE_B_METHODS = [
    "frozen_concat",
    "lora_concat",
    "frozen_token_coattn",
    "legacy_global",
    "token_coattn",
    "token_aux",
    "param_mlp",
    "full_revised",
]

STAGE_C_METHODS = [
    "legacy_global",
    "text_only",
    "vision_only",
    "token_coattn",
    "token_aux",
    "param_mlp",
    "full_revised",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run revised RAVEL staged experiments.")
    parser.add_argument("--stage", choices=["stage_b", "stage_c", "all"], default="all")
    parser.add_argument("--datasets", nargs="+", default=["mvsa_multiple", "hfm_existing", "hfm_deleak"])
    parser.add_argument("--seeds", nargs="+", type=int, default=[1, 3, 5, 7, 11])
    parser.add_argument("--stage-b-seed", type=int, default=7)
    parser.add_argument("--methods", nargs="+", default=None, choices=sorted(METHODS))
    parser.add_argument("--epochs", type=int, default=5)
    parser.add_argument("--patience", type=int, default=2)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--grad-accum-steps", type=int, default=2)
    parser.add_argument("--max-length", type=int, default=96)
    parser.add_argument("--learning-rate", type=float, default=5e-5)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--num-workers", type=int, default=4)
    parser.add_argument("--output-root", default="ravel_revision_results")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--limit-train-samples", type=int, default=None)
    parser.add_argument("--limit-val-samples", type=int, default=None)
    parser.add_argument("--limit-test-samples", type=int, default=None)
    parser.add_argument(
        "--hfm-deleak-manifest",
        default="ravel_revision_results/data_audit/hfm_split_manifest_deleaked.csv",
    )
    return parser.parse_args()


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def write_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: Sequence[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(fieldnames), extrasaction="ignore")
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def append_csv(path: Path, rows: Iterable[Dict[str, Any]], fieldnames: Sequence[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    exists = path.exists()
    with path.open("a", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=list(fieldnames), extrasaction="ignore")
        if not exists:
            writer.writeheader()
        for row in rows:
            writer.writerow(row)


def freeze_non_lora(module: nn.Module) -> None:
    for name, parameter in module.named_parameters():
        parameter.requires_grad = "lora_" in name


def set_module_trainable(module: Optional[nn.Module], trainable: bool) -> None:
    if module is None:
        return
    for parameter in module.parameters():
        parameter.requires_grad = trainable


def apply_method_trainability(model: nn.Module, method: MethodSpec) -> None:
    """Freeze modules that are intentionally bypassed by an ablation target."""
    target = method.train_target
    if target == "no_coattn":
        for name in [
            "fusion",
            "visual_head",
            "text_head",
            "refinement",
            "extra_mlp_control",
        ]:
            set_module_trainable(getattr(model, name, None), False)
        return

    if target == "primary":
        for name in ["visual_head", "text_head", "refinement", "extra_mlp_control"]:
            set_module_trainable(getattr(model, name, None), False)
        return

    if target == "primary_aux":
        for name in ["refinement", "extra_mlp_control"]:
            set_module_trainable(getattr(model, name, None), False)
        return

    if target == "param_mlp":
        set_module_trainable(getattr(model, "refinement", None), False)
        return

    if target == "disagreement_feature":
        set_module_trainable(getattr(model, "extra_mlp_control", None), False)
        return

    if target in {"consistency_loss", "infonce_alignment"}:
        for name in ["refinement", "extra_mlp_control"]:
            set_module_trainable(getattr(model, name, None), False)
        return


def maybe_limit(samples: List[Any], limit: Optional[int]) -> List[Any]:
    if limit is None or limit <= 0 or len(samples) <= limit:
        return samples
    buckets: Dict[Any, List[Any]] = {}
    for sample in samples:
        label = getattr(sample, "label", getattr(sample, "combined_majority", ""))
        buckets.setdefault(label, []).append(sample)
    per_label = max(1, limit // max(1, len(buckets)))
    selected: List[Any] = []
    for label in sorted(buckets, key=str):
        selected.extend(buckets[label][:per_label])
    selected.extend(samples[: max(0, limit - len(selected))])
    return selected[:limit]


def load_dataset(
    dataset_key: str,
    seed: int,
    batch_size: int,
    max_length: int,
    num_workers: int,
    method: MethodSpec,
    hfm_deleak_manifest: str,
    limits: Tuple[Optional[int], Optional[int], Optional[int]],
) -> Tuple[Any, Dict[str, Any], Any, Any, Any, List[Any], List[Any], List[Any], int, List[str]]:
    if dataset_key == "mvsa_multiple":
        from src.mvsa_multiple_pipeline import (
            CLARAModel,
            DEFAULT_MVSA_MULTIPLE_CONFIG,
            LABEL_ID_TO_NAME,
            MVSALoader,
            create_dataloaders,
            summarize_splits,
        )

        cfg = dict(DEFAULT_MVSA_MULTIPLE_CONFIG)
        cfg.update(
            {
                "architecture": method.architecture,
                "enable_clip_lora": method.enable_lora,
                "batch_size": batch_size,
                "max_length": max_length,
                "num_workers": num_workers,
                "pin_memory": True,
                "persistent_workers": bool(num_workers > 0),
                "prefetch_factor": 2,
                "seed": seed,
                "learning_rate": 5e-5,
                "weight_decay": 0.01,
                "use_mixup_negative": False,
                "use_weighted_sampler": False,
                "text_unfreeze_mode": "freeze_all",
                "unfreeze_epoch": 0,
                "paper_exact_counts": True,
            }
        )
        loader = MVSALoader(cfg["text_dir"], cfg["label_file"])
        loader.load(
            preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")),
            require_unanimous=bool(cfg["require_unanimous"]),
            require_cross_agree=bool(cfg["require_cross_agree"]),
            paper_exact_counts=True,
        )
        train_samples, val_samples, test_samples = loader.split(
            train_ratio=float(cfg["train_ratio"]),
            val_ratio=float(cfg["val_ratio"]),
            seed=seed,
            paper_811=True,
        )
        train_samples = maybe_limit(train_samples, limits[0])
        val_samples = maybe_limit(val_samples, limits[1])
        test_samples = maybe_limit(test_samples, limits[2])
        processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
        tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])
        train_loader, val_loader, test_loader = create_dataloaders(
            train_samples=train_samples,
            val_samples=val_samples,
            test_samples=test_samples,
            clip_processor=processor,
            tokenizer=tokenizer,
            batch_size=batch_size,
            max_length=max_length,
            num_workers=num_workers,
            pin_memory=True,
            persistent_workers=bool(num_workers > 0),
            prefetch_factor=2,
            use_mixup_negative=False,
            mixup_alpha=0.0,
            negative_class_boost=1.0,
            min_ratio_negative=0.0,
            weighted_train_sampler=False,
        )
        label_names = [LABEL_ID_TO_NAME[idx] for idx in range(int(cfg["num_classes"]))]
        return (
            CLARAModel,
            cfg,
            train_loader,
            val_loader,
            test_loader,
            train_samples,
            val_samples,
            test_samples,
            int(cfg["num_classes"]),
            label_names,
        )

    if dataset_key == "mvsa_single":
        from src.mvsa_single_pipeline import (
            CLARAModel,
            DEFAULT_MVSA_SINGLE_CONFIG,
            LABEL_ID_TO_NAME,
            MVSASingleLoader,
            create_dataloaders,
        )

        cfg = dict(DEFAULT_MVSA_SINGLE_CONFIG)
        cfg.update(
            {
                "architecture": method.architecture,
                "enable_clip_lora": method.enable_lora,
                "enable_text_lora": method.enable_lora,
                "batch_size": batch_size,
                "max_length": max_length,
                "num_workers": num_workers,
                "pin_memory": True,
                "persistent_workers": bool(num_workers > 0),
                "prefetch_factor": 2,
                "seed": seed,
                "learning_rate": 5e-5,
                "weight_decay": 0.01,
                "use_mixup_negative": False,
                "use_weighted_sampler": False,
                "text_unfreeze_mode": "freeze_all",
                "unfreeze_epoch": 0,
                "loss_type": "ce",
                "ce_use_class_weights": False,
                "label_smoothing": 0.0,
            }
        )
        loader = MVSASingleLoader(cfg["text_dir"], cfg["label_file"])
        samples = loader.load(
            preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")),
            require_unanimous=bool(cfg.get("require_unanimous", False)),
            require_cross_agree=bool(cfg.get("require_cross_agree", False)),
        )
        train_samples, val_samples, test_samples = loader.split(
            train_ratio=float(cfg["train_ratio"]),
            val_ratio=float(cfg["val_ratio"]),
            seed=seed,
        )
        train_samples = maybe_limit(train_samples, limits[0])
        val_samples = maybe_limit(val_samples, limits[1])
        test_samples = maybe_limit(test_samples, limits[2])
        processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
        tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])
        train_loader, val_loader, test_loader = create_dataloaders(
            train_samples=train_samples,
            val_samples=val_samples,
            test_samples=test_samples,
            clip_processor=processor,
            tokenizer=tokenizer,
            batch_size=batch_size,
            max_length=max_length,
            num_workers=num_workers,
            pin_memory=True,
            persistent_workers=bool(num_workers > 0),
            prefetch_factor=2,
            use_mixup_negative=False,
            mixup_alpha=0.0,
            negative_class_boost=1.0,
            min_ratio_negative=0.0,
            weighted_train_sampler=False,
        )
        label_names = [LABEL_ID_TO_NAME[idx] for idx in range(int(cfg["num_classes"]))]
        _ = samples  # Keep the load step explicit for provenance and error checking.
        return (
            CLARAModel,
            cfg,
            train_loader,
            val_loader,
            test_loader,
            train_samples,
            val_samples,
            test_samples,
            int(cfg["num_classes"]),
            label_names,
        )

    if dataset_key in {"hfm_existing", "hfm_deleak"}:
        from src.hfm_pipeline import (
            CLARAModel,
            DEFAULT_HFM_CONFIG,
            LABEL_ID_TO_NAME,
            HFMLoader,
            create_dataloaders,
        )

        cfg = dict(DEFAULT_HFM_CONFIG)
        cfg.update(
            {
                "architecture": method.architecture,
                "enable_clip_lora": method.enable_lora,
                "batch_size": batch_size,
                "max_length": max_length,
                "num_workers": num_workers,
                "pin_memory": True,
                "seed": seed,
                "learning_rate": 5e-5,
                "weight_decay": 0.01,
                "text_unfreeze_mode": "freeze_all",
                "num_classes": 2,
            }
        )
        loader = HFMLoader(cfg["text_dir"], cfg["image_root"])
        if dataset_key == "hfm_deleak":
            loader.load_from_manifest(hfm_deleak_manifest)
        else:
            loader.load()
        train_samples = maybe_limit(loader.get_split("train"), limits[0])
        val_samples = maybe_limit(loader.get_split("val"), limits[1])
        test_samples = maybe_limit(loader.get_split("test"), limits[2])
        processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
        tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])
        train_loader, val_loader, test_loader = create_dataloaders(
            train_samples=train_samples,
            val_samples=val_samples,
            test_samples=test_samples,
            clip_processor=processor,
            tokenizer=tokenizer,
            batch_size=batch_size,
            max_length=max_length,
            num_workers=num_workers,
            pin_memory=True,
            weighted_train_sampler=False,
        )
        label_names = [LABEL_ID_TO_NAME[idx] for idx in range(int(cfg["num_classes"]))]
        return (
            CLARAModel,
            cfg,
            train_loader,
            val_loader,
            test_loader,
            train_samples,
            val_samples,
            test_samples,
            int(cfg["num_classes"]),
            label_names,
        )

    raise ValueError(f"Unsupported dataset: {dataset_key}")


def _js_vector_torch(p: torch.Tensor, q: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    p = p.clamp(min=eps)
    q = q.clamp(min=eps)
    m = (0.5 * (p + q)).clamp(min=eps)
    return 0.5 * p * (p / m).log() + 0.5 * q * (q / m).log()


def _attention_entropy(attn: torch.Tensor, eps: float = 1e-8) -> torch.Tensor:
    probs = attn.clamp(min=eps)
    entropy = -(probs * probs.log()).sum(dim=-1)
    support = attn.size(-1)
    if support > 1:
        entropy = entropy / math.log(float(support))
    return entropy.mean(dim=tuple(range(1, entropy.ndim)))


def disagreement_feature(
    outputs: Dict[str, torch.Tensor],
    batch: Dict[str, torch.Tensor],
    formulation: str,
    num_classes: int,
) -> torch.Tensor:
    visual_probs = outputs["visual_probs"]
    text_probs = outputs["text_probs"]
    if formulation == "signed_difference":
        return visual_probs - text_probs
    if formulation == "absolute_difference":
        return (visual_probs - text_probs).abs()
    if formulation == "js_divergence":
        return _js_vector_torch(visual_probs, text_probs)
    if formulation == "log_probability_ratio":
        return (visual_probs.clamp(min=1e-8).log() - text_probs.clamp(min=1e-8).log())
    if formulation == "attention_discrepancy":
        v2t_list = outputs.get("attention_v2t", [])
        t2v_list = outputs.get("attention_t2v", [])
        if not v2t_list or not t2v_list:
            raise RuntimeError("attention_discrepancy requires return_attention=True outputs.")
        v2t = torch.stack(v2t_list, dim=0).mean(dim=0)
        t2v = torch.stack(t2v_list, dim=0).mean(dim=0)
        scalar = (_attention_entropy(v2t) - _attention_entropy(t2v)).unsqueeze(-1)
        return scalar.expand(-1, num_classes)
    raise ValueError(f"Unsupported disagreement formulation: {formulation}")


def consistency_js_loss(outputs: Dict[str, torch.Tensor]) -> torch.Tensor:
    return _js_vector_torch(outputs["visual_probs"], outputs["text_probs"]).sum(dim=-1).mean()


def infonce_loss(outputs: Dict[str, torch.Tensor], temperature: float) -> torch.Tensor:
    visual = F.normalize(outputs["visual_global"].float(), dim=-1)
    text = F.normalize(outputs["text_global"].float(), dim=-1)
    logits = visual @ text.T / max(float(temperature), 1e-6)
    labels = torch.arange(logits.size(0), device=logits.device)
    return 0.5 * (F.cross_entropy(logits, labels) + F.cross_entropy(logits.T, labels))


def logits_for_target(
    model: nn.Module,
    batch: Dict[str, torch.Tensor],
    method_or_target: Any,
    criterion: nn.Module,
    labels: torch.Tensor,
    num_classes: Optional[int] = None,
) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
    if isinstance(method_or_target, MethodSpec):
        method = method_or_target
        target = method.train_target
    else:
        method = None
        target = str(method_or_target)
    pixel_values = batch["pixel_values"]
    input_ids = batch["input_ids"]
    attention_mask = batch["attention_mask"]

    if target == "legacy_full":
        outputs = model(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask)
        logits = outputs["logits"]
        return logits, criterion(logits, labels), outputs, {}

    if target == "no_coattn":
        logits = model.logits_without_coattention(pixel_values, input_ids, attention_mask)
        return logits, criterion(logits, labels), {}, {}

    if target in {"text_only", "vision_only"} and hasattr(model, "encode_modalities"):
        encoded = model.encode_modalities(pixel_values, input_ids, attention_mask)
        if target == "text_only":
            logits = model.text_head(encoded["text_global"])
        else:
            logits = model.visual_head(encoded["visual_global"])
        return logits, criterion(logits, labels), {}, {}

    outputs = model(pixel_values=pixel_values, input_ids=input_ids, attention_mask=attention_mask)

    if target == "full":
        loss, parts = token_loss(
            outputs,
            labels,
            criterion,
            lambda_primary=method.lambda_primary if method is not None else 0.5,
            lambda_unimodal=method.lambda_unimodal if method is not None else 0.25,
        )
        return outputs["logits"], loss, outputs, parts

    if target == "primary":
        logits = outputs["pred_logits"]
        return logits, criterion(logits, labels), outputs, {}

    if target == "primary_aux":
        logits = outputs["pred_logits"]
        loss = criterion(logits, labels)
        loss = loss + 0.25 * (
            criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels)
        )
        return logits, loss, outputs, {}

    if target == "param_mlp":
        logits = model.extra_mlp_control(outputs["fused"])
        loss = criterion(logits, labels)
        loss = loss + 0.25 * criterion(outputs["pred_logits"], labels)
        loss = loss + 0.25 * (
            criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels)
        )
        return logits, loss, outputs, {}

    if target == "disagreement_feature":
        if method is None:
            raise ValueError("disagreement_feature requires a MethodSpec.")
        if num_classes is None:
            num_classes = int(outputs["visual_probs"].shape[-1])
        if method.disagreement_formulation == "attention_discrepancy":
            outputs = model(
                pixel_values=pixel_values,
                input_ids=input_ids,
                attention_mask=attention_mask,
                return_attention=True,
            )
        feature = disagreement_feature(
            outputs,
            batch,
            formulation=method.disagreement_formulation,
            num_classes=int(num_classes),
        ).to(dtype=outputs["fused"].dtype)
        logits = model.refinement(outputs["fused"], feature)
        loss = criterion(logits, labels)
        loss = loss + method.lambda_primary * criterion(outputs["pred_logits"], labels)
        loss = loss + method.lambda_unimodal * (
            criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels)
        )
        outputs = dict(outputs)
        outputs["logits"] = logits
        outputs["disagreement"] = feature
        return logits, loss, outputs, {}

    if target == "consistency_loss":
        if method is None:
            raise ValueError("consistency_loss requires a MethodSpec.")
        logits = outputs["pred_logits"]
        loss = criterion(logits, labels)
        loss = loss + method.lambda_unimodal * (
            criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels)
        )
        loss = loss + method.consistency_beta * consistency_js_loss(outputs)
        outputs = dict(outputs)
        outputs["logits"] = logits
        outputs["disagreement"] = torch.zeros_like(outputs["visual_probs"])
        return logits, loss, outputs, {}

    if target == "infonce_alignment":
        if method is None:
            raise ValueError("infonce_alignment requires a MethodSpec.")
        logits = outputs["pred_logits"]
        loss = criterion(logits, labels)
        loss = loss + method.lambda_unimodal * (
            criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels)
        )
        loss = loss + method.infonce_alpha * infonce_loss(outputs, method.infonce_temperature)
        outputs = dict(outputs)
        outputs["logits"] = logits
        outputs["disagreement"] = torch.zeros_like(outputs["visual_probs"])
        return logits, loss, outputs, {}

    if target == "no_disagreement":
        logits = model.refinement(outputs["fused"], torch.zeros_like(outputs["disagreement"]))
        loss = criterion(logits, labels)
        loss = loss + 0.5 * criterion(outputs["pred_logits"], labels)
        loss = loss + 0.25 * (
            criterion(outputs["visual_logits"], labels) + criterion(outputs["text_logits"], labels)
        )
        return logits, loss, outputs, {}

    raise ValueError(target)


def expected_calibration_error(probs: np.ndarray, y_true: np.ndarray, bins: int = 15) -> float:
    conf = probs.max(axis=1)
    pred = probs.argmax(axis=1)
    correct = (pred == y_true).astype(float)
    edges = np.linspace(0.0, 1.0, bins + 1)
    ece = 0.0
    for low, high in zip(edges[:-1], edges[1:]):
        mask = (conf > low) & (conf <= high)
        if not mask.any():
            continue
        ece += (mask.mean()) * abs(correct[mask].mean() - conf[mask].mean())
    return float(ece)


def adaptive_ece(probs: np.ndarray, y_true: np.ndarray, bins: int = 15) -> float:
    conf = probs.max(axis=1)
    pred = probs.argmax(axis=1)
    correct = (pred == y_true).astype(float)
    order = np.argsort(conf)
    chunks = np.array_split(order, min(bins, len(order)))
    ece = 0.0
    for chunk in chunks:
        if len(chunk) == 0:
            continue
        ece += (len(chunk) / len(conf)) * abs(correct[chunk].mean() - conf[chunk].mean())
    return float(ece)


def brier_score(probs: np.ndarray, y_true: np.ndarray, num_classes: int) -> float:
    one_hot = np.eye(num_classes)[y_true]
    return float(np.mean(np.sum((probs - one_hot) ** 2, axis=1)))


def nll_score(probs: np.ndarray, y_true: np.ndarray) -> float:
    return float(-np.mean(np.log(np.clip(probs[np.arange(len(y_true)), y_true], 1e-12, 1.0))))


def aurc_score(probs: np.ndarray, y_true: np.ndarray) -> Tuple[float, float, float, float, float]:
    conf = probs.max(axis=1)
    pred = probs.argmax(axis=1)
    errors = (pred != y_true).astype(float)
    order = np.argsort(-conf)
    sorted_errors = errors[order]
    cum_errors = np.cumsum(sorted_errors)
    coverage = np.arange(1, len(errors) + 1) / max(1, len(errors))
    risk = cum_errors / np.arange(1, len(errors) + 1)
    trapezoid = getattr(np, "trapezoid", None)
    if trapezoid is None:
        trapezoid = getattr(np, "trapz")
    aurc = float(trapezoid(risk, coverage)) if len(errors) > 1 else float(risk[-1])

    def risk_at(cov: float) -> float:
        idx = max(0, min(len(errors) - 1, math.ceil(cov * len(errors)) - 1))
        return float(risk[idx])

    def acc_at(cov: float) -> float:
        return 1.0 - risk_at(cov)

    return aurc, risk_at(0.8), risk_at(0.9), acc_at(0.8), acc_at(0.9)


def compute_metrics(
    logits: np.ndarray,
    y_true: np.ndarray,
    num_classes: int,
) -> Dict[str, float]:
    probs = torch.softmax(torch.tensor(logits, dtype=torch.float32), dim=-1).numpy()
    pred = probs.argmax(axis=1)
    metrics: Dict[str, float] = {
        "accuracy": float(accuracy_score(y_true, pred)),
        "macro_precision": float(precision_score(y_true, pred, average="macro", zero_division=0)),
        "macro_recall": float(recall_score(y_true, pred, average="macro", zero_division=0)),
        "macro_f1": float(f1_score(y_true, pred, average="macro", zero_division=0)),
        "weighted_f1": float(f1_score(y_true, pred, average="weighted", zero_division=0)),
        "raw_ece": expected_calibration_error(probs, y_true),
        "adaptive_ece": adaptive_ece(probs, y_true),
        "nll": nll_score(probs, y_true),
        "brier": brier_score(probs, y_true, num_classes),
    }
    aurc, r80, r90, a80, a90 = aurc_score(probs, y_true)
    metrics.update(
        {
            "aurc": aurc,
            "risk_at_80_coverage": r80,
            "risk_at_90_coverage": r90,
            "accuracy_at_80_coverage": a80,
            "accuracy_at_90_coverage": a90,
        }
    )
    try:
        if num_classes == 2:
            metrics["auroc"] = float(roc_auc_score(y_true, probs[:, 1]))
        else:
            metrics["auroc"] = float(roc_auc_score(y_true, probs, multi_class="ovr", average="macro"))
    except Exception:
        metrics["auroc"] = float("nan")
    return metrics


def samples_to_ids(samples: List[Any]) -> List[str]:
    ids: List[str] = []
    for idx, sample in enumerate(samples):
        ids.append(str(getattr(sample, "image_id", getattr(sample, "sample_id", idx))))
    return ids


def evaluate(
    model: nn.Module,
    loader: Any,
    samples: List[Any],
    method: MethodSpec,
    device: torch.device,
    criterion: nn.Module,
    num_classes: int,
    dataset_key: str,
    seed: int,
) -> Tuple[Dict[str, float], List[Dict[str, Any]], np.ndarray, np.ndarray]:
    model.eval()
    logits_all: List[np.ndarray] = []
    labels_all: List[np.ndarray] = []
    visual_probs_all: List[np.ndarray] = []
    text_probs_all: List[np.ndarray] = []
    primary_logits_all: List[np.ndarray] = []
    refined_logits_all: List[np.ndarray] = []

    with torch.no_grad():
        for batch in loader:
            labels = batch["labels"].long().to(device)
            batch_gpu = {
                key: value.to(device, non_blocking=True) if torch.is_tensor(value) else value
                for key, value in batch.items()
            }
            logits, _loss, outputs, _parts = logits_for_target(
                model, batch_gpu, method, criterion, labels, num_classes=num_classes
            )
            logits_all.append(logits.float().cpu().numpy())
            labels_all.append(labels.cpu().numpy())
            if isinstance(outputs, dict) and "visual_probs" in outputs:
                visual_probs_all.append(outputs["visual_probs"].float().cpu().numpy())
                text_probs_all.append(outputs["text_probs"].float().cpu().numpy())
                primary_logits_all.append(outputs["pred_logits"].float().cpu().numpy())
                refined_logits_all.append(outputs["logits"].float().cpu().numpy())

    logits_np = np.concatenate(logits_all, axis=0)
    y_np = np.concatenate(labels_all, axis=0)
    metrics = compute_metrics(logits_np, y_np, num_classes)
    probs = torch.softmax(torch.tensor(logits_np, dtype=torch.float32), dim=-1).numpy()
    pred = probs.argmax(axis=1)
    sample_ids = samples_to_ids(samples)
    visual_probs = np.concatenate(visual_probs_all, axis=0) if visual_probs_all else np.full_like(probs, np.nan)
    text_probs = np.concatenate(text_probs_all, axis=0) if text_probs_all else np.full_like(probs, np.nan)
    primary_logits = (
        np.concatenate(primary_logits_all, axis=0) if primary_logits_all else np.full_like(logits_np, np.nan)
    )
    refined_logits = (
        np.concatenate(refined_logits_all, axis=0) if refined_logits_all else np.full_like(logits_np, np.nan)
    )
    primary_probs = torch.softmax(torch.tensor(primary_logits, dtype=torch.float32), dim=-1).numpy()
    refined_probs = torch.softmax(torch.tensor(refined_logits, dtype=torch.float32), dim=-1).numpy()
    disagreement = visual_probs - text_probs

    rows: List[Dict[str, Any]] = []
    for i in range(len(y_np)):
        row: Dict[str, Any] = {
            "sample_id": sample_ids[i] if i < len(sample_ids) else str(i),
            "dataset": dataset_key,
            "split": "test",
            "seed": seed,
            "method": method.key,
            "true_label": int(y_np[i]),
            "predicted_label": int(pred[i]),
            "prediction_entropy": float(-(probs[i] * np.log(np.clip(probs[i], 1e-12, 1.0))).sum()),
            "max_confidence": float(probs[i].max()),
            "primary_prediction": int(np.nanargmax(primary_probs[i])) if np.isfinite(primary_probs[i]).all() else "",
            "primary_confidence": float(np.nanmax(primary_probs[i])) if np.isfinite(primary_probs[i]).all() else "",
            "refined_prediction": int(np.nanargmax(refined_probs[i])) if np.isfinite(refined_probs[i]).all() else "",
            "refined_confidence": float(np.nanmax(refined_probs[i])) if np.isfinite(refined_probs[i]).all() else "",
            "correct_primary": int(np.nanargmax(primary_probs[i]) == y_np[i]) if np.isfinite(primary_probs[i]).all() else "",
            "correct_refined": int(np.nanargmax(refined_probs[i]) == y_np[i]) if np.isfinite(refined_probs[i]).all() else "",
            "number_of_images": 1,
        }
        for c in range(num_classes):
            row[f"logit_class_{c}"] = float(logits_np[i, c])
            row[f"prob_class_{c}"] = float(probs[i, c])
            row[f"visual_prob_class_{c}"] = float(visual_probs[i, c]) if np.isfinite(visual_probs[i, c]) else ""
            row[f"text_prob_class_{c}"] = float(text_probs[i, c]) if np.isfinite(text_probs[i, c]) else ""
            row[f"disagreement_class_{c}"] = float(disagreement[i, c]) if np.isfinite(disagreement[i, c]) else ""
        if np.isfinite(disagreement[i]).all():
            row["disagreement_l1"] = float(np.abs(disagreement[i]).sum())
        else:
            row["disagreement_l1"] = ""
        rows.append(row)

    return metrics, rows, logits_np, y_np


def train_one(
    model: nn.Module,
    train_loader: Any,
    val_loader: Any,
    method: MethodSpec,
    device: torch.device,
    epochs: int,
    patience: int,
    grad_accum_steps: int,
    learning_rate: float,
    weight_decay: float,
    num_classes: int,
) -> Tuple[Dict[str, Any], Dict[str, torch.Tensor], List[Dict[str, Any]]]:
    criterion = nn.CrossEntropyLoss()
    optimizer = AdamW(
        [parameter for parameter in model.parameters() if parameter.requires_grad],
        lr=learning_rate,
        weight_decay=weight_decay,
    )
    use_bf16 = bool(device.type == "cuda" and torch.cuda.is_bf16_supported())
    amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
    def trainable_state_dict() -> Dict[str, torch.Tensor]:
        trainable_names = {name for name, parameter in model.named_parameters() if parameter.requires_grad}
        return {
            key: value.detach().cpu().clone()
            for key, value in model.state_dict().items()
            if key in trainable_names
        }

    best_state: Dict[str, torch.Tensor] = {}
    best_val = -1.0
    best_epoch = 0
    bad_epochs = 0
    history: List[Dict[str, Any]] = []

    for epoch in range(1, epochs + 1):
        model.train()
        train_losses: List[float] = []
        optimizer.zero_grad(set_to_none=True)
        start = time.time()
        for step, batch in enumerate(train_loader, start=1):
            labels = batch["labels"].long().to(device, non_blocking=True)
            batch_gpu = {
                key: value.to(device, non_blocking=True) if torch.is_tensor(value) else value
                for key, value in batch.items()
            }
            with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=device.type == "cuda"):
                _logits, loss, _outputs, _parts = logits_for_target(
                    model, batch_gpu, method, criterion, labels, num_classes=num_classes
                )
                loss = loss / max(1, grad_accum_steps)
            loss.backward()
            if step % max(1, grad_accum_steps) == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                optimizer.step()
                optimizer.zero_grad(set_to_none=True)
            train_losses.append(float(loss.detach().item()) * max(1, grad_accum_steps))
            if step % 250 == 0:
                print(
                    f"  epoch={epoch} step={step}/{len(train_loader)} "
                    f"loss={float(np.mean(train_losses[-50:])):.4f}",
                    flush=True,
                )
        if len(train_loader) % max(1, grad_accum_steps) != 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad(set_to_none=True)

        val_metrics, _rows, _logits, _labels = evaluate(
            model=model,
            loader=val_loader,
            samples=[],
            method=method,
            device=device,
            criterion=criterion,
            num_classes=num_classes,
            dataset_key="val",
            seed=0,
        )
        val_score = val_metrics["macro_f1"]
        history_row = {
            "epoch": epoch,
            "train_loss": float(np.mean(train_losses)) if train_losses else float("nan"),
            "val_macro_f1": val_metrics["macro_f1"],
            "val_weighted_f1": val_metrics["weighted_f1"],
            "val_accuracy": val_metrics["accuracy"],
            "val_raw_ece": val_metrics["raw_ece"],
            "elapsed_seconds": time.time() - start,
        }
        history.append(history_row)
        print(
            f"epoch={epoch} train_loss={history_row['train_loss']:.4f} "
            f"val_mF1={val_metrics['macro_f1']:.4f} val_ECE={val_metrics['raw_ece']:.4f}",
            flush=True,
        )
        if val_score > best_val:
            best_val = val_score
            best_epoch = epoch
            best_state = trainable_state_dict()
            bad_epochs = 0
        else:
            bad_epochs += 1
            if bad_epochs >= patience:
                break

    return {"best_epoch": best_epoch, "best_val_macro_f1": best_val}, best_state, history


def run_single(args: argparse.Namespace, dataset_key: str, method_key: str, seed: int) -> Optional[Dict[str, Any]]:
    method = METHODS[method_key]
    out_root = Path(args.output_root)
    run_dir = out_root / "runs" / dataset_key / method.key / f"seed_{seed}"
    metrics_path = run_dir / "metrics.json"
    pred_path = out_root / "predictions" / dataset_key / f"{method.key}_seed_{seed}.csv"
    history_path = run_dir / "history.csv"
    ckpt_path = run_dir / "checkpoint.pt"
    if metrics_path.exists() and pred_path.exists() and not args.overwrite:
        print(f"SKIP completed {dataset_key} {method.key} seed={seed}", flush=True)
        return json.loads(metrics_path.read_text(encoding="utf-8"))

    print(f"RUN dataset={dataset_key} method={method.key} seed={seed}", flush=True)
    if args.dry_run:
        return None

    set_seed(seed)
    device = torch.device(args.device if torch.cuda.is_available() or args.device == "cpu" else "cpu")
    (
        model_cls,
        cfg,
        train_loader,
        val_loader,
        test_loader,
        train_samples,
        val_samples,
        test_samples,
        num_classes,
        label_names,
    ) = load_dataset(
        dataset_key=dataset_key,
        seed=seed,
        batch_size=args.batch_size,
        max_length=args.max_length,
        num_workers=args.num_workers,
        method=method,
        hfm_deleak_manifest=args.hfm_deleak_manifest,
        limits=(args.limit_train_samples, args.limit_val_samples, args.limit_test_samples),
    )
    cfg.update(
        {
            "architecture": method.architecture,
            "enable_clip_lora": method.enable_lora,
            "enable_text_lora": method.enable_lora,
            "seed": seed,
            "batch_size": args.batch_size,
            "max_length": args.max_length,
        }
    )
    model = model_cls(cfg).to(device)
    if hasattr(model, "vision_lora"):
        freeze_non_lora(model.vision_lora)
    if hasattr(model, "text"):
        freeze_non_lora(model.text)
    apply_method_trainability(model, method)
    stats = model.parameter_stats()

    train_info, best_state, history = train_one(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        method=method,
        device=device,
        epochs=args.epochs,
        patience=args.patience,
        grad_accum_steps=args.grad_accum_steps,
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
        num_classes=num_classes,
    )
    if best_state:
        model.load_state_dict(best_state, strict=False)

    criterion = nn.CrossEntropyLoss()
    test_metrics, prediction_rows, _logits, _labels = evaluate(
        model=model,
        loader=test_loader,
        samples=test_samples,
        method=method,
        device=device,
        criterion=criterion,
        num_classes=num_classes,
        dataset_key=dataset_key,
        seed=seed,
    )

    run_dir.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "model_state": {
                key: value.detach().cpu()
                for key, value in model.state_dict().items()
                if key in {name for name, parameter in model.named_parameters() if parameter.requires_grad}
            },
            "checkpoint_type": "trainable_parameters_only",
            "cfg": cfg,
            "dataset": dataset_key,
            "method": method.key,
            "seed": seed,
            "metrics": test_metrics,
            "train_info": train_info,
        },
        ckpt_path,
    )
    write_csv(history_path, history, ["epoch", "train_loss", "val_macro_f1", "val_weighted_f1", "val_accuracy", "val_raw_ece", "elapsed_seconds"])
    pred_fields = [
        "sample_id",
        "dataset",
        "split",
        "seed",
        "method",
        "true_label",
        "predicted_label",
    ]
    for c in range(num_classes):
        pred_fields.extend(
            [
                f"logit_class_{c}",
                f"prob_class_{c}",
                f"visual_prob_class_{c}",
                f"text_prob_class_{c}",
                f"disagreement_class_{c}",
            ]
        )
    pred_fields.extend(
        [
            "disagreement_l1",
            "prediction_entropy",
            "max_confidence",
            "primary_prediction",
            "primary_confidence",
            "refined_prediction",
            "refined_confidence",
            "correct_primary",
            "correct_refined",
            "number_of_images",
        ]
    )
    write_csv(pred_path, prediction_rows, pred_fields)

    metric_payload: Dict[str, Any] = {
        "dataset": dataset_key,
        "method": method.key,
        "method_display": method.display,
        "configuration": method.train_target,
        "disagreement_formulation": method.disagreement_formulation,
        "additional_loss": method.additional_loss,
        "lambda_primary": method.lambda_primary,
        "lambda_unimodal": method.lambda_unimodal,
        "consistency_beta": method.consistency_beta if method.additional_loss == "js_consistency" else None,
        "infonce_alpha": method.infonce_alpha if method.additional_loss == "infonce_alignment" else None,
        "infonce_temperature": method.infonce_temperature if method.additional_loss == "infonce_alignment" else None,
        "seed": seed,
        "num_classes": num_classes,
        "label_names": label_names,
        "train_samples": len(train_samples),
        "val_samples": len(val_samples),
        "test_samples": len(test_samples),
        "total_params": stats["total"],
        "trainable_params": stats["trainable"],
        "best_epoch": train_info["best_epoch"],
        **test_metrics,
    }
    metrics_path.write_text(json.dumps(metric_payload, indent=2), encoding="utf-8")
    print(
        f"DONE {dataset_key} {method.key} seed={seed} "
        f"mF1={test_metrics['macro_f1']:.4f} wF1={test_metrics['weighted_f1']:.4f} "
        f"ECE={test_metrics['raw_ece']:.4f}",
        flush=True,
    )
    return metric_payload


def aggregate(output_root: Path) -> None:
    metric_rows: List[Dict[str, Any]] = []
    for path in sorted((output_root / "runs").glob("*/*/seed_*/metrics.json")):
        try:
            metric_rows.append(json.loads(path.read_text(encoding="utf-8")))
        except Exception:
            continue
    fields = [
        "dataset",
        "method",
        "method_display",
        "configuration",
        "seed",
        "accuracy",
        "macro_precision",
        "macro_recall",
        "macro_f1",
        "weighted_f1",
        "auroc",
        "raw_ece",
        "adaptive_ece",
        "nll",
        "brier",
        "aurc",
        "risk_at_80_coverage",
        "risk_at_90_coverage",
        "accuracy_at_80_coverage",
        "accuracy_at_90_coverage",
        "total_params",
        "trainable_params",
        "best_epoch",
        "train_samples",
        "val_samples",
        "test_samples",
    ]
    write_csv(output_root / "aggregate_results" / "main_results.csv", metric_rows, fields)

    grouped: Dict[Tuple[str, str], List[Dict[str, Any]]] = {}
    for row in metric_rows:
        grouped.setdefault((row["dataset"], row["method"]), []).append(row)
    summary_rows: List[Dict[str, Any]] = []
    for (dataset, method), rows in sorted(grouped.items()):
        out: Dict[str, Any] = {"dataset": dataset, "method": method, "runs": len(rows)}
        for metric in ["accuracy", "macro_f1", "weighted_f1", "raw_ece", "nll", "brier", "aurc"]:
            vals = np.array([float(row[metric]) for row in rows if row.get(metric) is not None], dtype=float)
            if vals.size:
                out[f"{metric}_mean"] = float(vals.mean())
                out[f"{metric}_std"] = float(vals.std(ddof=1)) if vals.size > 1 else 0.0
        summary_rows.append(out)
    write_csv(
        output_root / "aggregate_results" / "summary_by_method.csv",
        summary_rows,
        [
            "dataset",
            "method",
            "runs",
            "accuracy_mean",
            "accuracy_std",
            "macro_f1_mean",
            "macro_f1_std",
            "weighted_f1_mean",
            "weighted_f1_std",
            "raw_ece_mean",
            "raw_ece_std",
            "nll_mean",
            "nll_std",
            "brier_mean",
            "brier_std",
            "aurc_mean",
            "aurc_std",
        ],
    )


def planned_runs(args: argparse.Namespace) -> List[Tuple[str, str, int]]:
    runs: List[Tuple[str, str, int]] = []
    if args.methods:
        method_keys = args.methods
    elif args.stage == "stage_b":
        method_keys = STAGE_B_METHODS
    elif args.stage == "stage_c":
        method_keys = STAGE_C_METHODS
    else:
        method_keys = sorted(set(STAGE_B_METHODS + STAGE_C_METHODS), key=(STAGE_B_METHODS + STAGE_C_METHODS).index)

    if args.stage == "stage_b":
        seeds = [args.stage_b_seed]
    else:
        seeds = args.seeds
    for dataset in args.datasets:
        for seed in seeds:
            for method in method_keys:
                runs.append((dataset, method, seed))
    return runs


def main() -> None:
    args = parse_args()
    output_root = Path(args.output_root)
    output_root.mkdir(parents=True, exist_ok=True)
    plan = planned_runs(args)
    print(f"Planned runs: {len(plan)}", flush=True)
    for dataset_key, method_key, seed in plan:
        run_single(args, dataset_key, method_key, seed)
        aggregate(output_root)
    aggregate(output_root)
    print(f"Wrote aggregate results to {output_root / 'aggregate_results'}", flush=True)


if __name__ == "__main__":
    main()