CAGE_expression_inference-apvit/models/AffectNet8_Maxvit_VA/train2.py

"""
train2.py — Ablation study: 6 experiments (backbone comparison)

Pipeline IDENTICAL to train.py (train on full train_set, validate on official
val_set, save best model by lowest val loss, 20 epochs).
Only the model architecture differs across experiments.

All models use ImageNet pretrained weights + standard VA regression head + CCC loss.

Experiments:
  1) VGG-16           + CCC only
  2) MobileNetV2      + CCC only
  3) ResNet-18        + CCC only
  4) EfficientNetV2-S + CCC only
  5) Swin-T           + CCC only
  6) ConvNeXt-Tiny    + CCC only

Output per experiment:
  experiments/<exp_name>/
    ├── model.pt
    ├── train_history.csv
    └── *.png (10 plots)

Summary:
  experiments/
    ├── all_experiments_summary.csv
    └── summary_*.png
"""

import pandas as pd
import os
import numpy as np
import torch
import torchvision.transforms as transforms
import torchvision.models as models
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
import torch.optim as optim
from PIL import Image
from torch.optim import lr_scheduler
from tqdm import tqdm
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt


# ============================================================================
# Paths (identical to train.py)
# ============================================================================
SCRIPT_DIR = Path(__file__).resolve().parent
EXPERIMENTS_DIR = SCRIPT_DIR / "experiments"
EXPERIMENTS_DIR.mkdir(parents=True, exist_ok=True)


def resolve_images_dir(split_name, explicit_env_var, data_root, extract_root):
    explicit = os.getenv(explicit_env_var, "").strip()
    candidates = []
    if explicit:
        candidates.append(Path(explicit))

    candidates.append(Path(data_root) / f"{split_name}_set" / "images")
    candidates.append(Path(extract_root) / f"{split_name}_extracted" / "images")

    extract_split_root = Path(extract_root) / f"{split_name}_extracted"
    if extract_split_root.exists():
        for p in extract_split_root.rglob("images"):
            if p.is_dir():
                candidates.append(p)

    for p in candidates:
        if p.exists() and p.is_dir():
            return str(p)

    tried = "\n".join([str(p) for p in candidates])
    raise FileNotFoundError(
        f"Could not find images folder for split='{split_name}'. Tried:\n{tried}"
    )


# ============================================================================
# Data paths (identical to train.py)
# ============================================================================
DATA_ROOT = os.getenv("AFFECTNET_ROOT", "/workspace/data_affectnet/AffectNet")
EXTRACT_ROOT = os.getenv("AFFECTNET_EXTRACT_ROOT", f"{DATA_ROOT}/extracted")
ANNO_ROOT = os.getenv("AFFECTNET_ANNO_ROOT", "../../affectnet_annotations")
IMAGE_FOLDER = resolve_images_dir("train", "AFFECTNET_TRAIN_IMAGES", DATA_ROOT, EXTRACT_ROOT)
IMAGE_FOLDER_TEST = resolve_images_dir("val", "AFFECTNET_VAL_IMAGES", DATA_ROOT, EXTRACT_ROOT)
train_annotations_path = os.getenv(
    "AFFECTNET_TRAIN_ANNO", f"{ANNO_ROOT}/train_set_annotation_without_lnd.csv"
)
valid_annotations_path = os.getenv(
    "AFFECTNET_VAL_ANNO", f"{ANNO_ROOT}/val_set_annotation_without_lnd.csv"
)
train_annotations_df = pd.read_csv(train_annotations_path)
valid_annotations_df = pd.read_csv(valid_annotations_path)


