src/train_segmentation_v4.py

"""v4 SOTA-aimed trainer for BraTS whole-tumor segmentation.

Upgrades over v3 (2D SMP U-Net + ResNet34):
  - 3D not 2D: MONAI UNet with full volumetric patches. Inter-slice context
    is the single biggest 2D quality gap; this fixes it.
  - 128**3 patches: better than any 2D resolution.
  - Gradient accumulation: batch=2 micro x 8 = effective batch 16.
  - Deep supervision: multi-resolution loss heads (MONAI DynUNet supports
    deep supervision natively; we use it instead of plain UNet).
  - Larger capacity: DynUNet with filters=(32, 64, 128, 256, 320, 320).
  - Heavy 3D augmentation: TorchIO (flip, affine, elastic, bias field,
    intensity, noise, blur, ghosting, motion).
  - K-fold CV (--folds N): trains N models with patient-level splits, saves
    each fold's best, then computes ensemble + TTA prediction at test time.
  - 8-way TTA at inference: identity + 3 axis flips + 4 axis-pair flips
    averaged.
  - FP16 mixed precision, AdamW + cosine schedule with warmup, crash-resilient
    per-epoch checkpoints + --resume.

Output structure:

    segmentation_artifacts/brats3d_v4/
        fold_0/best_model.pt last.pt history.json
        fold_1/...
        ...
        evaluation_metrics.json   (ensemble + per-fold metrics)
        training_curves.png

Expected wall-time on RTX 4060 mobile (8 GB VRAM, FP16):
  - Single fold @ 50 epochs ~ 6-10 hours.
  - 5-fold ensemble: 30-50 hours total.

Use --folds 1 for a quick smoke test of the full stack before committing to
the multi-day 5-fold run.
"""

from __future__ import annotations

import argparse
import json
import os
import random
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
import torchio as tio
from monai.losses import DeepSupervisionLoss, DiceCELoss
from monai.networks.nets import DynUNet
from monai.transforms import (
    Compose,
    CropForegroundd,
    EnsureChannelFirstd,
    NormalizeIntensityd,
    RandSpatialCropd,
    SpatialPadd,
    ToTensord,
)
from sklearn.model_selection import KFold
from torch.utils.data import DataLoader, Dataset

_REPO_ROOT = Path(__file__).resolve().parents[1]
if str(_REPO_ROOT) not in sys.path:
    sys.path.append(str(_REPO_ROOT))


