scripts/run_stress_tests.py

#!/usr/bin/env python3
"""Run controlled conflict/corruption stress tests from trained RAVEL checkpoints."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import sys
from dataclasses import replace
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from PIL import Image
from sklearn.metrics import accuracy_score, f1_score
from torch.utils.data import DataLoader, Dataset
from transformers import CLIPProcessor, DebertaV2Tokenizer
from transformers.utils import logging as hf_logging

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

from run_revised_experiments import (  # noqa: E402
    METHODS,
    apply_method_trainability,
    aurc_score,
    brier_score,
    evaluate,
    expected_calibration_error,
    freeze_non_lora,
    load_dataset,
    nll_score,
    set_seed,
)

hf_logging.set_verbosity_error()


DATASETS = ["mvsa_multiple", "hfm_deleak"]
METHODS_E08 = ["legacy_global", "token_aux", "param_mlp", "full_revised"]
SEEDS = [1, 3, 5, 7, 11]
CONDITIONS = [
    "original",
    "within_class_image_shuffle",
    "cross_class_image_shuffle",
    "within_class_text_shuffle",
    "cross_class_text_shuffle",
    "blank_image",
    "empty_text",
]
IMAGE_CACHE: Dict[str, torch.Tensor] = {}
CLIP_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32)
CLIP_STD = np.array([0.26862954, 0.26130258, 0.27577711], dtype=np.float32)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Run E08 stress tests from checkpoints.")
    parser.add_argument("--output-root", default="ravel_revision_results")
    parser.add_argument("--datasets", nargs="+", default=DATASETS, choices=DATASETS)
    parser.add_argument("--methods", nargs="+", default=METHODS_E08, choices=METHODS_E08)
    parser.add_argument("--seeds", nargs="+", type=int, default=SEEDS)
    parser.add_argument("--conditions", nargs="+", default=CONDITIONS, choices=CONDITIONS)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--batch-size", type=int, default=32)
    parser.add_argument("--max-length", type=int, default=None)
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--hfm-deleak-manifest", default="ravel_revision_results/data_audit/hfm_split_manifest_deleaked.csv")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--max-runs", type=int, default=None)
    return parser.parse_args()


def stable_seed(dataset: str, seed: int, condition: str) -> int:
    digest = hashlib.sha256(f"{dataset}:{seed}:{condition}".encode("utf-8")).hexdigest()
    return (int(digest[:12], 16) + int(seed)) % (2**32 - 1)


def sample_id(sample: Any, index: int) -> str:
    return str(getattr(sample, "image_id", getattr(sample, "sample_id", index)))


def sample_label(dataset: str, sample: Any) -> int:
    if dataset.startswith("mvsa"):
        mapping = {"positive": 0, "neutral": 1, "negative": 2}
        return int(mapping[str(sample.combined_majority).lower()])
    return int(sample.label)


def sample_text(sample: Any) -> str:
    return str(getattr(sample, "text", ""))


def image_tensor(image_path: str) -> torch.Tensor:
    key = str(image_path)
    cached = IMAGE_CACHE.get(key)
    if cached is not None:
        return cached
    try:
        image = Image.open(image_path).convert("RGB").resize((224, 224))
    except Exception:
        image = Image.new("RGB", (224, 224), (0, 0, 0))
    arr = np.asarray(image, dtype=np.float32) / 255.0
    arr = (arr - CLIP_MEAN) / CLIP_STD
    tensor = torch.from_numpy(np.transpose(arr, (2, 0, 1))).float()
    IMAGE_CACHE[key] = tensor
    return tensor


class StressPairDataset(Dataset):
    def __init__(self, dataset: str, samples: Sequence[Any]):
        self.dataset = dataset
        self.samples = list(samples)

    def __len__(self) -> int:
        return len(self.samples)

    def __getitem__(self, idx: int) -> Dict[str, Any]:
        sample = self.samples[idx]
        return {
            "pixel_values": image_tensor(sample.image_path),
            "text": sample_text(sample),
            "labels": sample_label(self.dataset, sample),
        }


def set_mvsa_text(sample: Any, text: str) -> Any:
    new_sample = replace(sample)
    setattr(new_sample, "_text_cache", text)
    return new_sample


def clone_with_image(dataset: str, sample: Any, image_path: str) -> Any:
    return replace(sample, image_path=image_path)


def clone_with_text(dataset: str, sample: Any, text: str) -> Any:
    if dataset.startswith("mvsa"):
        return set_mvsa_text(sample, text)
    return replace(sample, text=text)


def choose_donors(dataset: str, samples: Sequence[Any], seed: int, condition: str) -> List[int]:
    rng = np.random.default_rng(stable_seed(dataset, seed, condition))
    labels = np.array([sample_label(dataset, sample) for sample in samples])
    donor_indices: List[int] = []
    for idx, label in enumerate(labels):
        if condition.startswith("within_class"):
            candidates = np.flatnonzero(labels == label)
        else:
            candidates = np.flatnonzero(labels != label)
        candidates = candidates[candidates != idx]
        if candidates.size == 0:
            donor_indices.append(idx)
        else:
            donor_indices.append(int(rng.choice(candidates)))
    return donor_indices


def make_stress_samples(
    dataset: str,
    samples: Sequence[Any],
    seed: int,
    condition: str,
) -> Tuple[List[Any], List[str]]:
    if condition == "original":
        return list(samples), ["" for _ in samples]

    if condition == "blank_image":
        missing_path = str(PROJECT_ROOT / "ravel_revision_results" / "_stress_blank_missing_image.jpg")
        return [clone_with_image(dataset, sample, missing_path) for sample in samples], ["" for _ in samples]

    if condition == "empty_text":
        return [clone_with_text(dataset, sample, "") for sample in samples], ["" for _ in samples]

    donor_indices = choose_donors(dataset, samples, seed, condition)
    out: List[Any] = []
    donor_ids: List[str] = []
    for idx, donor_idx in enumerate(donor_indices):
        sample = samples[idx]
        donor = samples[donor_idx]
        donor_ids.append(sample_id(donor, donor_idx))
        if condition.endswith("image_shuffle"):
            out.append(clone_with_image(dataset, sample, donor.image_path))
        elif condition.endswith("text_shuffle"):
            out.append(clone_with_text(dataset, sample, sample_text(donor)))
        else:
            raise ValueError(condition)
    return out, donor_ids


def build_loader(
    dataset: str,
    samples: List[Any],
    cfg: Dict[str, Any],
    batch_size: int,
    max_length: int,
    num_workers: int,
) -> DataLoader:
    tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])
    def collate_fn(batch: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
        texts = [item["text"] for item in batch]
        labels = torch.tensor([item["labels"] for item in batch], dtype=torch.long)
        text_inputs = tokenizer(
            texts,
            max_length=max_length,
            padding=True,
            truncation=True,
            return_tensors="pt",
        )
        return {
            "pixel_values": torch.stack([item["pixel_values"] for item in batch], dim=0),
            "input_ids": text_inputs["input_ids"],
            "attention_mask": text_inputs["attention_mask"],
            "labels": labels,
        }

    return DataLoader(
        StressPairDataset(dataset, samples),
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
        collate_fn=collate_fn,
    )


def prob_columns(df: pd.DataFrame) -> List[str]:
    cols = [col for col in df.columns if col.startswith("prob_class_")]
    if not cols:
        cols = [col for col in df.columns if col.startswith("prob_") and col[5:].isdigit()]
    return sorted(cols, key=lambda name: int(name.rsplit("_", 1)[-1]))


def class_vector(row: pd.Series, prefix: str, num_classes: int) -> List[Optional[float]]:
    out: List[Optional[float]] = []
    for c in range(num_classes):
        for name in (f"{prefix}{c}", f"{prefix}_class_{c}"):
            if name in row and pd.notna(row[name]):
                out.append(float(row[name]))
                break
        else:
            out.append(None)
    return out


def to_optional_float(value: Any) -> float:
    try:
        if value == "":
            return float("nan")
        return float(value)
    except Exception:
        return float("nan")


def rows_to_stress_frame(
    rows: List[Dict[str, Any]],
    dataset: str,
    method: str,
    seed: int,
    condition: str,
    donor_ids: Sequence[str],
    num_classes: int,
) -> pd.DataFrame:
    out_rows: List[Dict[str, Any]] = []
    for idx, row in enumerate(rows):
        probs = [float(row.get(f"prob_class_{c}", np.nan)) for c in range(num_classes)]
        visual_probs = [to_optional_float(row.get(f"visual_prob_class_{c}", np.nan)) for c in range(num_classes)]
        text_probs = [to_optional_float(row.get(f"text_prob_class_{c}", np.nan)) for c in range(num_classes)]
        true_label = int(row["true_label"])
        pred = int(row["predicted_label"])
        confidence = float(np.nanmax(probs))
        nll = -math.log(max(float(probs[true_label]), 1e-12))
        if all(np.isfinite(float(x)) for x in visual_probs + text_probs):
            tv = 0.5 * float(np.sum(np.abs(np.array(visual_probs, dtype=float) - np.array(text_probs, dtype=float))))
        else:
            tv = float("nan")
        payload: Dict[str, Any] = {
            "sample_id": row["sample_id"],
            "donor_sample_id": donor_ids[idx] if idx < len(donor_ids) else "",
            "dataset": dataset,
            "condition": condition,
            "method": method,
            "seed": seed,
            "true_label": true_label,
            "predicted_label": pred,
            "confidence": confidence,
            "entropy": float(row.get("prediction_entropy", np.nan)),
            "raw_ece_component": abs(float(pred == true_label) - confidence),
            "nll": nll,
            "visual_probs": json.dumps([None if pd.isna(x) else float(x) for x in visual_probs]),
            "text_probs": json.dumps([None if pd.isna(x) else float(x) for x in text_probs]),
            "tv_disagreement": tv,
        }
        for c in range(num_classes):
            payload[f"logit_{c}"] = row.get(f"logit_class_{c}", np.nan)
            payload[f"prob_{c}"] = probs[c]
            payload[f"visual_prob_{c}"] = visual_probs[c]
            payload[f"text_prob_{c}"] = text_probs[c]
        out_rows.append(payload)
    return pd.DataFrame(out_rows)


def original_to_stress_frame(path: Path, dataset: str, method: str, seed: int) -> pd.DataFrame:
    df = pd.read_csv(path)
    pcols = prob_columns(df)
    num_classes = len(pcols)
    rows: List[Dict[str, Any]] = []
    for _, row in df.iterrows():
        probs = [float(row[col]) for col in pcols]
        true_label = int(row["true_label"])
        pred = int(row["predicted_label"])
        confidence = float(np.max(probs))
        visual_probs = class_vector(row, "visual_prob", num_classes)
        text_probs = class_vector(row, "text_prob", num_classes)
        if all(x is not None and not pd.isna(x) for x in visual_probs + text_probs):
            tv = 0.5 * float(np.sum(np.abs(np.array(visual_probs, dtype=float) - np.array(text_probs, dtype=float))))
        else:
            tv = float("nan")
        payload: Dict[str, Any] = {
            "sample_id": row["sample_id"],
            "donor_sample_id": "",
            "dataset": dataset,
            "condition": "original",
            "method": method,
            "seed": seed,
            "true_label": true_label,
            "predicted_label": pred,
            "confidence": confidence,
            "entropy": float(row.get("prediction_entropy", -(np.array(probs) * np.log(np.clip(probs, 1e-12, 1))).sum())),
            "raw_ece_component": abs(float(pred == true_label) - confidence),
            "nll": -math.log(max(float(probs[true_label]), 1e-12)),
            "visual_probs": json.dumps([None if x is None or pd.isna(x) else float(x) for x in visual_probs]),
            "text_probs": json.dumps([None if x is None or pd.isna(x) else float(x) for x in text_probs]),
            "tv_disagreement": tv,
        }
        for c in range(num_classes):
            payload[f"logit_{c}"] = row.get(f"logit_class_{c}", np.nan)
            payload[f"prob_{c}"] = probs[c]
            payload[f"visual_prob_{c}"] = visual_probs[c]
            payload[f"text_prob_{c}"] = text_probs[c]
        rows.append(payload)
    return pd.DataFrame(rows)


def load_model_and_data(
    args: argparse.Namespace,
    dataset: str,
    method_key: str,
    seed: int,
) -> Tuple[nn.Module, Dict[str, Any], List[Any], int, int, torch.device]:
    root = Path(args.output_root)
    run_dir = root / "runs" / dataset / method_key / f"seed_{seed}"
    ckpt_path = run_dir / "checkpoint.pt"
    if not ckpt_path.exists():
        raise FileNotFoundError(ckpt_path)
    ckpt = torch.load(ckpt_path, map_location="cpu")
    ckpt_cfg = ckpt.get("cfg", {}) if isinstance(ckpt, dict) else {}
    max_length = int(args.max_length or ckpt_cfg.get("max_length") or 96)
    method = METHODS[method_key]
    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,
        seed=seed,
        batch_size=args.batch_size,
        max_length=max_length,
        num_workers=0,
        method=method,
        hfm_deleak_manifest=args.hfm_deleak_manifest,
        limits=(None, None, None),
    )
    cfg.update(ckpt_cfg)
    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": 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)
    model.load_state_dict(ckpt.get("model_state", {}), strict=False)
    model.eval()
    return model, cfg, test_samples, num_classes, max_length, device


def run_condition(
    args: argparse.Namespace,
    model: nn.Module,
    cfg: Dict[str, Any],
    test_samples: List[Any],
    num_classes: int,
    max_length: int,
    device: torch.device,
    dataset: str,
    method: str,
    seed: int,
    condition: str,
) -> pd.DataFrame:
    stress_samples, donor_ids = make_stress_samples(dataset, test_samples, seed, condition)
    loader = build_loader(dataset, stress_samples, cfg, args.batch_size, max_length, args.num_workers)
    criterion = nn.CrossEntropyLoss()
    metrics, rows, _logits, _labels = evaluate(
        model=model,
        loader=loader,
        samples=test_samples,
        method=METHODS[method],
        device=device,
        criterion=criterion,
        num_classes=num_classes,
        dataset_key=dataset,
        seed=seed,
    )
    df = rows_to_stress_frame(rows, dataset, method, seed, condition, donor_ids, num_classes)
    df.attrs["metrics"] = metrics
    return df


def metrics_from_frame(df: pd.DataFrame) -> Dict[str, float]:
    pcols = prob_columns(df)
    probs = df[pcols].astype(float).to_numpy()
    y = df["true_label"].astype(int).to_numpy()
    pred = df["predicted_label"].astype(int).to_numpy()
    num_classes = probs.shape[1]
    aurc, *_ = aurc_score(probs, y)
    return {
        "accuracy": float(accuracy_score(y, pred)),
        "macro_f1": float(f1_score(y, pred, average="macro", zero_division=0)),
        "weighted_f1": float(f1_score(y, pred, average="weighted", zero_division=0)),
        "raw_ece": expected_calibration_error(probs, y),
        "nll": nll_score(probs, y),
        "brier": brier_score(probs, y, num_classes),
        "aurc": float(aurc),
        "mean_confidence": float(df["confidence"].astype(float).mean()),
        "mean_entropy": float(df["entropy"].astype(float).mean()),
        "mean_tv_disagreement": float(df["tv_disagreement"].astype(float).mean()) if df["tv_disagreement"].notna().any() else float("nan"),
    }


def write_original_if_needed(args: argparse.Namespace, dataset: str, method: str, seed: int) -> None:
    root = Path(args.output_root)
    out_path = root / "stress_predictions" / dataset / "original" / f"{method}_seed_{seed}.csv"
    if out_path.exists() and not args.overwrite:
        return
    src = root / "predictions" / dataset / f"{method}_seed_{seed}.csv"
    if not src.exists():
        raise FileNotFoundError(src)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    original_to_stress_frame(src, dataset, method, seed).to_csv(out_path, index=False)


def aggregate_outputs(output_root: Path, datasets: Sequence[str], methods: Sequence[str], seeds: Sequence[int], conditions: Sequence[str]) -> None:
    rows: List[Dict[str, Any]] = []
    for dataset in datasets:
        for condition in conditions:
            for method in methods:
                for seed in seeds:
                    path = output_root / "stress_predictions" / dataset / condition / f"{method}_seed_{seed}.csv"
                    if not path.exists():
                        continue
                    df = pd.read_csv(path)
                    if df.empty:
                        continue
                    metrics = metrics_from_frame(df)
                    rows.append(
                        {
                            "dataset": dataset,
                            "condition": condition,
                            "method": method,
                            "seed": seed,
                            "accuracy": metrics["accuracy"],
                            "macro_f1": metrics["macro_f1"],
                            "weighted_f1": metrics["weighted_f1"],
                            "f1": metrics["weighted_f1"] if dataset.startswith("mvsa") else metrics["macro_f1"],
                            "raw_ece": metrics["raw_ece"],
                            "nll": metrics["nll"],
                            "brier": metrics["brier"],
                            "aurc": metrics["aurc"],
                            "mean_confidence": metrics["mean_confidence"],
                            "mean_entropy": metrics["mean_entropy"],
                            "prediction_entropy": metrics["mean_entropy"],
                            "mean_tv_disagreement": metrics["mean_tv_disagreement"],
                            "status": "COMPLETE",
                        }
                    )
    agg_dir = output_root / "aggregate_results"
    agg_dir.mkdir(parents=True, exist_ok=True)
    result = pd.DataFrame(rows)
    result.to_csv(agg_dir / "stress_test_results.csv", index=False)
    if result.empty:
        pd.DataFrame().to_csv(agg_dir / "stress_test_summary.csv", index=False)
        return
    summary_rows: List[Dict[str, Any]] = []
    for (dataset, condition, method), group in result.groupby(["dataset", "condition", "method"], dropna=False):
        out: Dict[str, Any] = {
            "dataset": dataset,
            "condition": condition,
            "method": method,
            "num_seeds": int(group["seed"].nunique()),
            "status": "COMPLETE" if int(group["seed"].nunique()) >= 5 else "PARTIAL",
        }
        for metric in ["accuracy", "macro_f1", "weighted_f1", "f1", "raw_ece", "nll", "brier", "aurc", "mean_confidence", "mean_entropy", "mean_tv_disagreement"]:
            vals = group[metric].astype(float)
            out[f"{metric}_mean"] = float(vals.mean())
            out[f"{metric}_std"] = float(vals.std(ddof=1)) if len(vals) > 1 else 0.0
        out["prediction_entropy_mean"] = out["mean_entropy_mean"]
        out["prediction_entropy_std"] = out["mean_entropy_std"]
        summary_rows.append(out)
    pd.DataFrame(summary_rows).to_csv(agg_dir / "stress_test_summary.csv", index=False)


def main() -> None:
    args = parse_args()
    root = Path(args.output_root)
    planned = [(d, m, s) for d in args.datasets for m in args.methods for s in args.seeds]
    if args.max_runs is not None:
        planned = planned[: args.max_runs]
    print(f"Planned E08 model runs: {len(planned)}", flush=True)
    for dataset, method, seed in planned:
        for condition in args.conditions:
            if condition == "original":
                write_original_if_needed(args, dataset, method, seed)
        pending = [
            condition
            for condition in args.conditions
            if condition != "original"
            and (
                args.overwrite
                or not (root / "stress_predictions" / dataset / condition / f"{method}_seed_{seed}.csv").exists()
            )
        ]
        if not pending:
            print(f"SKIP E08 {dataset} {method} seed={seed}", flush=True)
            continue
        print(f"RUN E08 dataset={dataset} method={method} seed={seed} conditions={','.join(pending)}", flush=True)
        model, cfg, test_samples, num_classes, max_length, device = load_model_and_data(args, dataset, method, seed)
        for condition in pending:
            out_path = root / "stress_predictions" / dataset / condition / f"{method}_seed_{seed}.csv"
            df = run_condition(args, model, cfg, test_samples, num_classes, max_length, device, dataset, method, seed, condition)
            out_path.parent.mkdir(parents=True, exist_ok=True)
            df.to_csv(out_path, index=False)
            metrics = metrics_from_frame(df)
            print(
                f"DONE E08 {dataset} {condition} {method} seed={seed} "
                f"F1={metrics['weighted_f1' if dataset.startswith('mvsa') else 'macro_f1']:.4f} "
                f"ECE={metrics['raw_ece']:.4f} conf={metrics['mean_confidence']:.4f}",
                flush=True,
            )
        del model
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    aggregate_outputs(root, args.datasets, args.methods, args.seeds, args.conditions)
    rows = pd.read_csv(root / "aggregate_results" / "stress_test_results.csv")
    print(f"E08 stress rows: {len(rows)}", flush=True)


if __name__ == "__main__":
    main()