# ============================================================================
# Parameters (identical to train.py)
# ============================================================================
BATCHSIZE = int(os.getenv("BATCHSIZE", "64"))
NUM_EPOCHS = 20
LR = 4e-5
NUM_WORKERS = int(os.getenv("NUM_WORKERS", "0"))
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("DATA_ROOT              =", DATA_ROOT)
print("EXTRACT_ROOT           =", EXTRACT_ROOT)
print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER)
print("AFFECTNET_VAL_IMAGES   =", IMAGE_FOLDER_TEST)
print("EXPERIMENTS_DIR        =", EXPERIMENTS_DIR)
print("Train set size         =", len(train_annotations_df))
print("Val set size           =", len(valid_annotations_df))
print("DEVICE                 =", DEVICE)
print("BATCH_SIZE             =", BATCHSIZE)
print("NUM_EPOCHS             =", NUM_EPOCHS)
print("LR                     =", LR)


# ============================================================================
# 6 Experiment configurations (backbone comparison, all CCC loss)
# ============================================================================
EXPERIMENTS = [
    {
        "name": "exp1_vgg16_ccc",
        "desc": "VGG-16 + ImageNet + CCC only",
        "model_type": "vgg16",
        "loss_type": "ccc_only",
    },
    {
        "name": "exp2_mobilenetv2_ccc",
        "desc": "MobileNetV2 + ImageNet + CCC only",
        "model_type": "mobilenetv2",
        "loss_type": "ccc_only",
    },
    {
        "name": "exp3_resnet18_ccc",
        "desc": "ResNet-18 + ImageNet + CCC only",
        "model_type": "resnet18",
        "loss_type": "ccc_only",
    },
    {
        "name": "exp4_efficientnetv2_ccc",
        "desc": "EfficientNetV2-S + ImageNet + CCC only",
        "model_type": "efficientnetv2",
        "loss_type": "ccc_only",
    },
    {
        "name": "exp5_swin_ccc",
        "desc": "Swin-T + ImageNet + CCC only",
        "model_type": "swin",
        "loss_type": "ccc_only",
    },
    {
        "name": "exp6_convnext_ccc",
        "desc": "ConvNeXt-Tiny + ImageNet + CCC only",
        "model_type": "convnext",
        "loss_type": "ccc_only",
    },
]


# ============================================================================
# Plot helpers (from train.py)
# ============================================================================
def save_line_plot(x, ys, labels, title, xlabel, ylabel, save_path):
    plt.figure(figsize=(10, 6))
    for y, label in zip(ys, labels):
        plt.plot(x, y, marker="o", linewidth=2, markersize=4, label=label)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_bar_plot(labels, values, title, xlabel, ylabel, save_path, annotate=True):
    plt.figure(figsize=(12, 6))
    bars = plt.bar(labels, values)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, axis="y", alpha=0.3)
    if annotate:
        for bar, val in zip(bars, values):
            plt.text(
                bar.get_x() + bar.get_width() / 2,
                bar.get_height(),
                f"{val:.4f}",
                ha="center",
                va="bottom",
                fontsize=8,
            )
    plt.xticks(rotation=30, ha="right")
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_scatter(x, y, title, xlabel, ylabel, save_path, annotate_labels=None):
    plt.figure(figsize=(10, 6))
    plt.scatter(x, y)
    if annotate_labels is not None:
        for xi, yi, txt in zip(x, y, annotate_labels):
            plt.annotate(str(txt), (xi, yi), fontsize=9)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


# ============================================================================
# Dataset & transforms (identical to train.py)
# ============================================================================
class CustomDataset(Dataset):
    def __init__(self, dataframe, root_dir, transform=None, balance=False):
        self.dataframe = dataframe
        self.transform = transform
        self.root_dir = root_dir
        self.balance = balance

        if self.balance:
            self.dataframe = self.balance_dataset()

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

    def __getitem__(self, idx):
        img_id = int(float(self.dataframe["number"].iloc[idx]))
        image_path = os.path.join(self.root_dir, f"{img_id}.jpg")
        image = Image.open(image_path)

        classes = torch.tensor(self.dataframe.iloc[idx, 1], dtype=torch.int8)
        valence = torch.tensor(self.dataframe.iloc[idx, 2], dtype=torch.float16)
        arousal = torch.tensor(self.dataframe.iloc[idx, 3], dtype=torch.float16)

        if self.transform:
            image = self.transform(image)

        return image, classes, valence, arousal

    def balance_dataset(self):
        balanced_df = self.dataframe.groupby("exp", group_keys=False).apply(
            lambda x: x.sample(self.dataframe["exp"].value_counts().min())
        )
        return balanced_df