class BratsNPZDataset(Dataset):
    """Loads pre-baked .npz volumes from dataset_brats_3d/<split>/*.npz.

    Each .npz has:
      image: (C=4, D, H, W) float32 z-scored
      mask:  (D, H, W) uint8 binary
    """

    def __init__(self, npz_paths: list[Path], patch_size: int = 128, train: bool = True, max_tries: int = 16):
        self.paths = list(npz_paths)
        self.patch_size = patch_size
        self.train = train
        self.max_tries = max_tries

        if train:
            # TorchIO accepts torch tensors via tio.Subject; we build a tiny pipeline
            # of 3D augmentations applied on the cropped patch (not the whole volume,
            # to keep wall time bearable).
            self.aug = tio.Compose([
                tio.RandomFlip(axes=(0, 1, 2), p=0.5),
                tio.RandomAffine(scales=(0.9, 1.1), degrees=10, translation=5, p=0.5),
                tio.RandomElasticDeformation(num_control_points=5, max_displacement=5, p=0.2),
                tio.RandomBiasField(coefficients=0.3, order=3, p=0.3),
                tio.RandomNoise(std=(0.0, 0.05), p=0.3),
                tio.RandomBlur(std=(0.0, 1.0), p=0.2),
                tio.RandomGamma(log_gamma=(-0.2, 0.2), p=0.3),
            ])
        else:
            self.aug = None

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

    def _sample_patch(self, image: np.ndarray, mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """Sample a 128**3 patch. During training, bias toward tumor-containing
        patches (sample center near a tumor voxel) so the model sees positives."""
        C, D, H, W = image.shape
        p = self.patch_size
        pad = [(0, max(0, p - D)), (0, max(0, p - H)), (0, max(0, p - W))]
        if any(b[1] > 0 for b in pad):
            image = np.pad(image, [(0, 0)] + pad, mode='constant')
            mask = np.pad(mask, pad, mode='constant')
            D, H, W = image.shape[1:]

        if self.train and mask.sum() > 0:
            # Pick a random tumor voxel for biased center sampling 80% of the time.
            for _ in range(self.max_tries):
                if random.random() < 0.8:
                    zs, ys, xs = np.where(mask > 0)
                    idx = random.randint(0, len(zs) - 1)
                    cz, cy, cx = int(zs[idx]), int(ys[idx]), int(xs[idx])
                else:
                    cz = random.randint(p // 2, D - p // 2)
                    cy = random.randint(p // 2, H - p // 2)
                    cx = random.randint(p // 2, W - p // 2)
                z0 = max(0, min(D - p, cz - p // 2))
                y0 = max(0, min(H - p, cy - p // 2))
                x0 = max(0, min(W - p, cx - p // 2))
                m_patch = mask[z0:z0 + p, y0:y0 + p, x0:x0 + p]
                if m_patch.sum() > 0 or random.random() < 0.2:
                    return image[:, z0:z0 + p, y0:y0 + p, x0:x0 + p], m_patch
        # Eval / fallback: center crop
        z0 = max(0, (D - p) // 2)
        y0 = max(0, (H - p) // 2)
        x0 = max(0, (W - p) // 2)
        return image[:, z0:z0 + p, y0:y0 + p, x0:x0 + p], mask[z0:z0 + p, y0:y0 + p, x0:x0 + p]

    def __getitem__(self, idx):
        data = np.load(str(self.paths[idx]))
        image = data['image'].astype(np.float32)
        mask = data['mask'].astype(np.float32)
        image_patch, mask_patch = self._sample_patch(image, mask)
        image_t = torch.from_numpy(image_patch)
        mask_t = torch.from_numpy(mask_patch).unsqueeze(0)  # (1, D, H, W)

        if self.aug is not None:
            subject = tio.Subject(
                image=tio.ScalarImage(tensor=image_t),
                mask=tio.LabelMap(tensor=mask_t),
            )
            subject = self.aug(subject)
            image_t = subject['image'].tensor
            mask_t = subject['mask'].tensor.float()
            mask_t = (mask_t > 0.5).float()

        return image_t, mask_t


def dice_score(probs: torch.Tensor, targets: torch.Tensor, smooth: float = 1e-6) -> torch.Tensor:
    p = probs.contiguous().view(probs.size(0), -1)
    t = targets.contiguous().view(targets.size(0), -1)
    inter = (p * t).sum(dim=1)
    return ((2.0 * inter + smooth) / (p.sum(dim=1) + t.sum(dim=1) + smooth)).mean()


def iou_score(probs: torch.Tensor, targets: torch.Tensor, smooth: float = 1e-6) -> torch.Tensor:
    p = probs.contiguous().view(probs.size(0), -1)
    t = targets.contiguous().view(targets.size(0), -1)
    inter = (p * t).sum(dim=1)
    union = p.sum(dim=1) + t.sum(dim=1) - inter
    return ((inter + smooth) / (union + smooth)).mean()


@torch.no_grad()
def sliding_window_inference(model, image: torch.Tensor, patch_size: int = 128,
                              overlap: float = 0.5, amp_enabled: bool = True) -> torch.Tensor:
    """Whole-volume inference by sliding 128**3 windows with `overlap` overlap.
    Returns sigmoid probabilities of shape (1, 1, D, H, W). MONAI has a
    sliding_window_inference function too; we hand-roll a simple one so we
    don't pull in extra deps at inference time."""
    from monai.inferers import sliding_window_inference as _sw
    return torch.sigmoid(_sw(
        inputs=image,
        roi_size=(patch_size, patch_size, patch_size),
        sw_batch_size=1,
        predictor=model,
        overlap=overlap,
        mode='gaussian',
    ))


@torch.no_grad()
def tta_predict(model, image: torch.Tensor, patch_size: int, amp_enabled: bool) -> torch.Tensor:
    """8-way TTA: identity + 3 axis flips + 4 axis-pair flips averaged."""
    flip_axis_sets = [
        (),
        (2,), (3,), (4,),
        (2, 3), (2, 4), (3, 4),
        (2, 3, 4),
    ]
    accum = None
    for ax in flip_axis_sets:
        x = torch.flip(image, dims=ax) if ax else image
        with torch.amp.autocast('cuda', enabled=amp_enabled):
            probs = sliding_window_inference(model, x, patch_size=patch_size, overlap=0.5, amp_enabled=amp_enabled)
        if ax:
            probs = torch.flip(probs, dims=ax)
        accum = probs if accum is None else accum + probs
    return accum / len(flip_axis_sets)


@torch.no_grad()
def evaluate_volumes(model, loader_npz_paths: list[Path], device, patch_size: int = 128,
                      threshold: float = 0.5, tta: bool = False, amp_enabled: bool = True) -> dict:
    model.eval()
    dice_sum = iou_sum = 0.0
    inter = pos_true = pos_pred = 0
    n = 0
    for p in loader_npz_paths:
        data = np.load(str(p))
        image = torch.from_numpy(data['image'].astype(np.float32)).unsqueeze(0).to(device)
        mask = torch.from_numpy(data['mask'].astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)
        if tta:
            probs = tta_predict(model, image, patch_size, amp_enabled)
        else:
            probs = sliding_window_inference(model, image, patch_size=patch_size, overlap=0.5, amp_enabled=amp_enabled)
        binp = (probs >= threshold).float()
        dice_sum += float(dice_score(binp, mask))
        iou_sum += float(iou_score(binp, mask))
        inter += int((binp * mask).sum().item())
        pos_true += int(mask.sum().item())
        pos_pred += int(binp.sum().item())
        n += 1
    if n == 0:
        return {}
    return {
        'n_volumes': n,
        'dice': dice_sum / n,
        'iou': iou_sum / n,
        'micro_dice': (2 * inter) / max(pos_true + pos_pred, 1),
        'micro_iou': inter / max(pos_true + pos_pred - inter, 1),
        'tta': tta,
    }


def build_model(deep_supervision: bool) -> torch.nn.Module:
    """DynUNet from MONAI: configurable nnU-Net-style architecture with deep
    supervision baked in."""
    kernel_size = [3, 3, 3, 3, 3, 3]
    strides = [1, 2, 2, 2, 2, [2, 2, 1]]  # last stride keeps tiny depth dim alive
    upsample_kernel_size = strides[1:]
    model = DynUNet(
        spatial_dims=3,
        in_channels=4,
        out_channels=1,
        kernel_size=kernel_size,
        strides=strides,
        upsample_kernel_size=upsample_kernel_size,
        filters=(32, 64, 128, 256, 320, 320),
        norm_name='instance',
        deep_supervision=deep_supervision,
        deep_supr_num=2,
        res_block=True,
    )
    return model


def train_one_fold(fold_idx: int, train_paths: list[Path], val_paths: list[Path],
                    test_paths: list[Path], args, fold_out: Path) -> dict:
    fold_out.mkdir(parents=True, exist_ok=True)
    torch.manual_seed(args.seed + fold_idx)
    np.random.seed(args.seed + fold_idx)
    random.seed(args.seed + fold_idx)

    if args.device == 'cuda' and not torch.cuda.is_available():
        args.device = 'cpu'
    device = torch.device(args.device)
    amp_enabled = (device.type == 'cuda') and not args.no_amp

    print(f'\n========== FOLD {fold_idx} ==========', flush=True)
    print(f'[fold {fold_idx}] device={device} amp={amp_enabled}'
          + (f' ({torch.cuda.get_device_name(0)})' if device.type == 'cuda' else ''), flush=True)
    print(f'[fold {fold_idx}] train={len(train_paths)} val={len(val_paths)} test={len(test_paths)}', flush=True)

    train_ds = BratsNPZDataset(train_paths, patch_size=args.patch_size, train=True)
    val_ds = BratsNPZDataset(val_paths, patch_size=args.patch_size, train=False)
    common = dict(batch_size=args.batch_size, num_workers=args.num_workers, pin_memory=(device.type == 'cuda'))
    train_loader = DataLoader(train_ds, shuffle=True, drop_last=True, **common)
    # For val during epochs we use patches (faster); for final test we do
    # whole-volume sliding window inference.
    val_loader = DataLoader(val_ds, shuffle=False, **common)

    model = build_model(deep_supervision=True).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f'[fold {fold_idx}] model: DynUNet 3D - {n_params:,} params', flush=True)

    # MONAI DiceCE handles binary segmentation with optional sigmoid; wrap in
    # DeepSupervisionLoss so each side output also contributes.
    base_loss = DiceCELoss(sigmoid=True, smooth_nr=1e-5, smooth_dr=1e-5, lambda_dice=0.6, lambda_ce=0.4)
    loss_fn = DeepSupervisionLoss(loss=base_loss, weights=None)  # None -> 1/(2^i) defaults

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max(1, args.epochs - args.warmup_epochs), eta_min=1e-6)
    scaler = torch.amp.GradScaler('cuda', enabled=amp_enabled)

    best_path = fold_out / 'best_model.pt'
    last_path = fold_out / 'last.pt'
    history_path = fold_out / 'history.json'
    history = {'train_loss': [], 'val_dice': [], 'val_iou': [], 'lr': []}
    best_val_dice = -1.0
    epochs_without_improve = 0
    start_epoch = 0

    if args.resume and last_path.exists():
        ckpt = torch.load(str(last_path), map_location=device, weights_only=False)
        model.load_state_dict(ckpt['state_dict'])
        if 'optimizer_state' in ckpt:
            optimizer.load_state_dict(ckpt['optimizer_state'])
        if 'scheduler_state' in ckpt:
            try:
                scheduler.load_state_dict(ckpt['scheduler_state'])
            except Exception:
                pass
        if amp_enabled and 'scaler_state' in ckpt:
            try:
                scaler.load_state_dict(ckpt['scaler_state'])
            except Exception:
                pass
        history = ckpt.get('history', history)
        best_val_dice = float(ckpt.get('best_val_dice', best_val_dice))
        epochs_without_improve = int(ckpt.get('epochs_without_improve', 0))
        start_epoch = int(ckpt.get('epoch', 0))
        print(f'[fold {fold_idx}] resumed at epoch {start_epoch} (best_val_dice={best_val_dice:.4f})', flush=True)

    for epoch in range(start_epoch, args.epochs):
        if epoch < args.warmup_epochs:
            for pg in optimizer.param_groups:
                pg['lr'] = args.learning_rate * (epoch + 1) / max(1, args.warmup_epochs)
        model.train()
        t0 = time.time()
        running_loss = 0.0
        n_steps = 0
        optimizer.zero_grad(set_to_none=True)
        for step, (x, y) in enumerate(train_loader):
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            with torch.amp.autocast('cuda', enabled=amp_enabled):
                outputs = model(x)
                # When deep_supervision=True, DynUNet returns a stacked tensor
                # of shape (heads, B, C, D, H, W). DeepSupervisionLoss handles
                # that layout natively.
                loss = loss_fn(outputs, y) / args.grad_accum_steps
            if amp_enabled:
                scaler.scale(loss).backward()
            else:
                loss.backward()
            if (step + 1) % args.grad_accum_steps == 0:
                if amp_enabled:
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    optimizer.step()
                optimizer.zero_grad(set_to_none=True)
            running_loss += float(loss) * args.grad_accum_steps
            n_steps += 1
        if epoch >= args.warmup_epochs:
            scheduler.step()

        # Validation: a few sliding-window forwards (no TTA each epoch - too slow)
        model.eval()
        d_sum = iou_sum = 0.0
        nv = 0
        with torch.no_grad():
            for vp in val_paths[: min(len(val_paths), args.val_subset)]:
                data = np.load(str(vp))
                vx = torch.from_numpy(data['image'].astype(np.float32)).unsqueeze(0).to(device)
                vy = torch.from_numpy(data['mask'].astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)
                with torch.amp.autocast('cuda', enabled=amp_enabled):
                    probs = sliding_window_inference(model, vx, patch_size=args.patch_size,
                                                       overlap=0.25, amp_enabled=amp_enabled)
                binp = (probs >= 0.5).float()
                d_sum += float(dice_score(binp, vy))
                iou_sum += float(iou_score(binp, vy))
                nv += 1
        val_dice = d_sum / max(nv, 1)
        val_iou = iou_sum / max(nv, 1)
        train_loss = running_loss / max(n_steps, 1)
        elapsed = time.time() - t0
        lr_now = optimizer.param_groups[0]['lr']
        history['train_loss'].append(train_loss)
        history['val_dice'].append(val_dice)
        history['val_iou'].append(val_iou)
        history['lr'].append(lr_now)
        print(f'[fold {fold_idx}][ep {epoch+1:02d}/{args.epochs}] '
              f'train_loss={train_loss:.4f}  val_dice@{nv}={val_dice:.4f}  val_iou={val_iou:.4f}  '
              f'lr={lr_now:.2e}  ({elapsed:.1f}s)', flush=True)

        if val_dice > best_val_dice:
            best_val_dice = val_dice
            epochs_without_improve = 0
            torch.save({
                'state_dict': model.state_dict(),
                'config': vars(args),
                'val_metrics': {'dice': val_dice, 'iou': val_iou, 'n_val_volumes': nv},
                'epoch': epoch + 1,
                'fold_idx': fold_idx,
            }, best_path)
            print(f'        -> new best val_dice={best_val_dice:.4f}; saved {best_path}', flush=True)
        else:
            epochs_without_improve += 1

        torch.save({
            'state_dict': model.state_dict(),
            'optimizer_state': optimizer.state_dict(),
            'scheduler_state': scheduler.state_dict(),
            'scaler_state': scaler.state_dict() if amp_enabled else None,
            'config': vars(args),
            'val_metrics': {'dice': val_dice, 'iou': val_iou},
            'epoch': epoch + 1,
            'history': history,
            'best_val_dice': best_val_dice,
            'epochs_without_improve': epochs_without_improve,
            'fold_idx': fold_idx,
        }, last_path)
        history_path.write_text(json.dumps(history, indent=2), encoding='utf-8')

        if epochs_without_improve >= args.patience:
            print(f'[fold {fold_idx}] Early stopping: no improvement in {args.patience} epochs.', flush=True)
            break

    if best_path.exists():
        ckpt = torch.load(str(best_path), map_location=device, weights_only=False)
        model.load_state_dict(ckpt['state_dict'])
    val_eval = evaluate_volumes(model, val_paths, device, patch_size=args.patch_size,
                                  tta=args.tta_eval, amp_enabled=amp_enabled)
    test_eval = evaluate_volumes(model, test_paths, device, patch_size=args.patch_size,
                                   tta=args.tta_eval, amp_enabled=amp_enabled)
    fold_metrics = {'fold': fold_idx, 'best_val_dice': best_val_dice, 'val': val_eval, 'test': test_eval}
    (fold_out / 'fold_metrics.json').write_text(json.dumps(fold_metrics, indent=2), encoding='utf-8')
    print(f'[fold {fold_idx}] FINAL: val={val_eval}\n    test={test_eval}', flush=True)
    return fold_metrics


@torch.no_grad()
def ensemble_evaluate(fold_dirs: list[Path], test_paths: list[Path], args) -> dict:
    """Average sigmoid predictions across all folds (with optional TTA per
    fold), then threshold."""
    if args.device == 'cuda' and not torch.cuda.is_available():
        args.device = 'cpu'
    device = torch.device(args.device)
    amp_enabled = (device.type == 'cuda') and not args.no_amp

    models = []
    for d in fold_dirs:
        best = d / 'best_model.pt'
        if not best.exists():
            continue
        ckpt = torch.load(str(best), map_location=device, weights_only=False)
        m = build_model(deep_supervision=False).to(device)
        m.load_state_dict(ckpt['state_dict'], strict=False)
        m.eval()
        models.append(m)
    if not models:
        return {}

    dice_sum = iou_sum = 0.0
    inter = pos_true = pos_pred = 0
    n = 0
    for p in test_paths:
        data = np.load(str(p))
        image = torch.from_numpy(data['image'].astype(np.float32)).unsqueeze(0).to(device)
        mask = torch.from_numpy(data['mask'].astype(np.float32)).unsqueeze(0).unsqueeze(0).to(device)
        ensemble_probs = None
        for m in models:
            probs = tta_predict(m, image, args.patch_size, amp_enabled) if args.tta_eval \
                else sliding_window_inference(m, image, patch_size=args.patch_size, overlap=0.5, amp_enabled=amp_enabled)
            ensemble_probs = probs if ensemble_probs is None else ensemble_probs + probs
        ensemble_probs /= len(models)
        binp = (ensemble_probs >= 0.5).float()
        dice_sum += float(dice_score(binp, mask))
        iou_sum += float(iou_score(binp, mask))
        inter += int((binp * mask).sum().item())
        pos_true += int(mask.sum().item())
        pos_pred += int(binp.sum().item())
        n += 1
    return {
        'n_models': len(models),
        'n_test_volumes': n,
        'dice': dice_sum / max(n, 1),
        'iou': iou_sum / max(n, 1),
        'micro_dice': (2 * inter) / max(pos_true + pos_pred, 1),
        'tta': args.tta_eval,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', default='dataset_brats_3d',
                        help='Output of prepare_brats_3d_dataset.py')
    parser.add_argument('--output_dir', default='segmentation_artifacts/brats3d_v4')
    parser.add_argument('--patch_size', type=int, default=128)
    parser.add_argument('--batch_size', type=int, default=2,
                        help='Micro batch. Effective batch = batch * grad_accum_steps.')
    parser.add_argument('--grad_accum_steps', type=int, default=8)
    parser.add_argument('--epochs', type=int, default=50)
    parser.add_argument('--warmup_epochs', type=int, default=3)
    parser.add_argument('--learning_rate', type=float, default=3e-4)
    parser.add_argument('--weight_decay', type=float, default=1e-5)
    parser.add_argument('--patience', type=int, default=15)
    parser.add_argument('--num_workers', type=int, default=0)
    parser.add_argument('--seed', type=int, default=42)
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--no_amp', action='store_true')
    parser.add_argument('--resume', action='store_true')
    parser.add_argument('--folds', type=int, default=5,
                        help='Number of K-fold cross-validation folds. Use 1 for a single-fold smoke test.')
    parser.add_argument('--tta_eval', action='store_true',
                        help='Apply 8-way TTA averaging during the final test eval and ensembling.')
    parser.add_argument('--val_subset', type=int, default=10,
                        help='How many val volumes to do sliding-window inference on each epoch '
                             '(full val is too slow). 0 = use all.')
    parser.add_argument('--only_ensemble', action='store_true',
                        help='Skip training, just run ensemble evaluation on already-trained folds.')
    args = parser.parse_args()

    data_dir = Path(args.data_dir)
    train_paths = sorted((data_dir / 'train').glob('*.npz'))
    val_paths = sorted((data_dir / 'val').glob('*.npz'))
    test_paths = sorted((data_dir / 'test').glob('*.npz'))
    if not train_paths or not val_paths or not test_paths:
        raise FileNotFoundError(
            f'Expected train/val/test/*.npz under {data_dir}. '
            'Run prepare_brats_3d_dataset.py first.'
        )
    print(f'[info] train_volumes={len(train_paths)}  val_volumes={len(val_paths)}  test_volumes={len(test_paths)}',
          flush=True)

    output_root = Path(args.output_dir)
    output_root.mkdir(parents=True, exist_ok=True)

    fold_dirs = []
    fold_results = []
    if args.folds == 1:
        # Single fold: use the canonical train/val/test split directly.
        fold_out = output_root / 'fold_0'
        fold_dirs.append(fold_out)
        if not args.only_ensemble:
            fold_results.append(train_one_fold(0, train_paths, val_paths, test_paths, args, fold_out))
    else:
        # K-fold over train+val (test stays held out across all folds).
        pool = train_paths + val_paths
        kf = KFold(n_splits=args.folds, shuffle=True, random_state=args.seed)
        for fold_idx, (tr_idx, va_idx) in enumerate(kf.split(pool)):
            fold_out = output_root / f'fold_{fold_idx}'
            fold_dirs.append(fold_out)
            if args.only_ensemble:
                continue
            tr_paths = [pool[i] for i in tr_idx]
            va_paths = [pool[i] for i in va_idx]
            fold_results.append(train_one_fold(fold_idx, tr_paths, va_paths, test_paths, args, fold_out))

    ens = ensemble_evaluate(fold_dirs, test_paths, args)
    final_payload = {'folds': fold_results, 'ensemble_test': ens, 'config': vars(args)}
    (output_root / 'evaluation_metrics.json').write_text(json.dumps(final_payload, indent=2), encoding='utf-8')
    print('\n[done] Ensemble + per-fold metrics:')
    print(json.dumps(final_payload, indent=2))


if __name__ == '__main__':
    main()