CAGE_expression_inference-apvit/models/AffectNet8_Maxvit_VA/train2_old.py

"""
train2.py – Ablation study: 6 experiments (3 models × 2 losses)

Models:
  1) MaxViT-T + AttentivePoolingHead + load Combined weights   (attn_combined)
  2) MaxViT-T + AttentivePoolingHead + NO Combined weights     (attn_nocombined)
  3) MaxViT-T + Standard head        + NO Combined weights     (std_nocombined)

Losses:
  A) 3*MSE_val + 3*MSE_aro + CCC_val + CCC_aro                (mse_ccc)
  B) CCC_val + CCC_aro                                         (ccc_only)

Same pipeline as train.py for each experiment. Each saves its own
folder with per-epoch CSV, plots, and model checkpoint.
"""

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
# =========================
SCRIPT_DIR = Path(__file__).resolve().parent
MODEL_PATH = os.getenv("MODEL_PATH", str(SCRIPT_DIR / "model.pt"))
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
# =========================
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
# =========================
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")

COMBINED_MODEL_PATH = os.getenv(
    "COMBINED_MODEL_PATH",
    str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()),
)

print("DATA_ROOT          =", DATA_ROOT)
print("EXTRACT_ROOT       =", EXTRACT_ROOT)
print("AFFECTNET_TRAIN_IMAGES =", IMAGE_FOLDER)
print("AFFECTNET_VAL_IMAGES   =", IMAGE_FOLDER_TEST)
print("MODEL_PATH         =", MODEL_PATH)
print("COMBINED_MODEL_PATH=", COMBINED_MODEL_PATH)
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)


# =========================
# Experiment configurations
# =========================
EXPERIMENTS = [
    {
        "name": "exp1_attn_combined_mse_ccc",
        "desc": "MaxViT-T + Attention + Combined weights + MSE+CCC",
        "use_attention": True,
        "load_combined": True,
        "loss_type": "mse_ccc",
    },
    {
        "name": "exp2_attn_nocombined_mse_ccc",
        "desc": "MaxViT-T + Attention + ImageNet only + MSE+CCC",
        "use_attention": True,
        "load_combined": False,
        "loss_type": "mse_ccc",
    },
    {
        "name": "exp3_standard_nocombined_mse_ccc",
        "desc": "MaxViT-T + Standard head + ImageNet only + MSE+CCC",
        "use_attention": False,
        "load_combined": False,
        "loss_type": "mse_ccc",
    },
    {
        "name": "exp4_attn_combined_ccc_only",
        "desc": "MaxViT-T + Attention + Combined weights + CCC only",
        "use_attention": True,
        "load_combined": True,
        "loss_type": "ccc_only",
    },
    {
        "name": "exp5_attn_nocombined_ccc_only",
        "desc": "MaxViT-T + Attention + ImageNet only + CCC only",
        "use_attention": True,
        "load_combined": False,
        "loss_type": "ccc_only",
    },
    {
        "name": "exp6_standard_nocombined_ccc_only",
        "desc": "MaxViT-T + Standard head + ImageNet only + CCC only",
        "use_attention": False,
        "load_combined": False,
        "loss_type": "ccc_only",
    },
]


# =========================
# Plot helpers
# =========================
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()


# =========================
# Dataset & transforms
# =========================
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 components
# =========================
class AttentivePoolingHead(nn.Module):
    def __init__(self, feat_dim, hidden_dim, dropout=0.3):
        super().__init__()
        self.attn_proj = nn.Linear(feat_dim, 1)
        self.mlp = nn.Sequential(
            nn.LayerNorm(feat_dim),
            nn.Dropout(dropout),
            nn.Linear(feat_dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, 2, bias=False),
        )

    def forward(self, x):
        tokens = x.flatten(2).transpose(1, 2)  # [B, HW, C]
        weights = torch.softmax(self.attn_proj(tokens), dim=1)  # [B, HW, 1]
        pooled = (tokens * weights).sum(dim=1)  # [B, C]
        return torch.tanh(self.mlp(pooled))