transform = transforms.Compose(
    [
        transforms.RandomHorizontalFlip(0.5),
        transforms.RandomGrayscale(0.01),
        transforms.RandomRotation(10),
        transforms.ColorJitter(
            brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1
        ),
        transforms.RandomPerspective(
            distortion_scale=0.2, p=0.5
        ),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        transforms.RandomErasing(
            p=0.5, scale=(0.02, 0.2), ratio=(0.3, 3.3), value="random"
        ),
    ]
)

transform_valid = transforms.Compose(
    [
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ]
)


# ============================================================================
# Model builders — all ImageNet pretrained, standard VA regression heads
# ============================================================================
def _make_head(in_features, dropout=0.3):
    """Standard VA regression head used by all models."""
    return nn.Sequential(
        nn.LayerNorm(in_features),
        nn.Linear(in_features, in_features),
        nn.Tanh(),
        nn.Dropout(dropout),
        nn.Linear(in_features, 2, bias=False),
    )


def build_model(exp_config):
    """
    Build model based on experiment config.
    All models use ImageNet pretrained weights + standard VA regression head.
    """
    model_type = exp_config["model_type"]

    if model_type == "vgg16":
        # VGG-16
        model = models.vgg16(weights="DEFAULT")
        in_feat = model.classifier[6].in_features  # 4096
        model.classifier[6] = _make_head(in_feat)

    elif model_type == "mobilenetv2":
        # MobileNetV2
        model = models.mobilenet_v2(weights="DEFAULT")
        in_feat = model.classifier[1].in_features  # 1280
        model.classifier = _make_head(in_feat)

    elif model_type == "resnet18":
        # ResNet-18
        model = models.resnet18(weights="DEFAULT")
        in_feat = model.fc.in_features  # 512
        model.fc = _make_head(in_feat)

    elif model_type == "efficientnetv2":
        # EfficientNetV2-S
        model = models.efficientnet_v2_s(weights="DEFAULT")
        in_feat = model.classifier[1].in_features  # 1280
        model.classifier = _make_head(in_feat)

    elif model_type == "swin":
        # Swin Transformer Tiny
        model = models.swin_t(weights="DEFAULT")
        in_feat = model.head.in_features  # 768
        model.head = _make_head(in_feat)

    elif model_type == "convnext":
        # ConvNeXt Tiny
        model = models.convnext_tiny(weights="DEFAULT")
        in_feat = model.classifier[2].in_features  # 768
        model.classifier[2] = _make_head(in_feat)

    else:
        raise ValueError(f"Unknown model_type: {model_type}")

    model.to(DEVICE)
    return model


# ============================================================================
# Loss functions & metrics (from train.py)
# ============================================================================
val_loss_fn = nn.MSELoss()
aro_loss_fn = nn.MSELoss()


def CCCLoss(x, y):
    """CCC-based loss (returns -CCC). Identical to train.py."""
    x_mean = torch.mean(x, dim=0)
    y_mean = torch.mean(y, dim=0)
    x_var = torch.var(x, dim=0)
    y_var = torch.var(y, dim=0)
    cov_matrix = torch.matmul(
        (x - x_mean).permute(*torch.arange(x.dim() - 1, -1, -1)), y - y_mean
    ) / (x.size(0) - 1)
    numerator = 2 * cov_matrix
    denominator = x_var + y_var + torch.pow((x_mean - y_mean), 2)
    ccc = torch.mean(numerator / denominator)
    return -ccc


def compute_loss(val_pred, aro_pred, val_true, aro_true, loss_type):
    """Compute training/validation loss based on experiment type."""
    if loss_type == "mse_ccc":
        return (
            3 * val_loss_fn(val_pred, val_true)
            + 3 * aro_loss_fn(aro_pred, aro_true)
            + CCCLoss(val_pred, val_true)
            + CCCLoss(aro_pred, aro_true)
        )
    else:  # ccc_only
        return (
            CCCLoss(val_pred, val_true)
            + CCCLoss(aro_pred, aro_true)
        )


