CAGE_expression_inference-apvit/models/AffectNet8_Maxvit_VA/train3.py

"""
train3.py – CAGE pipeline: 5-Fold StratifiedKFold (1 seed per fold)

Pipeline:
  - train_set is split into 80% train / 20% val via StratifiedKFold
  - Official val_set is used as test set (evaluated after each fold)
  - Model: MaxViT-T + Standard head + Combined weights
  - Loss:  3*MSE_val + 3*MSE_aro + CCC_val + CCC_aro
"""

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
from sklearn.model_selection import StratifiedKFold

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


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


# Load the annotations for training and validation from separate CSV files
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)


# Set 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")

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("PLOTS_DIR =", PLOTS_DIR)
print("Train set size =", len(train_annotations_df))
print("Official val_set size (final test only) =", len(valid_annotations_df))


# =========================
# 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=(10, 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=9,
            )
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_boxplot(data_list, labels, title, ylabel, save_path):
    plt.figure(figsize=(10, 6))
    plt.boxplot(data_list, labels=labels, vert=True)
    plt.title(title)
    plt.ylabel(ylabel)
    plt.grid(True, axis="y", alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path, dpi=200)
    plt.close()


def save_histogram(values, bins, title, xlabel, ylabel, save_path):
    plt.figure(figsize=(10, 6))
    plt.hist(values, bins=bins)
    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()


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()


def save_fold_plots(history_df, fold_idx, seed, fold_dir):
    epochs = history_df["epoch"].tolist()

    save_line_plot(
        epochs,
        [history_df["train_loss"].tolist()],
        ["Train Loss"],
        f"CAGE Fold {fold_idx} Seed {seed} - Train Loss",
        "Epoch",
        "Loss",
        fold_dir / f"fold_{fold_idx}_train_loss.png",
    )

    save_line_plot(
        epochs,
        [history_df["valid_loss"].tolist()],
        ["Validation Loss"],
        f"CAGE Fold {fold_idx} Seed {seed} - Validation Loss",
        "Epoch",
        "Loss",
        fold_dir / f"fold_{fold_idx}_valid_loss.png",
    )

    save_line_plot(
        epochs,
        [history_df["train_loss"].tolist(), history_df["valid_loss"].tolist()],
        ["Train Loss", "Validation Loss"],
        f"CAGE Fold {fold_idx} Seed {seed} - Train vs Validation Loss",
        "Epoch",
        "Loss",
        fold_dir / f"fold_{fold_idx}_train_valid_loss.png",
    )

    save_line_plot(
        epochs,
        [history_df["ccc_v"].tolist()],
        ["CCC_v"],
        f"CAGE Fold {fold_idx} Seed {seed} - CCC_v",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_v.png",
    )

    save_line_plot(
        epochs,
        [history_df["ccc_a"].tolist()],
        ["CCC_a"],
        f"CAGE Fold {fold_idx} Seed {seed} - CCC_a",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_a.png",
    )

    save_line_plot(
        epochs,
        [history_df["ccc_mean"].tolist()],
        ["CCC_mean"],
        f"CAGE Fold {fold_idx} Seed {seed} - CCC_mean",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_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"CAGE Fold {fold_idx} Seed {seed} - Combined CCC Curves",
        "Epoch",
        "CCC",
        fold_dir / f"fold_{fold_idx}_ccc_all.png",
    )

    save_line_plot(
        epochs,
        [history_df["lr"].tolist()],
        ["Learning Rate"],
        f"CAGE Fold {fold_idx} Seed {seed} - Learning Rate",
        "Epoch",
        "LR",
        fold_dir / f"fold_{fold_idx}_lr.png",
    )

    save_scatter(
        history_df["train_loss"].tolist(),
        history_df["valid_loss"].tolist(),
        f"CAGE Fold {fold_idx} Seed {seed} - Train Loss vs Valid Loss",
        "Train Loss",
        "Validation Loss",
        fold_dir / f"fold_{fold_idx}_train_vs_valid_scatter.png",
        annotate_labels=epochs,
    )

    save_scatter(
        history_df["ccc_v"].tolist(),
        history_df["ccc_a"].tolist(),
        f"CAGE Fold {fold_idx} Seed {seed} - CCC_v vs CCC_a",
        "CCC_v",
        "CCC_a",
        fold_dir / f"fold_{fold_idx}_ccc_v_vs_ccc_a.png",
        annotate_labels=epochs,
    )