def build_model(use_attention, load_combined):
    model = models.maxvit_t(weights="DEFAULT")
    block_channels = model.classifier[3].in_features

    if load_combined:
        # Temporarily set classifier to match Combined model architecture
        model.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.LayerNorm(block_channels),
            nn.Linear(block_channels, block_channels),
            nn.Tanh(),
            nn.Linear(block_channels, 10, bias=False),
        )
        model.to(DEVICE)
        model.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE))

    # Now set the final head
    if use_attention:
        model.classifier = AttentivePoolingHead(
            feat_dim=block_channels,
            hidden_dim=block_channels,
            dropout=0.3,
        )
    else:
        model.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.LayerNorm(block_channels),
            nn.Linear(block_channels, block_channels),
            nn.Tanh(),
            nn.Dropout(0.3),
            nn.Linear(block_channels, 2, bias=False),
        )

    model.to(DEVICE)
    return model


# =========================
# Loss & metrics
# =========================
mse_criterion_val = nn.MSELoss()
mse_criterion_aro = nn.MSELoss()


def CCCLoss(x, y):
    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):
    if loss_type == "mse_ccc":
        return (
            3 * mse_criterion_val(val_pred, val_true)
            + 3 * mse_criterion_aro(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):
    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
# =========================
def save_experiment_plots(history_df, exp_name, exp_dir):
    epochs = history_df["epoch"].tolist()

    save_line_plot(
        epochs,
        [history_df["train_loss"].tolist()],
        ["Train Loss"],
        f"{exp_name} - Train Loss",
        "Epoch", "Loss",
        exp_dir / "train_loss.png",
    )
    save_line_plot(
        epochs,
        [history_df["valid_loss"].tolist()],
        ["Validation Loss"],
        f"{exp_name} - Validation Loss",
        "Epoch", "Loss",
        exp_dir / "valid_loss.png",
    )
    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",
    )
    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",
    )
    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",
    )
    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",
    )
    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",
    )
    save_line_plot(
        epochs,
        [history_df["lr"].tolist()],
        ["Learning Rate"],
        f"{exp_name} - Learning Rate",
        "Epoch", "LR",
        exp_dir / "lr.png",
    )


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

    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",
    )
    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",
    )
    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",
    )
    save_bar_plot(
        names,
        summary_df["best_valid_loss"].tolist(),
        "Best Validation Loss by Experiment",
        "Experiment", "Loss",
        save_dir / "summary_best_valid_loss.png",
    )
    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():
        exp_dir = save_dir / row["name"]
        hist_csv = exp_dir / "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 valid loss curves from all experiments
    plt.figure(figsize=(12, 6))
    for _, row in summary_df.iterrows():
        exp_dir = save_dir / row["name"]
        hist_csv = exp_dir / "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()


# =========================
# Training function
# =========================
def run_experiment(exp_config, train_loader, valid_loader):
    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(SCRIPT_DIR / f"model_{name}.pt")

    print(f"\n{'='*60}")
    print(f"EXPERIMENT: {name}")
    print(f"  {desc}")
    print(f"  use_attention={exp_config['use_attention']}, "
          f"load_combined={exp_config['load_combined']}, "
          f"loss_type={loss_type}")
    print(f"  model_path={exp_model_path}")
    print(f"  plots_dir={exp_dir}")
    print(f"{'='*60}")

    model = build_model(exp_config["use_attention"], exp_config["load_combined"])
    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

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

    for epoch in range(NUM_EPOCHS):
        # ── Train ──────────────────────────────────────────────
        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} train"
        ):
            images = images.to(DEVICE)
            val_true = val_true.to(DEVICE)
            aro_true = aro_true.to(DEVICE)

            optimizer.zero_grad()
            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)

            train_loss += loss.item()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            scheduler.step()

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

        # ── Validate ───────────────────────────────────────────
        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 = images.to(DEVICE)
                val_true = val_true.to(DEVICE)
                aro_true = 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 = 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)
        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"Train Loss: {train_loss:.4f}, "
            f"Valid Loss: {valid_loss:.4f}, "
            f"CCC_v: {valid_ccc_v:.4f}, "
            f"CCC_a: {valid_ccc_a:.4f}, "
            f"CCC_mean: {valid_ccc_mean:.4f}, "
            f"LR: {current_lr:.8f}"
        )

        if valid_loss < best_valid_loss:
            best_valid_loss = valid_loss
            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 info
    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,
        "use_attention": exp_config["use_attention"],
        "load_combined": exp_config["load_combined"],
        "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)}")

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" + "=" * 60)
print("ALL EXPERIMENTS COMPLETE")
print("=" * 60)
print(summary_df.to_string(index=False))
print(f"\nSummary CSV: {EXPERIMENTS_DIR / 'all_experiments_summary.csv'}")
print(f"All plots:   {EXPERIMENTS_DIR}")