def ccc_numpy_1d(x, y, eps=1e-8):
    """CCC metric (numpy). Identical to train.py."""
    x = np.asarray(x).reshape(-1)
    y = np.asarray(y).reshape(-1)
    x_mean = x.mean()
    y_mean = y.mean()
    x_var = x.var()
    y_var = y.var()
    cov = ((x - x_mean) * (y - y_mean)).mean()
    return (2.0 * cov) / (x_var + y_var + (x_mean - y_mean) ** 2 + eps)


# ============================================================================
# Per-experiment plots (same plots as train.py)
# ============================================================================
def save_experiment_plots(history_df, exp_name, exp_dir):
    """Save the same set of plots as train.py for each experiment."""
    epochs = history_df["epoch"].tolist()

    # 1. Train Loss
    save_line_plot(
        epochs, [history_df["train_loss"].tolist()], ["Train Loss"],
        f"{exp_name} — Train Loss", "Epoch", "Loss",
        exp_dir / "train_loss.png",
    )
    # 2. Validation Loss
    save_line_plot(
        epochs, [history_df["valid_loss"].tolist()], ["Validation Loss"],
        f"{exp_name} — Validation Loss", "Epoch", "Loss",
        exp_dir / "valid_loss.png",
    )
    # 3. Train vs Validation Loss
    save_line_plot(
        epochs,
        [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()],
        ["Train Loss", "Validation Loss"],
        f"{exp_name} — Train vs Validation Loss", "Epoch", "Loss",
        exp_dir / "train_vs_valid_loss.png",
    )
    # 4. CCC_v
    save_line_plot(
        epochs, [history_df["ccc_v"].tolist()], ["CCC_v"],
        f"{exp_name} — Validation CCC_v", "Epoch", "CCC",
        exp_dir / "valid_ccc_v.png",
    )
    # 5. CCC_a
    save_line_plot(
        epochs, [history_df["ccc_a"].tolist()], ["CCC_a"],
        f"{exp_name} — Validation CCC_a", "Epoch", "CCC",
        exp_dir / "valid_ccc_a.png",
    )
    # 6. CCC_mean
    save_line_plot(
        epochs, [history_df["ccc_mean"].tolist()], ["CCC_mean"],
        f"{exp_name} — Validation CCC_mean", "Epoch", "CCC",
        exp_dir / "valid_ccc_mean.png",
    )
    # 7. All CCC on one plot
    save_line_plot(
        epochs,
        [
            history_df["ccc_v"].tolist(),
            history_df["ccc_a"].tolist(),
            history_df["ccc_mean"].tolist(),
        ],
        ["CCC_v", "CCC_a", "CCC_mean"],
        f"{exp_name} — All CCC Curves", "Epoch", "CCC",
        exp_dir / "valid_ccc_all.png",
    )
    # 8. Learning Rate
    save_line_plot(
        epochs, [history_df["lr"].tolist()], ["Learning Rate"],
        f"{exp_name} — Learning Rate", "Epoch", "LR",
        exp_dir / "lr.png",
    )
    # 9. Train Loss vs Valid Loss scatter
    save_scatter(
        history_df["train_loss"].tolist(),
        history_df["valid_loss"].tolist(),
        f"{exp_name} — Train Loss vs Valid Loss",
        "Train Loss", "Validation Loss",
        exp_dir / "train_vs_valid_scatter.png",
        annotate_labels=epochs,
    )
    # 10. CCC_v vs CCC_a scatter
    save_scatter(
        history_df["ccc_v"].tolist(),
        history_df["ccc_a"].tolist(),
        f"{exp_name} — CCC_v vs CCC_a",
        "CCC_v", "CCC_a",
        exp_dir / "ccc_v_vs_ccc_a_scatter.png",
        annotate_labels=epochs,
    )