def save_summary_plots(summary_df, save_dir):
    fold_labels = [f"Fold {x}" for x in summary_df["fold"].tolist()]

    save_bar_plot(
        fold_labels,
        summary_df["test_loss"].tolist(),
        "CAGE Final Test Loss by Fold",
        "Fold",
        "Loss",
        save_dir / "summary_test_loss_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["test_ccc_v"].tolist(),
        "CAGE Final Test CCC_v by Fold",
        "Fold",
        "CCC_v",
        save_dir / "summary_test_ccc_v_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["test_ccc_a"].tolist(),
        "CAGE Final Test CCC_a by Fold",
        "Fold",
        "CCC_a",
        save_dir / "summary_test_ccc_a_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["test_ccc_mean"].tolist(),
        "CAGE Final Test CCC_mean by Fold",
        "Fold",
        "CCC_mean",
        save_dir / "summary_test_ccc_mean_by_fold.png",
    )

    save_bar_plot(
        fold_labels,
        summary_df["best_epoch"].tolist(),
        "CAGE Best Epoch by Fold",
        "Fold",
        "Best Epoch",
        save_dir / "summary_best_epoch_by_fold.png",
        annotate=True,
    )

    save_boxplot(
        [
            summary_df["test_ccc_v"].tolist(),
            summary_df["test_ccc_a"].tolist(),
            summary_df["test_ccc_mean"].tolist(),
        ],
        ["CCC_v", "CCC_a", "CCC_mean"],
        "CAGE Distribution of Final CCC Across 5 Folds",
        "CCC",
        save_dir / "summary_ccc_boxplot.png",
    )

    save_histogram(
        summary_df["test_ccc_mean"].tolist(),
        bins=min(5, len(summary_df)),
        title="CAGE Histogram of Final CCC_mean",
        xlabel="CCC_mean",
        ylabel="Frequency",
        save_path=save_dir / "summary_ccc_mean_histogram.png",
    )

    save_scatter(
        summary_df["best_epoch"].tolist(),
        summary_df["test_ccc_mean"].tolist(),
        "CAGE Best Epoch vs Final CCC_mean",
        "Best Epoch",
        "Final CCC_mean",
        save_dir / "summary_best_epoch_vs_ccc_mean.png",
        annotate_labels=summary_df["fold"].tolist(),
    )


# **** Create dataset and data loaders ****
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]),
    ]
)

test_dataset = CustomDataset(
    dataframe=valid_annotations_df,
    root_dir=IMAGE_FOLDER_TEST,
    transform=transform_valid,
    balance=False,
)
test_loader = DataLoader(
    test_dataset, batch_size=BATCHSIZE, shuffle=False, num_workers=NUM_WORKERS
)


# ***** Define the model *****
def build_model():
    MODEL = models.maxvit_t(weights="DEFAULT")
    block_channels = MODEL.classifier[3].in_features
    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)
    COMBINED_MODEL_PATH = os.getenv(
        "COMBINED_MODEL_PATH",
        str((SCRIPT_DIR / "../AffectNet8_Maxvit_Combined/model.pt").resolve()),
    )
    MODEL.load_state_dict(torch.load(COMBINED_MODEL_PATH, map_location=DEVICE))
    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


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 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)