# ============================================================================
# Summary plots across all experiments
# ============================================================================
def save_summary_plots(summary_df, save_dir):
    """Save comparison plots across all experiments."""
    names = summary_df["name"].tolist()

    # Bar: Best CCC_v
    save_bar_plot(
        names, summary_df["best_ccc_v"].tolist(),
        "Best CCC_v by Experiment", "Experiment", "CCC_v",
        save_dir / "summary_best_ccc_v.png",
    )
    # Bar: Best CCC_a
    save_bar_plot(
        names, summary_df["best_ccc_a"].tolist(),
        "Best CCC_a by Experiment", "Experiment", "CCC_a",
        save_dir / "summary_best_ccc_a.png",
    )
    # Bar: Best CCC_mean
    save_bar_plot(
        names, summary_df["best_ccc_mean"].tolist(),
        "Best CCC_mean by Experiment", "Experiment", "CCC_mean",
        save_dir / "summary_best_ccc_mean.png",
    )
    # Bar: Best Validation Loss
    save_bar_plot(
        names, summary_df["best_valid_loss"].tolist(),
        "Best Validation Loss by Experiment", "Experiment", "Loss",
        save_dir / "summary_best_valid_loss.png",
    )
    # Bar: Best Epoch
    save_bar_plot(
        names, summary_df["best_epoch"].tolist(),
        "Best Epoch by Experiment", "Experiment", "Epoch",
        save_dir / "summary_best_epoch.png",
    )

    # Overlay: CCC_mean curves from all experiments
    plt.figure(figsize=(12, 6))
    for _, row in summary_df.iterrows():
        hist_csv = save_dir / row["name"] / "train_history.csv"
        if hist_csv.exists():
            h = pd.read_csv(hist_csv)
            plt.plot(
                h["epoch"], h["ccc_mean"],
                marker="o", linewidth=2, markersize=3,
                label=row["name"],
            )
    plt.title("Validation CCC_mean — All Experiments")
    plt.xlabel("Epoch")
    plt.ylabel("CCC_mean")
    plt.grid(True, alpha=0.3)
    plt.legend(fontsize=8)
    plt.tight_layout()
    plt.savefig(save_dir / "summary_ccc_mean_overlay.png", dpi=200)
    plt.close()

    # Overlay: Validation Loss curves from all experiments
    plt.figure(figsize=(12, 6))
    for _, row in summary_df.iterrows():
        hist_csv = save_dir / row["name"] / "train_history.csv"
        if hist_csv.exists():
            h = pd.read_csv(hist_csv)
            plt.plot(
                h["epoch"], h["valid_loss"],
                marker="o", linewidth=2, markersize=3,
                label=row["name"],
            )
    plt.title("Validation Loss — All Experiments")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.grid(True, alpha=0.3)
    plt.legend(fontsize=8)
    plt.tight_layout()
    plt.savefig(save_dir / "summary_valid_loss_overlay.png", dpi=200)
    plt.close()

    # Overlay: Train Loss curves from all experiments
    plt.figure(figsize=(12, 6))
    for _, row in summary_df.iterrows():
        hist_csv = save_dir / row["name"] / "train_history.csv"
        if hist_csv.exists():
            h = pd.read_csv(hist_csv)
            plt.plot(
                h["epoch"], h["train_loss"],
                marker="o", linewidth=2, markersize=3,
                label=row["name"],
            )
    plt.title("Train Loss — All Experiments")
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.grid(True, alpha=0.3)
    plt.legend(fontsize=8)
    plt.tight_layout()
    plt.savefig(save_dir / "summary_train_loss_overlay.png", dpi=200)
    plt.close()


# ============================================================================
# Training function — mirrors train.py's training loop EXACTLY
# ============================================================================
def run_experiment(exp_config, train_loader, valid_loader):
    """
    Run one experiment. Training loop structure is identical to train.py:
    - AdamW optimizer, CosineAnnealingLR(T_max=BATCHSIZE*NUM_EPOCHS)
    - AMP GradScaler + autocast
    - l2_reg/l1_reg computed per batch (not added to loss, kept for parity)
    - scheduler.step() called per batch (inside autocast, same as train.py)
    - Best model selected by minimum validation loss
    """
    name = exp_config["name"]
    desc = exp_config["desc"]
    loss_type = exp_config["loss_type"]

    exp_dir = EXPERIMENTS_DIR / name
    exp_dir.mkdir(parents=True, exist_ok=True)
    exp_model_path = str(exp_dir / "model.pt")

    print(f"\n{'='*70}")
    print(f"EXPERIMENT: {name}")
    print(f"  {desc}")
    print(f"  model_type={exp_config['model_type']}, "
          f"loss_type={loss_type}")
    print(f"  output: {exp_dir}")
    print(f"{'='*70}")

    # Build model
    model = build_model(exp_config)

    # Optimizer, scheduler, scaler — identical to train.py
    optimizer = optim.AdamW(model.parameters(), lr=LR)
    scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS)
    scaler = torch.cuda.amp.GradScaler()

    best_valid_loss = float("inf")
    best_epoch = -1
    l2_lambda = 0.00001  # kept for parity with train.py (not added to loss)
    l1_lambda = 0.00001

    history = {
        "epoch": [],
        "train_loss": [],
        "valid_loss": [],
        "ccc_v": [],
        "ccc_a": [],
        "ccc_mean": [],
        "lr": [],
    }

    for epoch in range(NUM_EPOCHS):
        # ── Train (identical structure to train.py) ──────────────────────
        model.train()
        train_loss = 0.0
        current_lr = optimizer.param_groups[0]["lr"]

        for images, _, val_true, aro_true in tqdm(
            train_loader, desc=f"[{name}] Epoch {epoch+1}/{NUM_EPOCHS}"
        ):
            images, val_true, aro_true = (
                images.to(DEVICE), val_true.to(DEVICE), aro_true.to(DEVICE)
            )
            optimizer.zero_grad()
            l2_reg = 0
            l1_reg = 0
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                outputs = model(images)
                val_pred = outputs[:, 0]
                aro_pred = outputs[:, 1]
                for param in model.parameters():
                    l2_reg += torch.norm(param, 2)
                    l1_reg += torch.norm(param, 1)
                loss = compute_loss(val_pred, aro_pred, val_true, aro_true, loss_type)
                train_loss += loss.item()
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                scheduler.step()

        train_loss /= max(1, len(train_loader))

        # ── Validate (identical structure to train.py) ───────────────────
        model.eval()
        valid_loss_sum = 0.0
        valid_val_true = []
        valid_val_pred = []
        valid_aro_true = []
        valid_aro_pred = []

        with torch.no_grad():
            for images, _, val_true, aro_true in valid_loader:
                images, val_true, aro_true = (
                    images.to(DEVICE), val_true.to(DEVICE), aro_true.to(DEVICE)
                )
                with torch.autocast(device_type="cuda", dtype=torch.float16):
                    outputs = model(images)
                    val_pred = outputs[:, 0]
                    aro_pred = outputs[:, 1]
                    loss = compute_loss(
                        val_pred, aro_pred, val_true, aro_true, loss_type
                    )
                    valid_loss_sum += loss.item()

                valid_val_true.extend(val_true.detach().cpu().float().numpy())
                valid_val_pred.extend(val_pred.detach().cpu().float().numpy())
                valid_aro_true.extend(aro_true.detach().cpu().float().numpy())
                valid_aro_pred.extend(aro_pred.detach().cpu().float().numpy())

        valid_loss_avg = valid_loss_sum / max(1, len(valid_loader))
        valid_ccc_v = ccc_numpy_1d(valid_val_pred, valid_val_true)
        valid_ccc_a = ccc_numpy_1d(valid_aro_pred, valid_aro_true)
        valid_ccc_mean = (valid_ccc_v + valid_ccc_a) / 2.0

        history["epoch"].append(epoch + 1)
        history["train_loss"].append(train_loss)
        history["valid_loss"].append(valid_loss_avg)
        history["ccc_v"].append(valid_ccc_v)
        history["ccc_a"].append(valid_ccc_a)
        history["ccc_mean"].append(valid_ccc_mean)
        history["lr"].append(current_lr)

        print(
            f"[{name}] Epoch [{epoch+1}/{NUM_EPOCHS}] - "
            f"Training Loss: {train_loss:.4f}, "
            f"Validation Loss: {valid_loss_avg:.4f}, "
            f"CCC_v: {valid_ccc_v:.4f}, "
            f"CCC_a: {valid_ccc_a:.4f}, "
            f"CCC_mean: {valid_ccc_mean:.4f}, "
            f"Learning Rate: {current_lr:.8f}"
        )

        if valid_loss_avg < best_valid_loss:
            best_valid_loss = valid_loss_avg
            best_epoch = epoch + 1
            print(f"  -> Saving best model at epoch {best_epoch}")
            torch.save(model.state_dict(), exp_model_path)

    # ── Save CSV + plots ─────────────────────────────────────────────────
    history_df = pd.DataFrame(history)
    history_df.to_csv(exp_dir / "train_history.csv", index=False)
    save_experiment_plots(history_df, name, exp_dir)

    # Print best epoch metrics
    best_idx = best_epoch - 1
    print(f"\n=== [{name}] Best Epoch Metrics ===")
    print(f"Best epoch:     {best_epoch}")
    print(f"Train Loss:     {history_df['train_loss'][best_idx]:.4f}")
    print(f"Valid Loss:     {history_df['valid_loss'][best_idx]:.4f}")
    print(f"Valid CCC_v:    {history_df['ccc_v'][best_idx]:.4f}")
    print(f"Valid CCC_a:    {history_df['ccc_a'][best_idx]:.4f}")
    print(f"Valid CCC_mean: {history_df['ccc_mean'][best_idx]:.4f}")

    return {
        "name": name,
        "desc": desc,
        "model_type": exp_config["model_type"],
        "loss_type": exp_config["loss_type"],
        "best_epoch": best_epoch,
        "best_valid_loss": float(history_df["valid_loss"][best_idx]),
        "best_ccc_v": float(history_df["ccc_v"][best_idx]),
        "best_ccc_a": float(history_df["ccc_a"][best_idx]),
        "best_ccc_mean": float(history_df["ccc_mean"][best_idx]),
    }