def evaluate_loader(model, loader):
    model.eval()
    eval_loss = 0.0
    eval_val_true = []
    eval_val_pred = []
    eval_aro_true = []
    eval_aro_pred = []

    with torch.no_grad():
        for images, _, val_true, aro_true in 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 = (
                    3 * val_loss(val_pred.cuda(), val_true.cuda())
                    + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda())
                    + CCCLoss(val_pred.cuda(), val_true.cuda())
                    + CCCLoss(aro_pred.cuda(), aro_true.cuda())
                )
                eval_loss += loss.item()

            eval_val_true.extend(val_true.detach().cpu().float().numpy())
            eval_val_pred.extend(val_pred.detach().cpu().float().numpy())
            eval_aro_true.extend(aro_true.detach().cpu().float().numpy())
            eval_aro_pred.extend(aro_pred.detach().cpu().float().numpy())

    ccc_v = ccc_numpy_1d(eval_val_pred, eval_val_true)
    ccc_a = ccc_numpy_1d(eval_aro_pred, eval_aro_true)
    ccc_mean = (ccc_v + ccc_a) / 2.0
    return eval_loss / max(1, len(loader)), ccc_v, ccc_a, ccc_mean


val_loss = nn.MSELoss()
aro_loss = nn.MSELoss()

# ***** Train the model *****
print("--- Start CAGE Stratified 5-Fold training (1 fold = 1 seed) ---")
seeds = [11, 32, 57, 23, 42]
fold_test_metrics = []
global_best_valid_loss = float("inf")
global_best_model_path = MODEL_PATH
l2_lambda = 0.00001
l1_lambda = 0.00001

for fold_idx, seed in enumerate(seeds, start=1):
    print(f"--- Fold {fold_idx}/5 with seed {seed} ---")
    fold_dir = PLOTS_DIR / f"fold_{fold_idx}_seed_{seed}"
    fold_dir.mkdir(parents=True, exist_ok=True)

    skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=seed)
    split_list = list(skf.split(train_annotations_df, train_annotations_df["exp"]))
    train_idx, val_idx = split_list[fold_idx - 1]

    fold_train_df = train_annotations_df.iloc[train_idx].reset_index(drop=True)
    fold_valid_df = train_annotations_df.iloc[val_idx].reset_index(drop=True)
    print(f"Fold train size = {len(fold_train_df)}, fold val size = {len(fold_valid_df)}")

    fold_train_df.to_csv(fold_dir / f"fold_{fold_idx}_train_split.csv", index=False)
    fold_valid_df.to_csv(fold_dir / f"fold_{fold_idx}_valid_split.csv", index=False)

    train_dataset = CustomDataset(
        dataframe=fold_train_df,
        root_dir=IMAGE_FOLDER,
        transform=transform,
        balance=True,
    )
    valid_dataset = CustomDataset(
        dataframe=fold_valid_df,
        root_dir=IMAGE_FOLDER,
        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
    )

    MODEL = build_model()
    optimizer = optim.AdamW(MODEL.parameters(), lr=LR)
    cosine_scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=BATCHSIZE * NUM_EPOCHS)
    scaler = torch.cuda.amp.GradScaler()
    fold_best_valid_loss = float("inf")
    fold_best_epoch = -1
    fold_model_path = str(SCRIPT_DIR / f"model_s{seed}_fold{fold_idx}.pt")

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

    for epoch in range(NUM_EPOCHS):
        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"CAGE Fold {fold_idx} train progress"
        ):
            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 = (
                    3 * val_loss(val_pred.cuda(), val_true.cuda())
                    + 3 * aro_loss(aro_pred.cuda(), aro_true.cuda())
                    + CCCLoss(val_pred.cuda(), val_true.cuda())
                    + CCCLoss(aro_pred.cuda(), aro_true.cuda())
                )
                train_loss += loss.item()
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
                cosine_scheduler.step()

        train_loss_mean = train_loss / len(train_loader)
        valid_loss_mean, ccc_v, ccc_a, ccc_mean = evaluate_loader(MODEL, valid_loader)

        history["epoch"].append(epoch + 1)
        history["train_loss"].append(train_loss_mean)
        history["valid_loss"].append(valid_loss_mean)
        history["ccc_v"].append(ccc_v)
        history["ccc_a"].append(ccc_a)
        history["ccc_mean"].append(ccc_mean)
        history["lr"].append(current_lr)

        print(
            f"Fold [{fold_idx}/5] Epoch [{epoch+1}/{NUM_EPOCHS}] - "
            f"Training Loss: {train_loss_mean:.4f}, "
            f"Validation Loss (fold 20%): {valid_loss_mean:.4f}, "
            f"CCC_v: {ccc_v:.4f}, "
            f"CCC_a: {ccc_a:.4f}, "
            f"CCC_mean: {ccc_mean:.4f}, "
            f"Learning Rate: {current_lr:.8f}, "
        )

        if valid_loss_mean < fold_best_valid_loss:
            fold_best_valid_loss = valid_loss_mean
            fold_best_epoch = epoch + 1
            torch.save(MODEL.state_dict(), fold_model_path)
            print(f"Saved best fold model: {fold_model_path}")

        if fold_best_valid_loss < global_best_valid_loss:
            global_best_valid_loss = fold_best_valid_loss
            global_best_model_path = fold_model_path

    history_df = pd.DataFrame(history)
    history_df.to_csv(fold_dir / f"fold_{fold_idx}_epoch_history.csv", index=False)
    save_fold_plots(history_df, fold_idx, seed, fold_dir)

    MODEL.load_state_dict(torch.load(fold_model_path, map_location=DEVICE))
    test_loss_mean, test_ccc_v, test_ccc_a, test_ccc_mean = evaluate_loader(MODEL, test_loader)

    fold_test_metrics.append(
        (
            seed,
            fold_idx,
            fold_best_epoch,
            fold_best_valid_loss,
            test_loss_mean,
            test_ccc_v,
            test_ccc_a,
            test_ccc_mean,
        )
    )

    print(
        f"Fold [{fold_idx}/5] Final Test (official val_set) - "
        f"Loss: {test_loss_mean:.4f}, "
        f"CCC_v: {test_ccc_v:.4f}, "
        f"CCC_a: {test_ccc_a:.4f}, "
        f"CCC_mean: {test_ccc_mean:.4f}"
    )

summary_df = pd.DataFrame(
    fold_test_metrics,
    columns=[
        "seed",
        "fold",
        "best_epoch",
        "best_valid_loss",
        "test_loss",
        "test_ccc_v",
        "test_ccc_a",
        "test_ccc_mean",
    ],
)

summary_df.to_csv(PLOTS_DIR / "fold_summary.csv", index=False)
save_summary_plots(summary_df, PLOTS_DIR)

print("--- Final summary over runs on official val_set ---")
test_losses = summary_df["test_loss"].to_numpy()
test_ccc_vs = summary_df["test_ccc_v"].to_numpy()
test_ccc_as = summary_df["test_ccc_a"].to_numpy()
test_ccc_means = summary_df["test_ccc_mean"].to_numpy()

print(
    f"Runs: {len(fold_test_metrics)}; Mean Loss: {np.mean(test_losses):.4f} (std {np.std(test_losses):.4f}), "
    f"Mean CCC_v: {np.mean(test_ccc_vs):.4f} (std {np.std(test_ccc_vs):.4f}), "
    f"Mean CCC_a: {np.mean(test_ccc_as):.4f} (std {np.std(test_ccc_as):.4f}), "
    f"Mean CCC_mean: {np.mean(test_ccc_means):.4f} (std {np.std(test_ccc_means):.4f})"
)

torch.save(torch.load(global_best_model_path, map_location=DEVICE), MODEL_PATH)
print(f"Best fold model copied to MODEL_PATH: {MODEL_PATH}")
print(f"All plots and CSV logs saved to: {PLOTS_DIR}")