# ============================================================================
# Main
# ============================================================================
print("\n--- Creating datasets (shared across all experiments) ---")

train_dataset = CustomDataset(
    dataframe=train_annotations_df,
    root_dir=IMAGE_FOLDER,
    transform=transform,
    balance=True,
)
valid_dataset = CustomDataset(
    dataframe=valid_annotations_df,
    root_dir=IMAGE_FOLDER_TEST,
    transform=transform_valid,
    balance=False,
)
train_loader = DataLoader(
    train_dataset, batch_size=BATCHSIZE, shuffle=True, num_workers=NUM_WORKERS
)
valid_loader = DataLoader(
    valid_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
)

print(f"Train dataset size (balanced): {len(train_dataset)}")
print(f"Valid dataset size:            {len(valid_dataset)}")
print(f"\n--- Running {len(EXPERIMENTS)} experiments ---")

all_results = []

for exp_config in EXPERIMENTS:
    result = run_experiment(exp_config, train_loader, valid_loader)
    all_results.append(result)

# ── Summary ──────────────────────────────────────────────────────────────
summary_df = pd.DataFrame(all_results)
summary_df.to_csv(EXPERIMENTS_DIR / "all_experiments_summary.csv", index=False)
save_summary_plots(summary_df, EXPERIMENTS_DIR)

print("\n" + "=" * 70)
print("ALL 6 EXPERIMENTS COMPLETE")
print("=" * 70)
print(summary_df.to_string(index=False))
print(f"\nSummary CSV: {EXPERIMENTS_DIR / 'all_experiments_summary.csv'}")
print(f"All output:  {EXPERIMENTS_DIR}")