#!/usr/bin/env python

"""Train S1 segmentation models with optional raw-AEF fusion.

This script is the NERSC-ready training entrypoint for the crop-based wetland
segmentation workflow. It keeps the current loss, metrics, and schedule shape
while supporting two comparable input modes:

- S1 only: a 3-band Sentinel-1 baseline for direct DW-target comparisons.
- S1 + raw AEF: a 67-band tensor where a learned 1x1 convolution projects the
  64 raw AlphaEarth bands before fusion with S1 inside the model.
- Training reads fixed 512x512 random windows directly from the source rasters.
- Epoch validation uses deterministic 512x512 windows, while final full-scene
  metrics stream overlapping windows and blend them without loading all 67
  input bands for a scene at once.
"""

from __future__ import annotations

import argparse
import copy as copy_module
import csv
import gc
import importlib
import json
import math
import random
import shutil
import time
from collections import Counter
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple

import numpy as np
import pandas as pd
import rasterio as rio
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset
from torchvision.transforms import InterpolationMode
from torchvision.transforms import functional as TVF

from fastai.imports import *
from fastai.torch_imports import *
from fastai.vision.all import *


DEFAULT_INPUT_DIR = Path("/pscratch/sd/r/rohit9/S1ML/training_mosaics/s1_multiband")
DEFAULT_AEF_DIR = Path("/pscratch/sd/r/rohit9/S1ML/training_embeddings/alphaearth_v1_annual_int8")
DEFAULT_LABEL_DIR = Path("/pscratch/sd/r/rohit9/S1ML/training_mosaics/dw_binary")
DEFAULT_ARTIFACT_DIR = Path(
    "/pscratch/sd/r/rohit9/S1ML/training_runs/s1aef_bottleneck_resnet34_crop512"
)
DEFAULT_HF_CACHE_DIR = Path("/pscratch/sd/r/rohit9/UFO/model_cache/huggingface")

SEGFORMER_MODEL_MAP = {
    "segformer_b0": "nvidia/segformer-b0-finetuned-ade-512-512",
    "segformer_b1": "nvidia/segformer-b1-finetuned-ade-512-512",
    "segformer_b2": "nvidia/segformer-b2-finetuned-ade-512-512",
    "segformer_b3": "nvidia/segformer-b3-finetuned-ade-512-512",
    "segformer_b4": "nvidia/segformer-b4-finetuned-ade-512-512",
    "segformer_b5": "nvidia/segformer-b5-finetuned-cityscapes-1024-1024",
}


@dataclass(frozen=True)
class TripletRecord:
    name: str
    s1_path: Path
    aef_path: Optional[Path]
    label_path: Path
    sample_id: str = ""
    group_id: str = ""
    source_id: str = ""
    prior_path: Optional[Path] = None


IGNORE_INDEX = 255


def valid_target_mask(targ: torch.Tensor) -> torch.Tensor:
    """True where Dynamic World observed the pixel (label is 0 or 1)."""
    return targ != IGNORE_INDEX


def sanitized_target(targ: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
    """Replace ignored labels with 0 so one-hot/indexing stays in range."""
    return torch.where(valid, targ, torch.zeros_like(targ))


def masked_mean(per_pixel: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
    weight = valid.float()
    return (per_pixel * weight).sum() / weight.sum().clamp(min=1.0)


def masked_cross_entropy(
    pred: torch.Tensor,
    targ: torch.Tensor,
    valid: torch.Tensor,
    gamma: Optional[float] = None,
) -> torch.Tensor:
    """Cross-entropy (optionally focal) averaged over observed pixels only.

    With no ignored pixels this equals the fastai CrossEntropyLossFlat /
    FocalLossFlat it replaces, so runs on older labels are unaffected.
    """
    target = sanitized_target(targ.long(), valid)
    per_pixel = F.cross_entropy(pred.float(), target, reduction="none")
    if gamma is not None:
        p_t = torch.exp(-per_pixel)
        per_pixel = torch.pow(1.0 - p_t, gamma) * per_pixel
    return masked_mean(per_pixel, valid)


def masked_dice_loss(
    pred: torch.Tensor,
    targ: torch.Tensor,
    valid: torch.Tensor,
    axis: int = 1,
    smooth: float = 1.0,
    reduction: str = "sum",
) -> torch.Tensor:
    """Dice over observed pixels only.

    fastai's DiceLoss one-hots the target, so a raw 255 raises IndexError, and it
    has no ignore support. Zeroing logits does not work either: softmax(0, 0) is
    0.5, so ignored pixels would still contribute. Zero the probabilities and the
    one-hot target instead - ignored pixels then add nothing to either term. The
    reduction matches fastai's default ("sum" over batch x class), so with no
    ignored pixels this is numerically identical to the loss it replaces.
    """
    probability = F.softmax(pred.float(), dim=axis)
    n_classes = probability.shape[axis]
    target = sanitized_target(targ.long(), valid)
    one_hot = F.one_hot(target, n_classes).permute(0, 3, 1, 2).float()
    weight = valid.unsqueeze(1).float()
    probability = probability * weight
    one_hot = one_hot * weight
    sum_dims = list(range(2, probability.ndim))
    intersection = torch.sum(probability * one_hot, dim=sum_dims)
    union = torch.sum(probability + one_hot, dim=sum_dims)
    dice_score = (2.0 * intersection + smooth) / (union + smooth)
    loss = 1.0 - dice_score
    return loss.sum() if reduction == "sum" else loss.mean()


def binarize_label_array(data) -> "np.ndarray":
    """0 = dry, 1 = water, IGNORE_INDEX = not observed.

    Non-finite values are treated as unobserved rather than dry, so a gap never
    becomes a training negative.
    """
    values = np.asarray(data)
    mask = np.full(values.shape, IGNORE_INDEX, dtype=np.int64)
    observed = values != IGNORE_INDEX
    if np.issubdtype(values.dtype, np.floating):
        observed &= np.isfinite(values)
    mask[observed] = (values[observed] > 0.5).astype(np.int64)
    return mask


class CombinedLoss(nn.Module):
    """Focal + Dice, carried over from the working baseline."""

    def __init__(self, axis=1, smooth=1.0, alpha=1.0, gamma=2.0):
        super().__init__()
        self.axis = axis
        self.alpha = alpha
        self.smooth = smooth
        self.gamma = gamma

    def forward(self, pred, targ):
        pred = pred.float()
        targ = targ.long()
        valid = valid_target_mask(targ)
        focal = masked_cross_entropy(pred, targ, valid, gamma=self.gamma)
        dice = masked_dice_loss(pred, targ, valid, axis=self.axis, smooth=self.smooth)
        return focal + (self.alpha * dice)

    def decodes(self, x):
        return x.argmax(dim=self.axis)

    def activation(self, x):
        return F.softmax(x, dim=self.axis)


def _lovasz_grad(gt_sorted: torch.Tensor) -> torch.Tensor:
    p = len(gt_sorted)
    gts = gt_sorted.sum()
    intersection = gts - gt_sorted.float().cumsum(0)
    union = gts + (1 - gt_sorted).float().cumsum(0)
    jaccard = 1.0 - intersection / union
    if p > 1:
        jaccard[1:p] = jaccard[1:p] - jaccard[0:-1]
    return jaccard


def _lovasz_softmax_flat(
    probas: torch.Tensor, labels: torch.Tensor, classes: str = "present"
) -> torch.Tensor:
    if probas.numel() == 0:
        return probas * 0.0
    num_classes = probas.size(1)
    losses = []
    class_to_sum = list(range(num_classes)) if classes in ("all", "present") else classes
    for class_idx in class_to_sum:
        foreground = (labels == class_idx).float()
        if classes == "present" and foreground.sum() == 0:
            continue
        errors = (foreground - probas[:, class_idx]).abs()
        errors_sorted, permutation = torch.sort(errors, 0, descending=True)
        foreground_sorted = foreground[permutation.data]
        grad = _lovasz_grad(foreground_sorted)
        losses.append((errors_sorted * grad).sum())
    if not losses:
        return torch.tensor(0.0, device=probas.device, requires_grad=True)
    return torch.stack(losses).mean()


class LovaszSoftmaxLoss(nn.Module):
    """Lovasz-Softmax loss for IoU-oriented segmentation tuning."""

    def __init__(self, axis: int = 1):
        super().__init__()
        self.axis = axis

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        probas = F.softmax(pred.float(), dim=self.axis)
        _, num_classes, _, _ = probas.shape
        probas_flat = probas.permute(0, 2, 3, 1).reshape(-1, num_classes)
        targ_flat = targ.long().reshape(-1)
        keep = targ_flat != IGNORE_INDEX
        probas_flat = probas_flat[keep]
        targ_flat = targ_flat[keep]
        return _lovasz_softmax_flat(probas_flat, targ_flat, classes="present")

    def decodes(self, x: torch.Tensor) -> torch.Tensor:
        return x.argmax(dim=self.axis)

    def activation(self, x: torch.Tensor) -> torch.Tensor:
        return F.softmax(x, dim=self.axis)


class CrossEntropyDiceLoss(nn.Module):
    """Cross-entropy + Dice baseline for a less recall-heavy alternative."""

    def __init__(self, axis: int = 1, smooth: float = 1.0, alpha: float = 1.0):
        super().__init__()
        self.axis = axis
        self.alpha = alpha
        self.smooth = smooth

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        pred = pred.float()
        targ = targ.long()
        valid = valid_target_mask(targ)
        cross_entropy = masked_cross_entropy(pred, targ, valid)
        dice = masked_dice_loss(pred, targ, valid, axis=self.axis, smooth=self.smooth)
        return cross_entropy + (self.alpha * dice)

    def decodes(self, x: torch.Tensor) -> torch.Tensor:
        return x.argmax(dim=self.axis)

    def activation(self, x: torch.Tensor) -> torch.Tensor:
        return F.softmax(x, dim=self.axis)


class TwoClassLossBase(nn.Module):
    """Shared FastAI-compatible helpers for experimental water losses."""

    def __init__(self, axis: int = 1):
        super().__init__()
        self.axis = axis

    def decodes(self, x: torch.Tensor) -> torch.Tensor:
        return x.argmax(dim=self.axis)

    def activation(self, x: torch.Tensor) -> torch.Tensor:
        return F.softmax(x, dim=self.axis)

    def water_probability_and_target(
        self, pred: torch.Tensor, targ: torch.Tensor
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Water probability, water target and the observed-pixel mask.

        Probability and target are zeroed where the label is ignored, so any
        sum over them skips unobserved pixels.
        """
        valid = valid_target_mask(targ.long())
        weight = valid.float()
        probability = F.softmax(pred.float(), dim=self.axis)[:, 1] * weight
        target = (targ.long() == 1).float() * weight
        return probability, target, valid


class CrossEntropySoftIoULoss(TwoClassLossBase):
    """Cross-entropy plus a direct soft Jaccard/IoU objective."""

    def __init__(self, axis: int = 1, smooth: float = 1.0, aux_weight: float = 1.0):
        super().__init__(axis=axis)
        self.smooth = smooth
        self.aux_weight = aux_weight

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        probability, target, valid = self.water_probability_and_target(pred, targ)
        dims = tuple(range(1, probability.ndim))
        intersection = (probability * target).sum(dim=dims)
        union = (probability + target - probability * target).sum(dim=dims)
        soft_iou_loss = 1.0 - ((intersection + self.smooth) / (union + self.smooth)).mean()
        return masked_cross_entropy(pred, targ, valid) + self.aux_weight * soft_iou_loss


class FocalTverskyCELoss(TwoClassLossBase):
    """CE plus focal Tversky, exposing FP/FN asymmetry for water mapping."""

    def __init__(
        self,
        axis: int = 1,
        alpha: float = 0.3,
        beta: float = 0.7,
        gamma: float = 1.33,
        smooth: float = 1.0,
        aux_weight: float = 1.0,
    ):
        super().__init__(axis=axis)
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.smooth = smooth
        self.aux_weight = aux_weight

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        probability, target, valid = self.water_probability_and_target(pred, targ)
        dims = tuple(range(1, probability.ndim))
        true_positive = (probability * target).sum(dim=dims)
        false_positive = (probability * (1.0 - target)).sum(dim=dims)
        false_negative = ((1.0 - probability) * target).sum(dim=dims)
        tversky = (true_positive + self.smooth) / (
            true_positive
            + self.alpha * false_positive
            + self.beta * false_negative
            + self.smooth
        )
        focal_tversky = torch.pow(1.0 - tversky, self.gamma).mean()
        return masked_cross_entropy(pred, targ, valid) + self.aux_weight * focal_tversky


class BoundaryAwareCEDiceLoss(TwoClassLossBase):
    """CE+Dice with an auxiliary soft edge-Dice term for shorelines."""

    def __init__(self, axis: int = 1, smooth: float = 1.0, aux_weight: float = 0.2):
        super().__init__(axis=axis)
        self.smooth = smooth
        self.aux_weight = aux_weight
        self.dice_loss = DiceLoss(axis, smooth)
        self.register_buffer(
            "sobel_x",
            torch.tensor(
                [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]]
            ).view(1, 1, 3, 3),
        )
        self.register_buffer(
            "sobel_y",
            torch.tensor(
                [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]]
            ).view(1, 1, 3, 3),
        )

    def edge_strength(self, image: torch.Tensor) -> torch.Tensor:
        image = image.unsqueeze(1)
        grad_x = F.conv2d(image, self.sobel_x.to(image), padding=1)
        grad_y = F.conv2d(image, self.sobel_y.to(image), padding=1)
        return torch.sqrt(grad_x.square() + grad_y.square() + 1e-6).clamp(0.0, 1.0)

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        probability, target, valid = self.water_probability_and_target(pred, targ)
        edge_weight = valid.float().unsqueeze(1)
        pred_edge = self.edge_strength(probability) * edge_weight
        target_edge = self.edge_strength(target) * edge_weight
        dims = tuple(range(1, pred_edge.ndim))
        intersection = (pred_edge * target_edge).sum(dim=dims)
        denominator = pred_edge.sum(dim=dims) + target_edge.sum(dim=dims)
        boundary_loss = 1.0 - ((2.0 * intersection + self.smooth) / (denominator + self.smooth)).mean()
        base = masked_cross_entropy(pred, targ, valid) + masked_dice_loss(
            pred, targ, valid, axis=self.axis, smooth=self.smooth
        )
        return base + self.aux_weight * boundary_loss


def _soft_erode(image: torch.Tensor) -> torch.Tensor:
    vertical = -F.max_pool2d(-image, kernel_size=(3, 1), stride=1, padding=(1, 0))
    horizontal = -F.max_pool2d(-image, kernel_size=(1, 3), stride=1, padding=(0, 1))
    return torch.minimum(vertical, horizontal)


def _soft_dilate(image: torch.Tensor) -> torch.Tensor:
    return F.max_pool2d(image, kernel_size=3, stride=1, padding=1)


def _soft_skeletonize(image: torch.Tensor, iterations: int) -> torch.Tensor:
    opened = _soft_dilate(_soft_erode(image))
    skeleton = F.relu(image - opened)
    for _ in range(iterations):
        image = _soft_erode(image)
        opened = _soft_dilate(_soft_erode(image))
        delta = F.relu(image - opened)
        skeleton = skeleton + F.relu(delta - skeleton * delta)
    return skeleton


class TopologyAwareCEDiceLoss(TwoClassLossBase):
    """CE+Dice+soft-clDice for connected thin water networks."""

    def __init__(
        self,
        axis: int = 1,
        smooth: float = 1.0,
        aux_weight: float = 0.2,
        skeleton_iterations: int = 10,
    ):
        super().__init__(axis=axis)
        self.smooth = smooth
        self.aux_weight = aux_weight
        self.skeleton_iterations = skeleton_iterations
        self.dice_loss = DiceLoss(axis, smooth)

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        probability, target, valid = self.water_probability_and_target(pred, targ)
        probability = probability.unsqueeze(1)
        target = target.unsqueeze(1)
        pred_skeleton = _soft_skeletonize(probability, self.skeleton_iterations)
        target_skeleton = _soft_skeletonize(target, self.skeleton_iterations)
        dims = tuple(range(1, probability.ndim))
        topology_precision = (
            (pred_skeleton * target).sum(dim=dims) + self.smooth
        ) / (pred_skeleton.sum(dim=dims) + self.smooth)
        topology_sensitivity = (
            (target_skeleton * probability).sum(dim=dims) + self.smooth
        ) / (target_skeleton.sum(dim=dims) + self.smooth)
        cldice = (
            2.0 * topology_precision * topology_sensitivity
            / (topology_precision + topology_sensitivity + 1e-6)
        ).mean()
        base = masked_cross_entropy(pred, targ, valid) + masked_dice_loss(
            pred, targ, valid, axis=self.axis, smooth=self.smooth
        )
        return base + self.aux_weight * (1.0 - cldice)


class SymmetricCrossEntropyDiceLoss(TwoClassLossBase):
    """Noise-robust symmetric cross-entropy plus Dice for weak DW labels."""

    def __init__(self, axis: int = 1, smooth: float = 1.0, aux_weight: float = 0.1):
        super().__init__(axis=axis)
        self.aux_weight = aux_weight
        self.smooth = smooth

    def forward(self, pred: torch.Tensor, targ: torch.Tensor) -> torch.Tensor:
        logits = pred.float()
        labels = targ.long()
        valid = valid_target_mask(labels)
        safe_labels = sanitized_target(labels, valid)
        probabilities = F.softmax(logits, dim=self.axis).clamp(1e-7, 1.0)
        one_hot = F.one_hot(safe_labels, num_classes=logits.shape[1]).permute(0, 3, 1, 2).float()
        reverse_ce_map = -(probabilities * torch.log(one_hot.clamp(1e-4, 1.0))).sum(dim=1)
        reverse_ce = masked_mean(reverse_ce_map, valid)
        return (
            masked_cross_entropy(logits, labels, valid)
            + self.aux_weight * reverse_ce
            + masked_dice_loss(logits, labels, valid, axis=self.axis, smooth=self.smooth)
        )


class MaskedDice(Metric):
    """fastai Dice restricted to observed pixels.

    The stock fastai Dice/JaccardCoeff multiply and add raw target values, so a
    label of 255 corrupts both the intersection and the union.
    """

    def __init__(self, axis: int = 1):
        self.axis = axis
        self.reset()

    def reset(self):
        self.inter = 0.0
        self.union = 0.0

    def accumulate(self, learn):
        pred = learn.pred.argmax(dim=self.axis)
        targ = learn.y
        valid = targ != IGNORE_INDEX
        water_pred = (pred == 1) & valid
        water_targ = (targ == 1) & valid
        self.inter += (water_pred & water_targ).float().sum().item()
        self.union += water_pred.float().sum().item() + water_targ.float().sum().item()

    @property
    def value(self):
        return 2.0 * self.inter / self.union if self.union > 0 else None

    @property
    def name(self):
        return "dice"


class MaskedJaccardCoeff(MaskedDice):
    """fastai JaccardCoeff restricted to observed pixels."""

    @property
    def value(self):
        denominator = self.union - self.inter
        return self.inter / denominator if denominator > 0 else None

    @property
    def name(self):
        return "jaccard_coeff"


class WaterIoU(Metric):
    """Per-class IoU for the water class (index 1)."""

    def __init__(self):
        self.inter = 0
        self.union = 0

    def reset(self):
        self.inter = 0
        self.union = 0

    def accumulate(self, learn):
        pred = learn.pred.argmax(dim=1)
        targ = learn.y
        valid = targ != IGNORE_INDEX
        water_pred = (pred == 1) & valid
        water_targ = (targ == 1) & valid
        self.inter += (water_pred & water_targ).float().sum().item()
        self.union += (water_pred | water_targ).float().sum().item()

    @property
    def value(self):
        return self.inter / self.union if self.union > 0 else 0.0

    @property
    def name(self):
        return "water_iou"


class SourceBalancedWaterIoU(Metric):
    """Mean water IoU across validation sources, giving each source equal weight."""

    def __init__(self, source_ids: Sequence[str]):
        self.source_ids = list(source_ids)
        self.reset()

    def reset(self):
        self.cursor = 0
        self.intersections: Dict[str, float] = {}
        self.unions: Dict[str, float] = {}

    def accumulate(self, learn):
        pred = learn.pred.argmax(dim=1)
        targ = learn.y
        batch_size = int(pred.shape[0])
        if self.cursor + batch_size > len(self.source_ids):
            raise RuntimeError(
                "Validation source tracking exceeded the configured source sequence."
            )
        for batch_idx in range(batch_size):
            source_id = self.source_ids[self.cursor + batch_idx] or "unspecified"
            valid = targ[batch_idx] != IGNORE_INDEX
            water_pred = (pred[batch_idx] == 1) & valid
            water_targ = (targ[batch_idx] == 1) & valid
            self.intersections[source_id] = self.intersections.get(source_id, 0.0) + (
                water_pred & water_targ
            ).float().sum().item()
            self.unions[source_id] = self.unions.get(source_id, 0.0) + (
                water_pred | water_targ
            ).float().sum().item()
        self.cursor += batch_size

    @property
    def value(self):
        source_ious = [
            self.intersections.get(source_id, 0.0) / union
            for source_id, union in self.unions.items()
            if union > 0
        ]
        return float(np.mean(source_ious)) if source_ious else 0.0

    @property
    def name(self):
        return "source_balanced_water_iou"


class TargetSourceWaterIoU(Metric):
    """Pooled water IoU for one validation source used as the checkpoint objective."""

    def __init__(self, source_ids: Sequence[str], target_source_id: str):
        self.source_ids = list(source_ids)
        self.target_source_id = target_source_id
        self.reset()

    def reset(self):
        self.cursor = 0
        self.intersection = 0.0
        self.union = 0.0

    def accumulate(self, learn):
        pred = learn.pred.argmax(dim=1)
        targ = learn.y
        batch_size = int(pred.shape[0])
        if self.cursor + batch_size > len(self.source_ids):
            raise RuntimeError(
                "Validation source tracking exceeded the configured source sequence."
            )
        for batch_idx in range(batch_size):
            source_id = self.source_ids[self.cursor + batch_idx] or "unspecified"
            if source_id != self.target_source_id:
                continue
            valid = targ[batch_idx] != IGNORE_INDEX
            water_pred = (pred[batch_idx] == 1) & valid
            water_targ = (targ[batch_idx] == 1) & valid
            self.intersection += (water_pred & water_targ).float().sum().item()
            self.union += (water_pred | water_targ).float().sum().item()
        self.cursor += batch_size

    @property
    def value(self):
        return self.intersection / self.union if self.union > 0 else 0.0

    @property
    def name(self):
        return "target_source_water_iou"


class NonFiniteLossCallback(Callback):
    """Stop loudly on NaN/Inf instead of silently cancelling the fit."""

    order = TerminateOnNaNCallback.order

    def after_batch(self):
        if not self.training or self.loss is None:
            return

        loss = self.loss.detach().float()
        if torch.isfinite(loss).all():
            return

        lr = None
        if getattr(self.learn, "opt", None) is not None and len(self.opt.hypers) > 0:
            lr = self.opt.hypers[-1].get("lr")

        debug_summary = {
            "phase": getattr(self.learn, "training_phase", "unknown"),
            "epoch": int(getattr(self.learn, "epoch", -1)),
            "iter": int(getattr(self.learn, "iter", -1)),
            "loss": float(loss.item()),
            "lr": None if lr is None else float(lr),
        }
        print({"non_finite_loss": debug_summary})
        raise RuntimeError("Non-finite loss detected during {phase}.".format(**debug_summary))


class EpochTimerCallback(Callback):
    """Track wall-clock epoch durations for benchmark summaries."""

    def before_fit(self):
        self.epoch_times = []
        self._epoch_start = None

    def before_epoch(self):
        self._epoch_start = time.perf_counter()

    def after_epoch(self):
        if self._epoch_start is None:
            return
        self.epoch_times.append(time.perf_counter() - self._epoch_start)
        self._epoch_start = None


def load_segformer_components():
    try:
        transformers = importlib.import_module("transformers")
    except ModuleNotFoundError as error:
        raise ModuleNotFoundError(
            "SegFormer support requires the 'transformers' package. "
            "Use a SegFormer-capable environment or keep --model_family=resnet34_unet."
        ) from error
    return transformers.SegformerConfig, transformers.SegformerForSemanticSegmentation


class AEFBottleneckUNet(nn.Module):
    """Wrap a U-Net with a learned raw-AEF bottleneck projection."""

    def __init__(
        self,
        unet_model: nn.Module,
        n_s1: int = 3,
        n_aef: int = 64,
        n_proj: int = 4,
        n_prior: int = 0,
        prior_dropout_p: float = 0.0,
        aef_dropout_p: float = 0.0,
    ):
        super().__init__()
        self.aef_dropout_p = aef_dropout_p
        self.n_s1 = n_s1
        self.n_aef = n_aef
        self.n_proj = n_proj
        self.n_prior = n_prior
        self.prior_dropout_p = prior_dropout_p
        self.aef_proj = nn.Conv2d(n_aef, n_proj, kernel_size=1, bias=False)
        nn.init.xavier_uniform_(self.aef_proj.weight, gain=0.1)
        self.unet = unet_model

    def forward(self, x):
        s1 = x[:, : self.n_s1]
        aef = x[:, self.n_s1 : self.n_s1 + self.n_aef]
        aef_projected = self.aef_proj(aef)
        if self.training and self.aef_dropout_p > 0:
            # factorial_v1: per-sample AEF dropout (whole embedding zeroed, no rescaling), training only
            keep = (
                torch.rand((aef_projected.shape[0], 1, 1, 1), device=aef_projected.device)
                >= self.aef_dropout_p
            ).to(aef_projected.dtype)
            aef_projected = aef_projected * keep
        components = [s1, aef_projected]
        if self.n_prior:
            prior = x[
                :,
                self.n_s1 + self.n_aef : self.n_s1 + self.n_aef + self.n_prior,
            ]
            if self.training and self.prior_dropout_p > 0:
                keep = (
                    torch.rand((prior.shape[0], 1, 1, 1), device=prior.device)
                    >= self.prior_dropout_p
                ).to(prior.dtype)
                prior = prior * keep
            components.append(prior)
        fused = torch.cat(components, dim=1)
        return self.unet(fused)


class AEFBottleneckSegFormer(nn.Module):
    """Wrap a SegFormer with the same learned AEF projection used by the U-Net."""

    def __init__(
        self,
        model_name: str,
        num_labels: int,
        n_s1: int = 3,
        n_aef: int = 64,
        n_proj: int = 4,
        hf_cache_dir: Optional[Path] = None,
        pretrained: bool = True,
        local_files_only: bool = False,
    ):
        super().__init__()
        self.n_s1 = n_s1
        self.n_aef = n_aef
        self.n_proj = n_proj
        self.aef_proj = nn.Conv2d(n_aef, n_proj, kernel_size=1, bias=False)
        nn.init.xavier_uniform_(self.aef_proj.weight, gain=0.1)

        SegformerConfig, SegformerForSemanticSegmentation = load_segformer_components()
        cache_dir = str(hf_cache_dir) if hf_cache_dir is not None else None
        if pretrained:
            config = SegformerConfig.from_pretrained(
                model_name,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            config.num_labels = num_labels
            self.segformer = SegformerForSemanticSegmentation.from_pretrained(
                model_name,
                config=config,
                ignore_mismatched_sizes=True,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
        else:
            config = SegformerConfig.from_pretrained(
                model_name,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
            )
            config.num_labels = num_labels
            self.segformer = SegformerForSemanticSegmentation(config)

        self._adapt_patch_embedding(in_channels=n_s1 + n_proj)

    def _adapt_patch_embedding(self, in_channels: int) -> None:
        old_proj = self.segformer.segformer.encoder.patch_embeddings[0].proj
        new_proj = nn.Conv2d(
            in_channels,
            old_proj.out_channels,
            kernel_size=old_proj.kernel_size,
            stride=old_proj.stride,
            padding=old_proj.padding,
            bias=old_proj.bias is not None,
        )
        with torch.no_grad():
            new_proj.weight.zero_()
            copied_channels = min(old_proj.in_channels, in_channels)
            if copied_channels > 0:
                new_proj.weight[:, :copied_channels] = old_proj.weight[:, :copied_channels]
            if in_channels > copied_channels:
                mean_weight = old_proj.weight.mean(dim=1, keepdim=True)
                new_proj.weight[:, copied_channels:] = mean_weight
            if old_proj.bias is not None and new_proj.bias is not None:
                new_proj.bias.copy_(old_proj.bias)
        self.segformer.segformer.encoder.patch_embeddings[0].proj = new_proj

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        s1 = x[:, : self.n_s1]
        aef = x[:, self.n_s1 : self.n_s1 + self.n_aef]
        aef_projected = self.aef_proj(aef)
        fused = torch.cat([s1, aef_projected], dim=1)
        logits = self.segformer(pixel_values=fused).logits
        if logits.shape[-2:] != fused.shape[-2:]:
            logits = F.interpolate(
                logits,
                size=fused.shape[-2:],
                mode="bilinear",
                align_corners=False,
            )
        return logits


class S1AEFCropDataset(Dataset):
    """Training dataset that samples aligned random crops from triplet tiles."""

    def __init__(
        self,
        records: Sequence[TripletRecord],
        crop_size: int,
        band_mean: np.ndarray,
        band_std: np.ndarray,
        n_s1_bands: int,
        rotate_max_deg: float,
        rotate_p: float,
        scale_jitter_min: float,
        scale_jitter_max: float,
        s1_gain_jitter: float,
        s1_bias_jitter: float,
        gaussian_noise_std: float,
        gaussian_noise_p: float,
        positive_crop_p: float,
        min_water_fraction: float,
        positive_crop_attempts: int,
    ):
        self.records = list(records)
        self.crop_size = crop_size
        self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1)
        self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1)
        self.n_s1_bands = n_s1_bands
        self.rotate_max_deg = rotate_max_deg
        self.rotate_p = rotate_p
        self.scale_jitter_min = scale_jitter_min
        self.scale_jitter_max = scale_jitter_max
        self.s1_gain_jitter = s1_gain_jitter
        self.s1_bias_jitter = s1_bias_jitter
        self.gaussian_noise_std = gaussian_noise_std
        self.gaussian_noise_p = gaussian_noise_p
        self.positive_crop_p = positive_crop_p
        self.min_water_fraction = min_water_fraction
        self.positive_crop_attempts = positive_crop_attempts

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

    def __getitem__(self, idx):
        image, mask = load_random_triplet_window(
            self.records[idx],
            self.crop_size,
            seek_positive=random.random() < self.positive_crop_p,
            min_water_fraction=self.min_water_fraction,
            attempts=self.positive_crop_attempts,
        )
        image, mask = apply_train_augmentations(
            image=image,
            mask=mask,
            crop_size=self.crop_size,
            n_s1_bands=self.n_s1_bands,
            rotate_max_deg=self.rotate_max_deg,
            rotate_p=self.rotate_p,
            scale_jitter_min=self.scale_jitter_min,
            scale_jitter_max=self.scale_jitter_max,
            s1_gain_jitter=self.s1_gain_jitter,
            s1_bias_jitter=self.s1_bias_jitter,
            gaussian_noise_std=self.gaussian_noise_std,
            gaussian_noise_p=self.gaussian_noise_p,
            normalized_band_count=int(self.band_mean.shape[0]),
        )
        image = normalize_image(image, self.band_mean, self.band_std)
        return image, mask.long()


class S1AEFValidationPatchDataset(Dataset):
    """Stable validation patches used by fastai during each training epoch."""

    def __init__(
        self,
        records: Sequence[TripletRecord],
        crop_size: int,
        band_mean: np.ndarray,
        band_std: np.ndarray,
    ):
        self.records = list(records)
        self.crop_size = crop_size
        self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1)
        self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1)

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

    def __getitem__(self, idx):
        image, mask = load_center_triplet_window(self.records[idx], self.crop_size)
        image = normalize_image(image, self.band_mean, self.band_std)
        return image, mask.long()


def read_tile_index(path: Path, records: Sequence[TripletRecord], split: str) -> List[Tuple[TripletRecord, int, int]]:
    """full_coverage_v1: tiles (record, top, left) of the given split with at least one observed label pixel."""
    table = pd.read_csv(path)
    table = table[(table["split"] == split) & (table["observed_frac"] > 0)]
    by_name = {record.name: record for record in records}
    table = table[table["tile_name"].isin(by_name)]
    missing = sorted(set(by_name) - set(table["tile_name"]))
    if missing:
        raise ValueError("{n} {split} records have no tile in {path}, e.g. {ex}".format(n=len(missing), split=split, path=path, ex=missing[:3]))
    return [(by_name[name], int(top), int(left)) for name, top, left in zip(table["tile_name"], table["top"], table["left"])]


class S1AEFTileDataset(Dataset):
    """full_coverage_v1 training set: every tile of every training chip once per epoch (no random crop)."""

    def __init__(self, tiles, crop_size, band_mean, band_std, n_s1_bands, rotate_max_deg, rotate_p, scale_jitter_min,
                 scale_jitter_max, s1_gain_jitter, s1_bias_jitter, gaussian_noise_std, gaussian_noise_p):
        self.tiles = list(tiles)
        self.records = sorted({tile[0].name: tile[0] for tile in self.tiles}.values(), key=lambda r: r.name)
        self.crop_size = crop_size
        self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1)
        self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1)
        self.aug = dict(n_s1_bands=n_s1_bands, rotate_max_deg=rotate_max_deg, rotate_p=rotate_p, scale_jitter_min=scale_jitter_min,
                        scale_jitter_max=scale_jitter_max, s1_gain_jitter=s1_gain_jitter, s1_bias_jitter=s1_bias_jitter,
                        gaussian_noise_std=gaussian_noise_std, gaussian_noise_p=gaussian_noise_p)

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

    def __getitem__(self, idx):
        record, top, left = self.tiles[idx]
        image, mask = load_triplet_window(record, top=top, left=left, height=self.crop_size, width=self.crop_size)
        image, mask = apply_train_augmentations(image=image, mask=mask, crop_size=self.crop_size,
                                                normalized_band_count=int(self.band_mean.shape[0]), **self.aug)
        image = normalize_image(image, self.band_mean, self.band_std)
        return image, mask.long()


class S1AEFValidationTileDataset(Dataset):
    """full_coverage_v1 validation: every tile of every validation chip, fixed order. Exposes the whole-chip records
    (for the windowed end-of-training evaluation) and one source id per tile (for the per-epoch metrics)."""

    def __init__(self, tiles, records, crop_size, band_mean, band_std):
        self.tiles = list(tiles)
        self.records = list(records)
        self.item_source_ids = [tile[0].source_id for tile in self.tiles]
        self.crop_size = crop_size
        self.band_mean = torch.from_numpy(band_mean.astype(np.float32)).view(-1, 1, 1)
        self.band_std = torch.from_numpy(band_std.astype(np.float32)).view(-1, 1, 1)

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

    def __getitem__(self, idx):
        record, top, left = self.tiles[idx]
        image, mask = load_triplet_window(record, top=top, left=left, height=self.crop_size, width=self.crop_size)
        image = normalize_image(image, self.band_mean, self.band_std)
        return image, mask.long()


def add_bool_arg(parser: argparse.ArgumentParser, name: str, default: bool, help_text: str) -> None:
    parser.add_argument("--{name}".format(name=name), dest=name, action="store_true", help=help_text)
    parser.add_argument(
        "--no-{name}".format(name=name),
        dest=name,
        action="store_false",
        help="Disable {name}".format(name=name.replace("_", " ")),
    )
    parser.set_defaults(**{name: default})


def using_aef(args: argparse.Namespace) -> bool:
    return args.n_aef_bands > 0


def using_prior(args: argparse.Namespace) -> bool:
    return bool(getattr(args, "use_permanent_prior", False))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Train S1 segmentation models with optional raw-AEF fusion.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument("--input_dir", type=Path, default=DEFAULT_INPUT_DIR)
    parser.add_argument("--aef_dir", type=Path, default=DEFAULT_AEF_DIR)
    parser.add_argument("--label_dir", type=Path, default=DEFAULT_LABEL_DIR)
    parser.add_argument(
        "--triplet_manifest",
        type=Path,
        default=None,
        help="CSV with explicit tile_name/name, s1_path, label_path, and optional aef_path rows.",
    )
    parser.add_argument(
        "--triplet_sample_id",
        default=None,
        help="Optional sample_id filter for --triplet_manifest. Required when it contains multiple labels.",
    )
    parser.add_argument(
        "--fixed_split_manifest",
        type=Path,
        default=None,
        help="CSV with split and tile_name columns. Locks train/validation membership across experiments.",
    )
    parser.add_argument("--artifact_dir", type=Path, default=DEFAULT_ARTIFACT_DIR)
    parser.add_argument(
        "--warm_start_run_dir",
        type=Path,
        default=None,
        help="Parent run directory with a compatible best checkpoint and band_stats.npz.",
    )
    parser.add_argument(
        "--warm_start_checkpoint",
        type=Path,
        default=None,
        help="Explicit checkpoint for strict warm-start loading.",
    )
    parser.add_argument(
        "--band_stats_path",
        type=Path,
        default=None,
        help="Precomputed band_stats.npz to copy into the new artifact directory.",
    )
    parser.add_argument(
        "--model_family",
        choices=["resnet34_unet", "resnet50_unet", "resnet101_unet", "segformer"],
        default="resnet34_unet",
    )
    parser.add_argument(
        "--segformer_variant",
        choices=list(SEGFORMER_MODEL_MAP.keys()),
        default="segformer_b4",
    )
    parser.add_argument("--hf_cache_dir", type=Path, default=DEFAULT_HF_CACHE_DIR)
    parser.add_argument("--benchmark_summary_path", type=Path, default=None)
    parser.add_argument("--mode", choices=["benchmark", "train"], default="train")
    parser.add_argument("--crop_size", type=int, default=512)
    parser.add_argument("--tile_overlap", type=int, default=64)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--num_workers", type=int, default=8)
    parser.add_argument(
        "--disable_pin_memory",
        action="store_true",
        help="Disable DataLoader pinned memory for clusters where CUDA pinning is unreliable.",
    )
    parser.add_argument("--device_id", type=int, default=0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--split_seed",
        type=int,
        default=None,
        help="Seed for the train/validation split. Defaults to --seed.",
    )
    parser.add_argument(
        "--train_subset_seed",
        type=int,
        default=None,
        help="Seed for deterministic train-record subsampling. Defaults to --seed.",
    )
    parser.add_argument("--valid_pct", type=float, default=0.2)
    parser.add_argument(
        "--max_train_records",
        type=int,
        default=None,
        help="Cap the number of training records after the validation split.",
    )
    parser.add_argument("--stats_sample_size", type=int, default=256)
    parser.add_argument("--freeze_epochs", type=int, default=2)
    parser.add_argument("--partial_unfreeze_epochs", type=int, default=2)
    parser.add_argument("--finetune_epochs", type=int, default=12)
    parser.add_argument("--finetune_tail_epochs", type=int, default=10)
    parser.add_argument("--warmup_lr", type=float, default=1e-4)
    parser.add_argument("--partial_unfreeze_lr", type=float, default=1e-5)
    parser.add_argument("--unfreeze_lr", type=float, default=3e-6)
    parser.add_argument("--finetune_tail_lr", type=float, default=1.5e-6)
    parser.add_argument(
        "--tuning_strategy",
        choices=["staged", "head_only", "decoder_aef", "last_encoder_decoder_aef", "full_low_lr"],
        default="staged",
        help="Explicit trainable scope for transfer-learning strategy comparisons.",
    )
    parser.add_argument("--strategy_epochs", type=int, default=6)
    parser.add_argument("--strategy_lr", type=float, default=1e-5)
    parser.add_argument("--weight_decay", type=float, default=1e-2)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--grad_accum_steps", type=int, default=1)
    parser.add_argument("--optimizer", choices=["ranger", "adamw"], default="ranger")
    parser.add_argument("--schedule_policy", choices=["one_cycle", "flat_cos"], default="one_cycle")
    parser.add_argument("--patience", type=int, default=4)
    parser.add_argument(
        "--selection_metric",
        choices=["dice", "water_iou", "source_balanced_water_iou", "target_source_water_iou"],
        default="water_iou",
    )
    parser.add_argument(
        "--target_source_id",
        default="global_postflood",
        help="Source to oversample when --target_source_fraction is set.",
    )
    parser.add_argument(
        "--target_source_fraction",
        type=float,
        default=None,
        help="Expected fraction of each training epoch drawn from --target_source_id.",
    )
    parser.add_argument("--min_delta", type=float, default=0.001)
    parser.add_argument(
        "--loss",
        choices=[
            "focal_dice",
            "ce_dice",
            "lovasz",
            "ce_soft_iou",
            "focal_tversky",
            "ce_dice_boundary",
            "ce_dice_cldice",
            "sce_dice",
        ],
        default="focal_dice",
    )
    parser.add_argument("--loss_aux_weight", type=float, default=None)
    parser.add_argument("--tversky_alpha", type=float, default=0.3)
    parser.add_argument("--tversky_beta", type=float, default=0.7)
    parser.add_argument("--tversky_gamma", type=float, default=1.33)
    parser.add_argument("--cldice_iterations", type=int, default=10)
    parser.add_argument("--rotate_max_deg", type=float, default=5.0)
    parser.add_argument("--rotate_p", type=float, default=0.5)
    parser.add_argument("--scale_jitter_min", type=float, default=1.0)
    parser.add_argument("--scale_jitter_max", type=float, default=1.0)
    parser.add_argument("--s1_gain_jitter", type=float, default=0.0)
    parser.add_argument("--s1_bias_jitter", type=float, default=0.0)
    parser.add_argument("--gaussian_noise_std", type=float, default=0.0)
    parser.add_argument("--gaussian_noise_p", type=float, default=0.0)
    parser.add_argument(
        "--positive_crop_p",
        type=float,
        default=0.0,
        help="Probability of trying to draw a crop with at least --min_water_fraction water.",
    )
    parser.add_argument("--min_water_fraction", type=float, default=0.005)
    parser.add_argument("--positive_crop_attempts", type=int, default=8)
    parser.add_argument("--train_tile_index", type=Path, default=None,
                        help="full_coverage_v1: tile index CSV; train on every tile of every chip each epoch.")
    parser.add_argument("--valid_tile_index", type=Path, default=None,
                        help="full_coverage_v1: tile index CSV; per-epoch validation on every tile of every validation chip.")
    parser.add_argument("--max_train_batches", type=int, default=None)
    parser.add_argument("--max_valid_batches", type=int, default=None)
    parser.add_argument("--probe_batch_sizes", type=int, nargs="+", default=[12, 8, 4, 2])
    parser.add_argument("--max_benchmark_batch_size", type=int, default=None)
    parser.add_argument("--eval_thresholds", type=float, nargs="+", default=None)
    parser.add_argument("--worst_tile_count", type=int, default=25)
    parser.add_argument("--n_s1_bands", type=int, default=3)
    parser.add_argument("--n_aef_bands", type=int, default=64)
    parser.add_argument("--n_proj_bands", type=int, default=4)
    parser.add_argument(
        "--aef_dropout_p",
        type=float,
        default=0.0,
        help="factorial_v1: per-sample probability of zeroing the projected AEF input during training.",
    )
    parser.add_argument(
        "--prior_dropout_p",
        type=float,
        default=0.0,
        help="Per-sample probability of zeroing the permanent-water prior during training.",
    )
    add_bool_arg(
        parser,
        "use_permanent_prior",
        False,
        "Append prior_path probability maps from the triplet manifest as a raw input channel.",
    )
    add_bool_arg(parser, "use_fp16", True, "Enable mixed precision on GPU.")
    add_bool_arg(parser, "pretrained", True, "Use ImageNet-pretrained encoder weights.")
    add_bool_arg(parser, "hf_local_files_only", False, "Restrict SegFormer loading to the local Hugging Face cache.")
    add_bool_arg(parser, "force_tiled_validation", False, "Force tiled inference during validation and threshold sweeps.")
    args = parser.parse_args()

    args.codes = ["other", "water"]
    if args.n_aef_bands < 0:
        parser.error("--n_aef_bands must be >= 0")
    if args.n_aef_bands == 0 and args.n_proj_bands != 0:
        parser.error("--n_proj_bands must be 0 when --n_aef_bands=0")
    if args.n_aef_bands == 0 and args.model_family == "segformer":
        parser.error("--model_family=segformer is not yet supported with --n_aef_bands=0")
    if using_prior(args) and args.model_family == "segformer":
        parser.error("--use_permanent_prior is currently supported only with ResNet U-Nets")
    if not 0 <= args.prior_dropout_p < 1:
        parser.error("--prior_dropout_p must be in [0, 1)")
    if args.prior_dropout_p > 0 and not using_prior(args):
        parser.error("--prior_dropout_p requires --use_permanent_prior")

    if args.model_family == "segformer":
        variant_name = args.segformer_variant.replace("segformer_", "")
        if using_aef(args):
            args.best_model_name = "s1aef_bottleneck_segformer_{variant}_best".format(
                variant=variant_name
            )
            args.history_csv = "s1aef_bottleneck_segformer_{variant}_history.csv".format(
                variant=variant_name
            )
            args.threshold_sweep_csv = "s1aef_bottleneck_segformer_{variant}_threshold_sweep.csv".format(
                variant=variant_name
            )
            args.eval_summary_json = "s1aef_bottleneck_segformer_{variant}_eval_summary.json".format(
                variant=variant_name
            )
            args.tile_metrics_csv = "s1aef_bottleneck_segformer_{variant}_tile_metrics.csv".format(
                variant=variant_name
            )
            args.worst_tiles_csv = "s1aef_bottleneck_segformer_{variant}_worst_tiles.csv".format(
                variant=variant_name
            )
    else:
        encoder_name = args.model_family.replace("_unet", "")
        if using_aef(args):
            args.best_model_name = "s1aef_bottleneck_{encoder}_best".format(encoder=encoder_name)
            args.history_csv = "s1aef_bottleneck_{encoder}_history.csv".format(encoder=encoder_name)
            args.threshold_sweep_csv = "s1aef_bottleneck_{encoder}_threshold_sweep.csv".format(encoder=encoder_name)
            args.eval_summary_json = "s1aef_bottleneck_{encoder}_eval_summary.json".format(encoder=encoder_name)
            args.tile_metrics_csv = "s1aef_bottleneck_{encoder}_tile_metrics.csv".format(encoder=encoder_name)
            args.worst_tiles_csv = "s1aef_bottleneck_{encoder}_worst_tiles.csv".format(encoder=encoder_name)
        else:
            args.best_model_name = "s1dw_{encoder}_best".format(encoder=encoder_name)
            args.history_csv = "s1dw_{encoder}_history.csv".format(encoder=encoder_name)
            args.threshold_sweep_csv = "s1dw_{encoder}_threshold_sweep.csv".format(encoder=encoder_name)
            args.eval_summary_json = "s1dw_{encoder}_eval_summary.json".format(encoder=encoder_name)
            args.tile_metrics_csv = "s1dw_{encoder}_tile_metrics.csv".format(encoder=encoder_name)
            args.worst_tiles_csv = "s1dw_{encoder}_worst_tiles.csv".format(encoder=encoder_name)
    if using_prior(args):
        encoder_name = args.model_family.replace("_unet", "")
        args.best_model_name = "s1aef_prior_{encoder}_best".format(encoder=encoder_name)
        args.history_csv = "s1aef_prior_{encoder}_history.csv".format(encoder=encoder_name)
        args.threshold_sweep_csv = "s1aef_prior_{encoder}_threshold_sweep.csv".format(encoder=encoder_name)
        args.eval_summary_json = "s1aef_prior_{encoder}_eval_summary.json".format(encoder=encoder_name)
        args.tile_metrics_csv = "s1aef_prior_{encoder}_tile_metrics.csv".format(encoder=encoder_name)
        args.worst_tiles_csv = "s1aef_prior_{encoder}_worst_tiles.csv".format(encoder=encoder_name)
    args.benchmark_summary_json = "benchmark_summary.json"
    if args.eval_thresholds is None:
        coarse_low = np.arange(0.30, 0.55, 0.05)
        focused_mid = np.arange(0.55, 0.71, 0.01)
        coarse_high = np.arange(0.75, 0.81, 0.05)
        args.eval_thresholds = sorted(
            {round(float(v), 2) for v in np.concatenate([coarse_low, focused_mid, coarse_high])}
        )
    args.use_tta_eval = True
    args.tta_eval_mode_count = 4
    args.input_dir = args.input_dir.expanduser().resolve()
    args.aef_dir = args.aef_dir.expanduser().resolve()
    args.label_dir = args.label_dir.expanduser().resolve()
    if args.triplet_manifest is not None:
        args.triplet_manifest = args.triplet_manifest.expanduser().resolve()
    if args.fixed_split_manifest is not None:
        args.fixed_split_manifest = args.fixed_split_manifest.expanduser().resolve()
    args.artifact_dir = args.artifact_dir.expanduser().resolve()
    if args.warm_start_run_dir is not None:
        args.warm_start_run_dir = args.warm_start_run_dir.expanduser().resolve()
    if args.warm_start_checkpoint is not None:
        args.warm_start_checkpoint = args.warm_start_checkpoint.expanduser().resolve()
    if args.band_stats_path is not None:
        args.band_stats_path = args.band_stats_path.expanduser().resolve()
    args.hf_cache_dir = args.hf_cache_dir.expanduser().resolve()
    if args.benchmark_summary_path is not None:
        args.benchmark_summary_path = args.benchmark_summary_path.expanduser().resolve()
    if args.crop_size < 1:
        parser.error("--crop_size must be >= 1")
    if args.tile_overlap < 0 or args.tile_overlap >= args.crop_size:
        parser.error("--tile_overlap must be >= 0 and < --crop_size")
    if args.scale_jitter_min <= 0 or args.scale_jitter_max <= 0:
        parser.error("--scale_jitter_min and --scale_jitter_max must be > 0")
    if args.scale_jitter_min > args.scale_jitter_max:
        parser.error("--scale_jitter_min must be <= --scale_jitter_max")
    if args.batch_size < 1:
        parser.error("--batch_size must be >= 1")
    if args.grad_accum_steps < 1:
        parser.error("--grad_accum_steps must be >= 1")
    if args.target_source_fraction is not None and not 0 < args.target_source_fraction <= 1:
        parser.error("--target_source_fraction must be greater than 0 and at most 1")
    if not 0 <= args.positive_crop_p <= 1:
        parser.error("--positive_crop_p must be in [0, 1]")
    if not 0 <= args.min_water_fraction <= 1:
        parser.error("--min_water_fraction must be in [0, 1]")
    if args.positive_crop_attempts < 1:
        parser.error("--positive_crop_attempts must be >= 1")
    if args.loss_aux_weight is not None and args.loss_aux_weight < 0:
        parser.error("--loss_aux_weight must be non-negative")
    if args.tversky_alpha < 0 or args.tversky_beta < 0:
        parser.error("--tversky_alpha and --tversky_beta must be non-negative")
    if args.tversky_gamma <= 0:
        parser.error("--tversky_gamma must be > 0")
    if args.cldice_iterations < 1:
        parser.error("--cldice_iterations must be >= 1")
    if args.strategy_epochs < 1:
        parser.error("--strategy_epochs must be >= 1")
    if args.strategy_lr <= 0:
        parser.error("--strategy_lr must be > 0")
    if args.max_benchmark_batch_size is not None and args.max_benchmark_batch_size < 1:
        parser.error("--max_benchmark_batch_size must be >= 1")
    if args.worst_tile_count < 1:
        parser.error("--worst_tile_count must be >= 1")
    if args.triplet_sample_id and args.triplet_manifest is None:
        parser.error("--triplet_sample_id requires --triplet_manifest")
    return args


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def seed_worker(worker_id: Optional[int] = None) -> None:
    worker_seed = torch.initial_seed() % (2**32)
    random.seed(worker_seed)
    np.random.seed(worker_seed)


def setup_environment(args: argparse.Namespace) -> torch.device:
    seed_everything(args.seed)
    args.artifact_dir.mkdir(parents=True, exist_ok=True)
    (args.artifact_dir / "models").mkdir(parents=True, exist_ok=True)

    if torch.cuda.is_available():
        device_count = torch.cuda.device_count()
        actual_device = args.device_id if args.device_id < device_count else 0
        torch.cuda.set_device(actual_device)
        device = torch.device("cuda:{idx}".format(idx=actual_device))
        torch.backends.cudnn.benchmark = True
        print(
            {
                "device": str(device),
                "device_name": torch.cuda.get_device_name(actual_device),
                "device_count": device_count,
            }
        )
    else:
        device = torch.device("cpu")
        print({"device": "cpu"})

    n_input = args.n_s1_bands + args.n_aef_bands
    n_unet_input = args.n_s1_bands + (args.n_proj_bands if using_aef(args) else 0)
    if using_aef(args):
        dataloader_bands = "{s1} S1 + {aef} AEF = {n}".format(
            s1=args.n_s1_bands, aef=args.n_aef_bands, n=n_input
        )
        model_input_bands = "{s1} S1 + {proj} projected = {n}".format(
            s1=args.n_s1_bands, proj=args.n_proj_bands, n=n_unet_input
        )
    else:
        dataloader_bands = "{s1} S1 = {n}".format(s1=args.n_s1_bands, n=n_input)
        model_input_bands = "{s1} S1 = {n}".format(s1=args.n_s1_bands, n=n_unet_input)
    print(
        {
            "artifact_dir": str(args.artifact_dir),
            "mode": args.mode,
            "model_family": args.model_family,
            "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None,
            "dataloader_bands": dataloader_bands,
            "model_input_bands": model_input_bands,
            "crop_size": args.crop_size,
            "batch_size": args.batch_size,
            "num_workers": args.num_workers,
            "loss": args.loss,
            "tuning_strategy": args.tuning_strategy,
            "grad_accum_steps": args.grad_accum_steps,
            "selection_metric": args.selection_metric,
            "target_source_id": args.target_source_id,
            "target_source_fraction": args.target_source_fraction,
            "eval_thresholds": args.eval_thresholds,
            "force_tiled_validation": args.force_tiled_validation,
            "use_fp16": args.use_fp16,
            "pretrained": args.pretrained,
        }
    )
    return device


def build_aef_index(aef_dir: Path) -> Dict[str, Path]:
    if not aef_dir.exists():
        raise FileNotFoundError("AEF directory does not exist: {path}".format(path=aef_dir))
    index: Dict[str, Path] = {}
    duplicates: List[str] = []
    for path in sorted(aef_dir.rglob("*.tif")):
        rel_parts = path.relative_to(aef_dir).parts
        if any(part.startswith(".") for part in rel_parts):
            continue
        if path.name in index:
            duplicates.append(path.name)
        else:
            index[path.name] = path
    if duplicates:
        raise RuntimeError(
            "Found duplicate AEF filenames across recursive folders. Examples: {examples}".format(
                examples=duplicates[:5]
            )
        )
    if not index:
        raise FileNotFoundError("No AEF GeoTIFFs found under {path}".format(path=aef_dir))
    return index


def collect_triplets(
    input_dir: Path,
    aef_dir: Optional[Path],
    label_dir: Path,
) -> List[TripletRecord]:
    if not input_dir.exists():
        raise FileNotFoundError("Input directory does not exist: {path}".format(path=input_dir))
    if not label_dir.exists():
        raise FileNotFoundError("Label directory does not exist: {path}".format(path=label_dir))

    use_aef = aef_dir is not None
    aef_index = build_aef_index(aef_dir) if use_aef else {}
    # Build the label index from directory entries once. Calling Path.exists()
    # for every recovered-label symlink is extremely slow on Perlmutter's
    # scratch filesystem and previously consumed the entire smoke-job limit.
    label_index = {path.name: path for path in label_dir.glob("*.tif")}
    s1_paths = sorted(input_dir.glob("*.tif"))
    triplets: List[TripletRecord] = []
    missing_aef: List[str] = []
    missing_labels: List[str] = []

    for s1_path in s1_paths:
        aef_path = aef_index.get(s1_path.name) if use_aef else None
        label_path = label_index.get(s1_path.name)
        if use_aef and aef_path is None:
            missing_aef.append(s1_path.name)
            continue
        if label_path is None:
            missing_labels.append(s1_path.name)
            continue
        triplets.append(
            TripletRecord(
                name=s1_path.name,
                s1_path=s1_path,
                aef_path=aef_path,
                label_path=label_path,
                group_id=s1_path.name,
            )
        )

    if not triplets:
        raise FileNotFoundError(
            "No aligned S1/AEF/label triplets found. Missing AEF: {ma}, missing labels: {ml}".format(
                ma=len(missing_aef), ml=len(missing_labels)
            )
        )

    print(
        {
            "triplet_count": len(triplets),
            "missing_aef_count": len(missing_aef) if use_aef else 0,
            "missing_label_count": len(missing_labels),
            "missing_aef_examples": missing_aef[:3] if use_aef else [],
            "missing_label_examples": missing_labels[:3],
        }
    )
    return triplets


def collect_triplets_from_manifest(
    manifest_path: Path,
    sample_id: Optional[str],
    expect_aef: bool,
    expect_prior: bool = False,
) -> List[TripletRecord]:
    if not manifest_path.exists():
        raise FileNotFoundError("Triplet manifest does not exist: {path}".format(path=manifest_path))
    with manifest_path.open(newline="") as handle:
        rows = list(csv.DictReader(handle))
    required = {"s1_path", "label_path"}
    if expect_aef:
        required.add("aef_path")
    if expect_prior:
        required.add("prior_path")
    missing = required.difference(rows[0].keys() if rows else [])
    if missing:
        raise ValueError("Triplet manifest is missing columns: {cols}".format(cols=", ".join(sorted(missing))))

    manifest_sample_ids = sorted({row.get("sample_id", "") for row in rows if row.get("sample_id", "")})
    if sample_id is None and len(manifest_sample_ids) > 1:
        raise ValueError(
            "Triplet manifest has multiple sample_id values; pass --triplet_sample_id. Examples: {ids}".format(
                ids=", ".join(manifest_sample_ids[:5])
            )
        )
    selected_rows = [row for row in rows if sample_id is None or row.get("sample_id", "") == sample_id]
    if not selected_rows:
        raise ValueError("Triplet manifest has no rows for sample_id {sample}".format(sample=sample_id or "<all>"))

    records: List[TripletRecord] = []
    missing_paths: List[str] = []
    for row in selected_rows:
        name = row.get("tile_name") or row.get("name") or Path(row["s1_path"]).name
        s1_path = Path(row["s1_path"]).expanduser()
        label_path = Path(row["label_path"]).expanduser()
        aef_value = row.get("aef_path", "")
        aef_path = Path(aef_value).expanduser() if expect_aef and aef_value else None
        prior_value = row.get("prior_path", "")
        prior_path = Path(prior_value).expanduser() if expect_prior and prior_value else None
        path_checks = [("s1", s1_path), ("label", label_path)]
        if expect_aef:
            if aef_path is None:
                missing_paths.append("{name}: empty aef_path".format(name=name))
                continue
            path_checks.append(("aef", aef_path))
        if expect_prior:
            if prior_path is None:
                missing_paths.append("{name}: empty prior_path".format(name=name))
                continue
            path_checks.append(("prior", prior_path))
        missing_for_row = ["{kind}={path}".format(kind=kind, path=path) for kind, path in path_checks if not path.exists()]
        if missing_for_row:
            missing_paths.append("{name}: {paths}".format(name=name, paths=", ".join(missing_for_row)))
            continue
        records.append(
            TripletRecord(
                name=name,
                s1_path=s1_path.resolve(),
                aef_path=aef_path.resolve() if aef_path is not None else None,
                label_path=label_path.resolve(),
                sample_id=row.get("sample_id", sample_id or ""),
                group_id=row.get("grid_id") or row.get("group_id") or name,
                source_id=row.get("source_id", ""),
                prior_path=prior_path.resolve() if prior_path is not None else None,
            )
        )

    if missing_paths:
        raise FileNotFoundError(
            "Triplet manifest has missing or incomplete paths for {n} row(s). Examples: {examples}".format(
                n=len(missing_paths),
                examples="; ".join(missing_paths[:5]),
            )
        )
    if not records:
        raise FileNotFoundError("Triplet manifest did not produce any usable records: {}".format(manifest_path))
    print(
        {
            "triplet_manifest": str(manifest_path),
            "triplet_sample_id": sample_id or (manifest_sample_ids[0] if len(manifest_sample_ids) == 1 else None),
            "triplet_count": len(records),
        }
    )
    return sorted(records, key=lambda record: record.name)


def split_grouped_triplets(
    records: Sequence[TripletRecord],
    valid_pct: float,
    seed: int,
) -> Tuple[List[TripletRecord], List[TripletRecord]]:
    if not records:
        raise ValueError("Cannot split an empty record list.")
    grouped: Dict[str, List[TripletRecord]] = {}
    for record in records:
        grouped.setdefault(record.group_id or record.name, []).append(record)
    if len(grouped) < 2:
        raise ValueError("At least two independent groups are required for a train/validation split.")

    rng = random.Random(seed)
    shuffled_groups = list(grouped.items())
    rng.shuffle(shuffled_groups)
    target_valid_count = max(1, int(round(len(records) * valid_pct)))
    valid_group_ids = set()
    valid_count = 0
    for group_id, group_records in shuffled_groups:
        if valid_count >= target_valid_count:
            break
        valid_group_ids.add(group_id)
        valid_count += len(group_records)

    valid_records = sorted(
        [record for record in records if (record.group_id or record.name) in valid_group_ids],
        key=lambda r: r.name,
    )
    train_records = sorted(
        [record for record in records if (record.group_id or record.name) not in valid_group_ids],
        key=lambda r: r.name,
    )
    if not train_records:
        raise ValueError("Train split is empty; reduce --valid_pct.")
    train_group_ids = {record.group_id or record.name for record in train_records}
    if train_group_ids.intersection(valid_group_ids):
        raise RuntimeError("Group leakage detected between train and validation splits.")
    return train_records, valid_records


def split_triplets(records: Sequence[TripletRecord], valid_pct: float, seed: int) -> Tuple[List[TripletRecord], List[TripletRecord]]:
    source_ids = sorted({record.source_id for record in records if record.source_id})
    if len(source_ids) <= 1:
        return split_grouped_triplets(records, valid_pct=valid_pct, seed=seed)

    train_records: List[TripletRecord] = []
    valid_records: List[TripletRecord] = []
    for source_id in source_ids:
        source_records = [record for record in records if record.source_id == source_id]
        source_train, source_valid = split_grouped_triplets(
            source_records,
            valid_pct=valid_pct,
            seed=seed,
        )
        train_records.extend(source_train)
        valid_records.extend(source_valid)

    assigned = len(train_records) + len(valid_records)
    if assigned != len(records):
        raise RuntimeError(
            "Source-stratified split assigned {assigned} of {total} records.".format(
                assigned=assigned,
                total=len(records),
            )
        )
    return sorted(train_records, key=lambda r: r.name), sorted(valid_records, key=lambda r: r.name)


def split_triplets_from_manifest(
    records: Sequence[TripletRecord], manifest_path: Path
) -> Tuple[List[TripletRecord], List[TripletRecord]]:
    """Apply an immutable train/validation membership manifest by tile name."""
    if not manifest_path.exists():
        raise FileNotFoundError("Fixed split manifest does not exist: {path}".format(path=manifest_path))
    with manifest_path.open(newline="") as handle:
        rows = list(csv.DictReader(handle))
    required = {"split", "tile_name"}
    missing_columns = required.difference(rows[0].keys() if rows else [])
    if missing_columns:
        raise ValueError(
            "Fixed split manifest is missing columns: {columns}".format(
                columns=", ".join(sorted(missing_columns))
            )
        )

    assignments: Dict[str, str] = {}
    for row in rows:
        name = row["tile_name"].strip()
        split = row["split"].strip().lower()
        if split not in {"train", "valid", "validation"}:
            raise ValueError("Unknown split {split!r} for {name}".format(split=split, name=name))
        split = "valid" if split == "validation" else split
        if name in assignments:
            raise ValueError("Duplicate tile in fixed split manifest: {name}".format(name=name))
        assignments[name] = split

    records_by_name = {record.name: record for record in records}
    if len(records_by_name) != len(records):
        raise ValueError("Training records contain duplicate tile names; fixed split is ambiguous.")
    missing_records = sorted(set(assignments).difference(records_by_name))
    unassigned_records = sorted(set(records_by_name).difference(assignments))
    if missing_records or unassigned_records:
        raise ValueError(
            "Fixed split population mismatch: missing_records={missing}, unassigned_records={unassigned}, "
            "missing_examples={missing_examples}, unassigned_examples={unassigned_examples}".format(
                missing=len(missing_records),
                unassigned=len(unassigned_records),
                missing_examples=missing_records[:5],
                unassigned_examples=unassigned_records[:5],
            )
        )

    train_records = sorted(
        [records_by_name[name] for name, split in assignments.items() if split == "train"],
        key=lambda record: record.name,
    )
    valid_records = sorted(
        [records_by_name[name] for name, split in assignments.items() if split == "valid"],
        key=lambda record: record.name,
    )
    if not train_records or not valid_records:
        raise ValueError("Fixed split manifest must contain at least one train and one valid tile.")
    print(
        {
            "fixed_split_manifest": str(manifest_path),
            "fixed_train_tiles": len(train_records),
            "fixed_valid_tiles": len(valid_records),
        }
    )
    return train_records, valid_records


def limit_train_records(
    records: Sequence[TripletRecord],
    max_records: Optional[int],
    subset_seed: int,
) -> List[TripletRecord]:
    selected = list(records)
    if max_records is None:
        return selected
    if max_records <= 0:
        raise ValueError("--max_train_records must be positive when set.")
    rng = random.Random(subset_seed)
    shuffled = list(selected)
    rng.shuffle(shuffled)
    return sorted(shuffled[: min(max_records, len(shuffled))], key=lambda r: r.name)


def write_split_manifest(
    artifact_dir: Path,
    train_records: Sequence[TripletRecord],
    valid_records: Sequence[TripletRecord],
) -> Path:
    split_path = artifact_dir / "triplet_split_manifest.csv"
    fields = (
        "split",
        "sample_id",
        "source_id",
        "group_id",
        "name",
        "s1_path",
        "aef_path",
        "label_path",
        "prior_path",
    )
    rows = [("train", record) for record in train_records] + [("valid", record) for record in valid_records]
    with split_path.open("w", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for split, record in rows:
            writer.writerow(
                {
                    "split": split,
                    "sample_id": record.sample_id,
                    "source_id": record.source_id,
                    "group_id": record.group_id or record.name,
                    "name": record.name,
                    "s1_path": str(record.s1_path),
                    "aef_path": "" if record.aef_path is None else str(record.aef_path),
                    "label_path": str(record.label_path),
                    "prior_path": "" if record.prior_path is None else str(record.prior_path),
                }
            )
    print(
        {
            "triplet_split_manifest": str(split_path),
            "train_tiles": len(train_records),
            "valid_tiles": len(valid_records),
            "train_groups": len({record.group_id or record.name for record in train_records}),
            "valid_groups": len({record.group_id or record.name for record in valid_records}),
        }
    )
    return split_path


def subset_records(records: Sequence[TripletRecord], max_batches: Optional[int], batch_size: int) -> List[TripletRecord]:
    selected = list(records)
    if max_batches is None:
        return selected
    max_items = max(batch_size * max_batches, batch_size)
    max_items = min(len(selected), max_items)
    records_by_source: Dict[str, List[TripletRecord]] = {}
    for record in selected:
        records_by_source.setdefault(record.source_id or "unspecified", []).append(record)
    if len(records_by_source) <= 1:
        return selected[:max_items]

    # Small benchmark/smoke caps must not accidentally remove an entire label
    # source merely because the full split is sorted by tile name.
    source_ids = sorted(records_by_source)
    source_positions = {source_id: 0 for source_id in source_ids}
    source_stratified: List[TripletRecord] = []
    while len(source_stratified) < max_items:
        added_record = False
        for source_id in source_ids:
            source_records = records_by_source[source_id]
            source_position = source_positions[source_id]
            if source_position >= len(source_records):
                continue
            source_stratified.append(source_records[source_position])
            source_positions[source_id] += 1
            added_record = True
            if len(source_stratified) >= max_items:
                break
        if not added_record:
            break
    return source_stratified


def open_geotiff_float(path: Path) -> torch.Tensor:
    with rio.open(str(path)) as src:
        data = src.read().astype(np.float32)
    data = np.nan_to_num(data, nan=0.0, posinf=0.0, neginf=0.0)
    return torch.from_numpy(data)


def open_mask_binary(path: Path) -> torch.Tensor:
    """0 = dry, 1 = water, IGNORE_INDEX = Dynamic World did not observe."""
    with rio.open(str(path)) as src:
        data = src.read(1).astype(np.float32)
    return torch.from_numpy(binarize_label_array(data))


def open_prior_probability(path: Path) -> torch.Tensor:
    with rio.open(str(path)) as src:
        data = src.read(1).astype(np.float32)
        nodata = src.nodata
    invalid = ~np.isfinite(data)
    if nodata is not None:
        invalid |= data == nodata
    data = np.where(invalid, 0.0, data)
    data = np.clip(data, 0.0, 1.0)
    return torch.from_numpy(data).unsqueeze(0)


def triplet_shape(record: TripletRecord) -> Tuple[int, int]:
    with rio.open(str(record.s1_path)) as s1_src:
        height, width = s1_src.height, s1_src.width
    with rio.open(str(record.label_path)) as label_src:
        label_shape = (label_src.height, label_src.width)
    if label_shape != (height, width):
        raise ValueError(
            "Triplet {name} has mismatched S1/label shapes: s1={s1}, label={label}".format(
                name=record.name,
                s1=(height, width),
                label=label_shape,
            )
        )
    if record.aef_path is not None:
        with rio.open(str(record.aef_path)) as aef_src:
            aef_shape = (aef_src.height, aef_src.width)
        if aef_shape != (height, width):
            raise ValueError(
                "Triplet {name} has mismatched S1/AEF shapes: s1={s1}, aef={aef}".format(
                    name=record.name,
                    s1=(height, width),
                    aef=aef_shape,
                )
                )
    if record.prior_path is not None:
        with rio.open(str(record.prior_path)) as prior_src:
            prior_shape = (prior_src.height, prior_src.width)
            prior_count = prior_src.count
        if prior_count != 1 or prior_shape != (height, width):
            raise ValueError(
                "Triplet {name} has mismatched prior: s1={s1}, prior_shape={prior}, prior_bands={bands}".format(
                    name=record.name,
                    s1=(height, width),
                    prior=prior_shape,
                    bands=prior_count,
                )
            )
    return height, width


def load_triplet_window(
    record: TripletRecord,
    top: int,
    left: int,
    height: int,
    width: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Read one aligned source window instead of materializing a full 67-band tile."""
    window = rio.windows.Window(col_off=left, row_off=top, width=width, height=height)
    with rio.open(str(record.s1_path)) as src:
        s1_data = src.read(window=window).astype(np.float32)
    with rio.open(str(record.label_path)) as src:
        mask_data = src.read(1, window=window).astype(np.float32)
    s1_missing = ~np.isfinite(s1_data[:2]).all(axis=0) | (s1_data[:2] == 0).any(axis=0)   # full_coverage_v1 fix 2
    mask_data = np.where(s1_missing, float(IGNORE_INDEX), mask_data).astype(np.float32)

    arrays = [s1_data]
    if record.aef_path is not None:
        with rio.open(str(record.aef_path)) as src:
            arrays.append(src.read(window=window).astype(np.float32))
    if record.prior_path is not None:
        with rio.open(str(record.prior_path)) as src:
            prior_data = src.read(1, window=window).astype(np.float32)
            prior_nodata = src.nodata
        prior_invalid = ~np.isfinite(prior_data)
        if prior_nodata is not None:
            prior_invalid |= prior_data == prior_nodata
        prior_data = np.where(prior_invalid, 0.0, prior_data)
        arrays.append(np.clip(prior_data, 0.0, 1.0)[None, :, :])
    image_data = np.concatenate(arrays, axis=0) if len(arrays) > 1 else arrays[0]
    image_data = np.nan_to_num(image_data, nan=0.0, posinf=0.0, neginf=0.0)
    image = torch.from_numpy(image_data).float()
    mask = torch.from_numpy(binarize_label_array(mask_data)).long()
    return image, mask


def load_random_triplet_window(
    record: TripletRecord,
    crop_size: int,
    seek_positive: bool = False,
    min_water_fraction: float = 0.005,
    attempts: int = 8,
) -> Tuple[torch.Tensor, torch.Tensor]:
    height, width = triplet_shape(record)
    if height < crop_size or width < crop_size:
        raise ValueError(
            "Crop size {crop} exceeds tile shape ({height}, {width}) for {name}".format(
                crop=crop_size,
                height=height,
                width=width,
                name=record.name,
            )
        )
    top = random.randint(0, height - crop_size)
    left = random.randint(0, width - crop_size)
    if seek_positive:
        with rio.open(str(record.label_path)) as label_src:
            for _ in range(attempts):
                candidate_top = random.randint(0, height - crop_size)
                candidate_left = random.randint(0, width - crop_size)
                window = rio.windows.Window(
                    col_off=candidate_left,
                    row_off=candidate_top,
                    width=crop_size,
                    height=crop_size,
                )
                candidate_mask = binarize_label_array(label_src.read(1, window=window))
                top, left = candidate_top, candidate_left
                water_pixels = float(np.count_nonzero(candidate_mask == 1))
                if water_pixels / candidate_mask.size >= min_water_fraction:
                    break
    return load_triplet_window(record, top=top, left=left, height=crop_size, width=crop_size)


def load_center_triplet_window(record: TripletRecord, crop_size: int) -> Tuple[torch.Tensor, torch.Tensor]:
    height, width = triplet_shape(record)
    if height < crop_size or width < crop_size:
        raise ValueError(
            "Crop size {crop} exceeds tile shape ({height}, {width}) for {name}".format(
                crop=crop_size,
                height=height,
                width=width,
                name=record.name,
            )
        )
    top = (height - crop_size) // 2
    left = (width - crop_size) // 2
    return load_triplet_window(record, top=top, left=left, height=crop_size, width=crop_size)


def load_triplet_tensors(record: TripletRecord) -> Tuple[torch.Tensor, torch.Tensor]:
    s1 = open_geotiff_float(record.s1_path)
    mask = open_mask_binary(record.label_path)
    mask[(s1[:2] == 0).any(dim=0)] = IGNORE_INDEX   # full_coverage_v1 fix 2 (NaN already set to 0 by open_geotiff_float)
    if record.aef_path is None:
        if s1.shape[1:] != mask.shape:
            raise ValueError(
                "Triplet {name} has mismatched shapes: s1={s1}, mask={mask}".format(
                    name=record.name, s1=tuple(s1.shape), mask=tuple(mask.shape)
                )
            )
        image = s1
    else:
        aef = open_geotiff_float(record.aef_path)
        if s1.shape[1:] != aef.shape[1:] or s1.shape[1:] != mask.shape:
            raise ValueError(
                "Triplet {name} has mismatched shapes: s1={s1}, aef={aef}, mask={mask}".format(
                    name=record.name, s1=tuple(s1.shape), aef=tuple(aef.shape), mask=tuple(mask.shape)
                )
            )
        image = torch.cat([s1, aef], dim=0)
    if record.prior_path is not None:
        prior = open_prior_probability(record.prior_path)
        if prior.shape[1:] != mask.shape:
            raise ValueError(
                "Triplet {name} has mismatched prior/mask shapes: prior={prior}, mask={mask}".format(
                    name=record.name,
                    prior=tuple(prior.shape),
                    mask=tuple(mask.shape),
                )
            )
        image = torch.cat([image, prior], dim=0)
    return image.float(), mask.long()


def compute_band_stats(
    records: Sequence[TripletRecord],
    max_files: Optional[int],
    normalized_band_count: Optional[int] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    all_records = list(records)
    source_ids = sorted({record.source_id for record in all_records if record.source_id})
    if max_files is None or max_files >= len(all_records):
        selected = all_records
    elif len(source_ids) <= 1:
        selected = all_records[:max_files]
    else:
        selected = []
        base_count, remainder = divmod(max_files, len(source_ids))
        for source_idx, source_id in enumerate(source_ids):
            source_records = [record for record in all_records if record.source_id == source_id]
            source_limit = base_count + (1 if source_idx < remainder else 0)
            selected.extend(source_records[: min(source_limit, len(source_records))])
        if len(selected) < max_files:
            selected_names = {record.name for record in selected}
            selected.extend(
                record for record in all_records if record.name not in selected_names
            )
            selected = selected[:max_files]
    if not selected:
        raise ValueError("Cannot compute stats without at least one training record.")

    print(
        {
            "band_stats_record_count": len(selected),
            "band_stats_source_counts": dict(Counter(record.source_id or "unspecified" for record in selected)),
        }
    )

    channel_sum = None
    channel_sumsq = None
    pixel_count = 0

    for record in selected:
        image, _ = load_triplet_tensors(record)
        if normalized_band_count is not None:
            image = image[:normalized_band_count]
        flat = image.numpy().reshape(image.shape[0], -1)
        if channel_sum is None:
            channel_sum = flat.sum(axis=1)
            channel_sumsq = np.square(flat).sum(axis=1)
        else:
            channel_sum += flat.sum(axis=1)
            channel_sumsq += np.square(flat).sum(axis=1)
        pixel_count += flat.shape[1]

    mean = channel_sum / pixel_count
    variance = np.maximum((channel_sumsq / pixel_count) - np.square(mean), 1e-8)
    std = np.sqrt(variance)
    return mean.astype(np.float32), std.astype(np.float32)


def get_band_stats(
    train_records: Sequence[TripletRecord],
    artifact_dir: Path,
    stats_sample_size: int,
    source_stats_path: Optional[Path] = None,
    expected_bands: Optional[int] = None,
) -> Tuple[np.ndarray, np.ndarray]:
    stats_cache = artifact_dir / "band_stats.npz"
    if source_stats_path is not None:
        if not source_stats_path.exists():
            raise FileNotFoundError("Band stats source does not exist: {path}".format(path=source_stats_path))
        if source_stats_path.resolve() != stats_cache.resolve():
            shutil.copy2(source_stats_path, stats_cache)
        save_json(
            artifact_dir / "band_stats_source.json",
            {"source_band_stats": str(source_stats_path.resolve()), "artifact_band_stats": str(stats_cache)},
        )
        print("Copied band stats from {source} to {dest}".format(source=source_stats_path, dest=stats_cache))
    if stats_cache.exists():
        cached = np.load(stats_cache)
        validate_band_stats(cached["mean"], cached["std"], expected_bands, stats_cache)
        print("Loaded cached band stats from {path}".format(path=stats_cache))
        return cached["mean"], cached["std"]
    band_mean, band_std = compute_band_stats(
        train_records,
        max_files=stats_sample_size,
        normalized_band_count=expected_bands,
    )
    validate_band_stats(band_mean, band_std, expected_bands, stats_cache)
    np.savez(stats_cache, mean=band_mean, std=band_std)
    print("Saved band stats to {path}".format(path=stats_cache))
    return band_mean, band_std


def validate_band_stats(mean: np.ndarray, std: np.ndarray, expected_bands: Optional[int], path: Path) -> None:
    if mean.shape != std.shape:
        raise ValueError("Band stats mean/std shape mismatch at {path}: {mean} vs {std}".format(path=path, mean=mean.shape, std=std.shape))
    if expected_bands is not None and len(mean) != expected_bands:
        raise ValueError(
            "Band stats at {path} have {actual} bands; expected {expected}".format(
                path=path,
                actual=len(mean),
                expected=expected_bands,
            )
        )


def normalize_image(image: torch.Tensor, band_mean: torch.Tensor, band_std: torch.Tensor) -> torch.Tensor:
    normalized_band_count = int(band_mean.shape[0])
    if image.shape[0] < normalized_band_count:
        raise ValueError(
            "Image has {actual} bands but normalization expects {expected}.".format(
                actual=image.shape[0],
                expected=normalized_band_count,
            )
        )
    normalized = image.clone()
    normalized[:normalized_band_count] = (
        image[:normalized_band_count] - band_mean
    ) / band_std
    return normalized


def random_crop(image: torch.Tensor, mask: torch.Tensor, crop_size: int) -> Tuple[torch.Tensor, torch.Tensor]:
    _, height, width = image.shape
    if height < crop_size or width < crop_size:
        raise ValueError(
            "Crop size {crop} exceeds tile shape ({h}, {w})".format(
                crop=crop_size, h=height, w=width
            )
        )
    top = random.randint(0, height - crop_size)
    left = random.randint(0, width - crop_size)
    return image[:, top : top + crop_size, left : left + crop_size], mask[top : top + crop_size, left : left + crop_size]


def apply_train_augmentations(
    image: torch.Tensor,
    mask: torch.Tensor,
    crop_size: int,
    n_s1_bands: int,
    rotate_max_deg: float,
    rotate_p: float,
    scale_jitter_min: float,
    scale_jitter_max: float,
    s1_gain_jitter: float,
    s1_bias_jitter: float,
    gaussian_noise_std: float,
    gaussian_noise_p: float,
    normalized_band_count: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if random.random() < 0.5:
        image = torch.flip(image, dims=(-1,))
        mask = torch.flip(mask, dims=(-1,))

    rotations = random.randint(0, 3)
    if rotations:
        image = torch.rot90(image, rotations, dims=(-2, -1))
        mask = torch.rot90(mask, rotations, dims=(-2, -1))

    if scale_jitter_min != 1.0 or scale_jitter_max != 1.0:
        scale = random.uniform(scale_jitter_min, scale_jitter_max)
        scaled_size = max(32, int(round(crop_size * scale)))
        if scaled_size != crop_size:
            image = F.interpolate(
                image.unsqueeze(0),
                size=(scaled_size, scaled_size),
                mode="bilinear",
                align_corners=False,
            ).squeeze(0)
            mask = F.interpolate(
                mask.unsqueeze(0).unsqueeze(0).float(),
                size=(scaled_size, scaled_size),
                mode="nearest",
            ).squeeze(0).squeeze(0)
            if scaled_size > crop_size:
                image, mask = random_crop(image, mask, crop_size)
            else:
                pad_h = crop_size - scaled_size
                pad_w = crop_size - scaled_size
                top = random.randint(0, pad_h)
                left = random.randint(0, pad_w)
                bottom = pad_h - top
                right = pad_w - left
                image = F.pad(image, (left, right, top, bottom), mode="reflect")
                mask = F.pad(mask, (left, right, top, bottom), mode="constant", value=0)

    if rotate_max_deg > 0 and random.random() < rotate_p:
        angle = random.uniform(-rotate_max_deg, rotate_max_deg)
        image = TVF.rotate(image, angle=angle, interpolation=InterpolationMode.BILINEAR)
        mask = TVF.rotate(
            mask.unsqueeze(0).float(),
            angle=angle,
            interpolation=InterpolationMode.NEAREST,
        ).squeeze(0)

    if s1_gain_jitter > 0 or s1_bias_jitter > 0:
        gain = 1.0 + random.uniform(-s1_gain_jitter, s1_gain_jitter)
        bias = random.uniform(-s1_bias_jitter, s1_bias_jitter)
        image[:n_s1_bands] = (image[:n_s1_bands] * gain) + bias

    if gaussian_noise_std > 0 and random.random() < gaussian_noise_p:
        image[:normalized_band_count] = image[:normalized_band_count] + (
            torch.randn_like(image[:normalized_band_count]) * gaussian_noise_std
        )

    return image.float(), mask.long()


def build_dataloaders(
    args: argparse.Namespace,
    train_records: Sequence[TripletRecord],
    valid_records: Sequence[TripletRecord],
    band_mean: np.ndarray,
    band_std: np.ndarray,
    device: torch.device,
) -> DataLoaders:
    train_dataset = S1AEFCropDataset(
        records=train_records,
        crop_size=args.crop_size,
        band_mean=band_mean,
        band_std=band_std,
        n_s1_bands=args.n_s1_bands,
        rotate_max_deg=args.rotate_max_deg,
        rotate_p=args.rotate_p,
        scale_jitter_min=args.scale_jitter_min,
        scale_jitter_max=args.scale_jitter_max,
        s1_gain_jitter=args.s1_gain_jitter,
        s1_bias_jitter=args.s1_bias_jitter,
        gaussian_noise_std=args.gaussian_noise_std,
        gaussian_noise_p=args.gaussian_noise_p,
        positive_crop_p=args.positive_crop_p,
        min_water_fraction=args.min_water_fraction,
        positive_crop_attempts=args.positive_crop_attempts,
    )
    valid_dataset = S1AEFValidationPatchDataset(
        valid_records,
        crop_size=args.crop_size,
        band_mean=band_mean,
        band_std=band_std,
    )

    if args.train_tile_index is not None:   # full_coverage_v1
        train_dataset = S1AEFTileDataset(
            read_tile_index(args.train_tile_index, train_records, "train"), crop_size=args.crop_size, band_mean=band_mean,
            band_std=band_std, n_s1_bands=args.n_s1_bands, rotate_max_deg=args.rotate_max_deg, rotate_p=args.rotate_p,
            scale_jitter_min=args.scale_jitter_min, scale_jitter_max=args.scale_jitter_max, s1_gain_jitter=args.s1_gain_jitter,
            s1_bias_jitter=args.s1_bias_jitter, gaussian_noise_std=args.gaussian_noise_std, gaussian_noise_p=args.gaussian_noise_p)
        print({"full_coverage": "train", "tiles": len(train_dataset), "chips": len(train_records),
               "tiles_by_source": dict(Counter(tile[0].source_id for tile in train_dataset.tiles))})
    if args.valid_tile_index is not None:   # full_coverage_v1
        valid_dataset = S1AEFValidationTileDataset(
            read_tile_index(args.valid_tile_index, valid_records, "valid"), valid_records, crop_size=args.crop_size,
            band_mean=band_mean, band_std=band_std)
        print({"full_coverage": "valid", "tiles": len(valid_dataset), "chips": len(valid_records),
               "tiles_by_source": dict(Counter(valid_dataset.item_source_ids))})
    pin_memory = device.type == "cuda" and not args.disable_pin_memory
    train_loader_kwargs = {
        "dataset": train_dataset,
        "bs": args.batch_size,
        "shuffle": True,
        "num_workers": args.num_workers,
        "pin_memory": pin_memory,
        "persistent_workers": args.num_workers > 0,
        "device": None,
        "wif": seed_worker,
    }
    if args.target_source_fraction is not None:
        target_indices = [
            idx for idx, record in enumerate(train_records) if record.source_id == args.target_source_id
        ]
        other_indices = [
            idx for idx, record in enumerate(train_records) if record.source_id != args.target_source_id
        ]
        if not target_indices or not other_indices:
            raise ValueError(
                "Source-aware sampling requires both target source {target} and non-target records. "
                "Observed counts: {counts}".format(
                    target=args.target_source_id,
                    counts=dict(Counter(record.source_id or "unspecified" for record in train_records)),
                )
            )
        target_weight = args.target_source_fraction / len(target_indices)
        other_weight = (1.0 - args.target_source_fraction) / len(other_indices)
        record_weights = [
            target_weight if record.source_id == args.target_source_id else other_weight
            for record in train_records
        ]

        source_sampler_rng = random.Random(args.seed)

        def source_weighted_indices():
            return source_sampler_rng.choices(
                list(range(len(record_weights))),
                weights=record_weights,
                k=len(record_weights),
            )

        train_loader_kwargs["shuffle"] = False
        train_loader_kwargs["get_idxs"] = source_weighted_indices
        print(
            {
                "source_sampling": "weighted_with_replacement",
                "target_source_id": args.target_source_id,
                "target_source_fraction": args.target_source_fraction,
                "epoch_samples": len(train_dataset),
                "source_record_counts": dict(
                    Counter(record.source_id or "unspecified" for record in train_records)
                ),
            }
        )
    else:
        print(
            {
                "source_sampling": "natural_without_replacement",
                "epoch_samples": len(train_dataset),
                "source_record_counts": dict(
                    Counter(record.source_id or "unspecified" for record in train_records)
                ),
            }
        )
    train_dl = DataLoader(**train_loader_kwargs)
    valid_dl = DataLoader(
        valid_dataset,
        bs=1,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=pin_memory,
        persistent_workers=args.num_workers > 0,
        device=None,
        wif=seed_worker,
    )
    dls = DataLoaders(train_dl, valid_dl, path=args.artifact_dir, device=device)
    dls.c = len(args.codes)
    dls.vocab = args.codes
    return dls


def format_lr(lr_max):
    if isinstance(lr_max, slice):
        return {"low": lr_max.start, "high": lr_max.stop}
    return lr_max


class GlobalBestSaveModelCallback(SaveModelCallback):
    """full_coverage_v1: best checkpoint over all training stages (the stock callback restarts its best at each stage)."""

    def before_fit(self):
        super().before_fit()
        best = getattr(self.learn, "global_best_metric", None)
        if best is not None:
            self.best = best

    def after_epoch(self):
        super().after_epoch()
        if self.new_best:
            self.learn.global_best_metric = self.best
            self.learn.global_best_where = {"phase": getattr(self.learn, "training_phase", None), "epoch_in_phase": int(self.epoch),
                                            "value": float(self.best)}
            print({"global_best": self.learn.global_best_where})


def build_callbacks(args: argparse.Namespace, epoch_timer: Optional[EpochTimerCallback] = None) -> List[Callback]:
    callbacks: List[Callback] = [
        GradientClip(args.grad_clip),
        NonFiniteLossCallback(),
        GlobalBestSaveModelCallback(
            monitor=args.selection_metric,
            comp=np.greater,
            min_delta=args.min_delta,
            fname=args.best_model_name,
        ),
        EarlyStoppingCallback(
            monitor=args.selection_metric,
            comp=np.greater,
            min_delta=args.min_delta,
            patience=args.patience,
        ),
    ]
    if args.grad_accum_steps > 1:
        callbacks.append(GradientAccumulation(n_acc=args.grad_accum_steps))
    if epoch_timer is not None:
        callbacks.append(epoch_timer)
    return callbacks


def recorder_history_to_df(learn: Learner, phase: str) -> pd.DataFrame:
    values = list(learn.recorder.values)
    if not values:
        return pd.DataFrame(columns=["phase"])

    column_count = len(values[0])
    metric_names = list(learn.recorder.metric_names[1:])
    if metric_names and metric_names[-1] == "time" and len(metric_names) == column_count + 1:
        metric_names = metric_names[:-1]
    if len(metric_names) != column_count:
        metric_names = ["value_{idx}".format(idx=idx) for idx in range(column_count)]
    return pd.DataFrame(values, columns=metric_names).assign(phase=phase)


def run_phase(
    learn: Learner,
    args: argparse.Namespace,
    phase: str,
    epochs: int,
    lr_max,
    epoch_timer: Optional[EpochTimerCallback] = None,
) -> pd.DataFrame:
    if epochs <= 0:
        return pd.DataFrame(columns=["phase"])

    print({"phase": phase, "epochs": epochs, "lr_max": format_lr(lr_max)})
    learn.training_phase = phase
    try:
        fit_kwargs = {
            "wd": args.weight_decay,
            "cbs": build_callbacks(args, epoch_timer=epoch_timer),
            "reset_opt": True,
        }
        if args.schedule_policy == "flat_cos":
            learn.fit_flat_cos(epochs, lr=lr_max, **fit_kwargs)
        else:
            learn.fit_one_cycle(epochs, lr_max=lr_max, **fit_kwargs)
    finally:
        learn.training_phase = None
    return recorder_history_to_df(learn, phase=phase)


def build_schedule(args: argparse.Namespace) -> List[Tuple[str, int, object]]:
    if args.mode == "benchmark":
        return [("freeze", 2, args.warmup_lr)]
    if args.tuning_strategy != "staged":
        return [(args.tuning_strategy, args.strategy_epochs, args.strategy_lr)]
    schedule: List[Tuple[str, int, object]] = [
        ("freeze", args.freeze_epochs, args.warmup_lr),
    ]
    bridge_epochs = min(args.partial_unfreeze_epochs, args.finetune_epochs)
    remaining_finetune_epochs = max(args.finetune_epochs - bridge_epochs, 0)
    if bridge_epochs > 0:
        schedule.append(
            (
                "partial_unfreeze",
                bridge_epochs,
                slice(args.partial_unfreeze_lr / 20, args.partial_unfreeze_lr),
            )
        )
    if remaining_finetune_epochs > 0:
        schedule.append(
            (
                "finetune",
                remaining_finetune_epochs,
                slice(args.unfreeze_lr / 50, args.unfreeze_lr),
            )
        )
    if args.finetune_tail_epochs > 0:
        schedule.append(
            (
                "finetune_tail",
                args.finetune_tail_epochs,
                slice(args.finetune_tail_lr / 10, args.finetune_tail_lr),
            )
        )
    return schedule


def set_requires_grad(module: nn.Module, trainable: bool) -> None:
    for parameter in module.parameters():
        parameter.requires_grad = trainable


def trainable_parameter_summary(model: nn.Module) -> Dict[str, int]:
    trainable = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
    total = sum(parameter.numel() for parameter in model.parameters())
    return {"trainable_parameters": int(trainable), "total_parameters": int(total)}


def resnet_unet_children(model: nn.Module) -> List[nn.Module]:
    if not isinstance(model, AEFBottleneckUNet):
        raise ValueError("Explicit tuning strategies currently require AEFBottleneckUNet.")
    children = list(model.unet.layers) if hasattr(model.unet, "layers") else list(model.unet.children())
    if len(children) < 3:
        raise ValueError("Unexpected DynamicUnet layout for tuning strategy.")
    return children


def enable_head_only(model: nn.Module) -> None:
    children = resnet_unet_children(model)
    parameterized_children = [module for module in children if any(True for _ in module.parameters())]
    if not parameterized_children:
        raise ValueError("DynamicUnet head did not expose trainable parameters.")
    head = parameterized_children[-1]
    set_requires_grad(head, True)


def enable_decoder(model: nn.Module) -> None:
    children = resnet_unet_children(model)
    for module in children[1:]:
        set_requires_grad(module, True)


def enable_aef_projection(model: nn.Module) -> None:
    if not isinstance(model, AEFBottleneckUNet):
        raise ValueError("AEF projection strategy requires AEFBottleneckUNet.")
    set_requires_grad(model.aef_proj, True)


def enable_last_encoder_stage(model: nn.Module) -> None:
    encoder = resnet_unet_children(model)[0]
    encoder_children = list(encoder.children())
    if not encoder_children:
        raise ValueError("DynamicUnet encoder did not expose child stages.")
    set_requires_grad(encoder_children[-1], True)


def apply_explicit_tuning_strategy(learn: Learner, strategy: str) -> None:
    set_requires_grad(learn.model, False)
    if strategy == "head_only":
        enable_head_only(learn.model)
    elif strategy == "decoder_aef":
        enable_decoder(learn.model)
        enable_aef_projection(learn.model)
    elif strategy == "last_encoder_decoder_aef":
        enable_decoder(learn.model)
        enable_aef_projection(learn.model)
        enable_last_encoder_stage(learn.model)
    elif strategy == "full_low_lr":
        set_requires_grad(learn.model, True)
    else:
        raise ValueError("Unknown explicit tuning strategy: {}".format(strategy))
    summary = trainable_parameter_summary(learn.model)
    if summary["trainable_parameters"] == 0:
        raise ValueError("Tuning strategy {} left no trainable parameters.".format(strategy))
    print({"tuning_strategy_scope": strategy, **summary})


def apply_phase_freeze(learn: Learner, phase: str) -> None:
    if phase in {"head_only", "decoder_aef", "last_encoder_decoder_aef", "full_low_lr"}:
        apply_explicit_tuning_strategy(learn, phase)
        return
    if phase == "freeze":
        learn.freeze()
    elif phase == "partial_unfreeze":
        learn.freeze_to(-2)
    else:
        learn.unfreeze()


def build_wrapped_splitter(base_splitter):
    def wrapped_splitter(model: AEFBottleneckUNet):
        groups = list(base_splitter(model.unet))
        if not groups:
            return L([L(model.parameters())])
        final_group = L(groups[-1]) + L(model.aef_proj.parameters())
        return L(groups[:-1]) + L([final_group])

    return wrapped_splitter


def build_segformer_splitter():
    def splitter(model: AEFBottleneckSegFormer):
        encoder = model.segformer.segformer.encoder
        groups: List[L] = []
        for patch_embedding, blocks, norm in zip(
            encoder.patch_embeddings,
            encoder.block,
            encoder.layer_norm,
        ):
            groups.append(
                L(patch_embedding.parameters()) + L(blocks.parameters()) + L(norm.parameters())
            )
        groups.append(L(model.segformer.decode_head.parameters()) + L(model.aef_proj.parameters()))
        return L(groups)

    return splitter


def build_loss(args: argparse.Namespace) -> nn.Module:
    loss_name = args.loss
    if loss_name == "focal_dice":
        return CombinedLoss()
    if loss_name == "ce_dice":
        return CrossEntropyDiceLoss()
    if loss_name == "lovasz":
        return LovaszSoftmaxLoss()
    if loss_name == "ce_soft_iou":
        weight = 1.0 if args.loss_aux_weight is None else args.loss_aux_weight
        return CrossEntropySoftIoULoss(aux_weight=weight)
    if loss_name == "focal_tversky":
        weight = 1.0 if args.loss_aux_weight is None else args.loss_aux_weight
        return FocalTverskyCELoss(
            alpha=args.tversky_alpha,
            beta=args.tversky_beta,
            gamma=args.tversky_gamma,
            aux_weight=weight,
        )
    if loss_name == "ce_dice_boundary":
        weight = 0.2 if args.loss_aux_weight is None else args.loss_aux_weight
        return BoundaryAwareCEDiceLoss(aux_weight=weight)
    if loss_name == "ce_dice_cldice":
        weight = 0.2 if args.loss_aux_weight is None else args.loss_aux_weight
        return TopologyAwareCEDiceLoss(
            aux_weight=weight,
            skeleton_iterations=args.cldice_iterations,
        )
    if loss_name == "sce_dice":
        weight = 0.1 if args.loss_aux_weight is None else args.loss_aux_weight
        return SymmetricCrossEntropyDiceLoss(aux_weight=weight)
    raise ValueError("Unknown loss: {name}".format(name=loss_name))


def build_optimizer(args: argparse.Namespace):
    if args.optimizer == "ranger":
        return ranger
    if args.optimizer == "adamw":
        return Adam
    raise ValueError("Unknown optimizer: {name}".format(name=args.optimizer))


def build_learner(
    args: argparse.Namespace,
    dls: DataLoaders,
    valid_records: Sequence[TripletRecord],
) -> Learner:
    n_prior_bands = 1 if using_prior(args) else 0
    n_total_in = args.n_s1_bands + args.n_aef_bands + n_prior_bands
    n_model_in = (
        args.n_s1_bands
        + (args.n_proj_bands if using_aef(args) else 0)
        + n_prior_bands
    )
    loss_func = build_loss(args)
    opt_func = build_optimizer(args)
    metrics = [
        MaskedJaccardCoeff(),
        MaskedDice(),
        WaterIoU(),
        SourceBalancedWaterIoU(getattr(dls.valid.dataset, "item_source_ids", [record.source_id for record in valid_records])),
        TargetSourceWaterIoU(
            getattr(dls.valid.dataset, "item_source_ids", [record.source_id for record in valid_records]),
            target_source_id=args.target_source_id,
        ),
    ]
    if args.model_family == "segformer":
        wrapped_model = AEFBottleneckSegFormer(
            model_name=SEGFORMER_MODEL_MAP[args.segformer_variant],
            num_labels=len(args.codes),
            n_s1=args.n_s1_bands,
            n_aef=args.n_aef_bands,
            n_proj=args.n_proj_bands,
            hf_cache_dir=args.hf_cache_dir,
            pretrained=args.pretrained,
            local_files_only=args.hf_local_files_only,
        )
        wrapped_model = wrapped_model.to(dls.device)
        learn = Learner(
            dls,
            wrapped_model,
            loss_func=loss_func,
            metrics=metrics,
            opt_func=opt_func,
            splitter=build_segformer_splitter(),
            path=args.artifact_dir,
            model_dir="models",
        )
        model_name = "AEFBottleneckSegFormer({variant})".format(
            variant=args.segformer_variant
        )
    else:
        encoder_arches = {
            "resnet34_unet": resnet34,
            "resnet50_unet": resnet50,
            "resnet101_unet": resnet101,
        }
        encoder_arch = encoder_arches[args.model_family]
        learn = unet_learner(
            dls,
            encoder_arch,
            normalize=False,
            n_in=n_model_in,
            n_out=len(args.codes),
            metrics=metrics,
            loss_func=loss_func,
            opt_func=opt_func,
            act_cls=Mish,
            path=args.artifact_dir,
            model_dir="models",
            pretrained=args.pretrained,
        )
        if using_aef(args):
            base_splitter = learn.splitter
            wrapped_model = AEFBottleneckUNet(
                unet_model=learn.model,
                n_s1=args.n_s1_bands,
                n_aef=args.n_aef_bands,
                n_proj=args.n_proj_bands,
                n_prior=n_prior_bands,
                prior_dropout_p=args.prior_dropout_p,
                aef_dropout_p=getattr(args, "aef_dropout_p", 0.0),
            )
            learn.model = wrapped_model.to(dls.device)
            learn.splitter = build_wrapped_splitter(base_splitter)
            model_name = "AEFBottleneckUNet({family})".format(family=args.model_family)
        else:
            model_name = "S1UNet({family})".format(family=args.model_family)
        learn.model = learn.model.to(dls.device)
    if args.use_fp16 and torch.cuda.is_available():
        learn = learn.to_fp16()

    model_summary = {
        "model": model_name,
        "dataloader_input_bands": n_total_in,
        "backbone_input_bands": n_model_in,
        "loss": args.loss,
        "loss_aux_weight": args.loss_aux_weight,
        "optimizer": args.optimizer,
        "schedule_policy": args.schedule_policy,
        "selection_metric": args.selection_metric,
        "tuning_strategy": args.tuning_strategy,
        "grad_accum_steps": args.grad_accum_steps,
        "use_permanent_prior": using_prior(args),
        "prior_dropout_p": args.prior_dropout_p,
        "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0),
        "precision": "fp16" if args.use_fp16 and torch.cuda.is_available() else "fp32",
        "best_model_path": str(args.artifact_dir / "models" / "{name}.pth".format(name=args.best_model_name)),
    }
    if using_aef(args):
        model_summary["bottleneck"] = "{aef}->{proj}".format(
            aef=args.n_aef_bands, proj=args.n_proj_bands
        )
    print(
        model_summary
    )
    return learn


def resolve_warm_start_checkpoint(args: argparse.Namespace) -> Optional[Path]:
    if args.warm_start_checkpoint is not None:
        return args.warm_start_checkpoint
    if args.warm_start_run_dir is None:
        return None
    return args.warm_start_run_dir / "models" / "{name}.pth".format(name=args.best_model_name)


def resolve_band_stats_source(args: argparse.Namespace) -> Optional[Path]:
    if args.band_stats_path is not None:
        return args.band_stats_path
    if args.warm_start_run_dir is not None:
        return args.warm_start_run_dir / "band_stats.npz"
    return None


def unwrap_state_dict(checkpoint: object) -> Dict[str, torch.Tensor]:
    if not isinstance(checkpoint, dict):
        raise ValueError("Warm-start checkpoint must contain a state dict mapping.")
    for key in ("model", "state_dict"):
        value = checkpoint.get(key)
        if isinstance(value, dict):
            checkpoint = value
            break
    if not checkpoint or not all(isinstance(key, str) for key in checkpoint.keys()):
        raise ValueError("Warm-start checkpoint does not expose string state-dict keys.")
    return checkpoint


def validate_warm_start_projection(
    state_dict: Dict[str, torch.Tensor],
    model: nn.Module,
    args: argparse.Namespace,
    checkpoint_path: Path,
) -> None:
    if not using_aef(args):
        return
    key = "aef_proj.weight"
    if key not in state_dict:
        raise ValueError("Warm-start checkpoint is missing {key}: {path}".format(key=key, path=checkpoint_path))
    checkpoint_shape = tuple(state_dict[key].shape)
    model_shape = tuple(model.state_dict()[key].shape)
    expected_shape = (args.n_proj_bands, args.n_aef_bands, 1, 1)
    if checkpoint_shape != expected_shape or model_shape != expected_shape:
        raise ValueError(
            "Warm-start AEF projection shape mismatch at {path}: checkpoint={ckpt}, model={model}, expected={expected}".format(
                path=checkpoint_path,
                ckpt=checkpoint_shape,
                model=model_shape,
                expected=expected_shape,
            )
        )


def maybe_warm_start(learn: Learner, args: argparse.Namespace) -> Optional[Path]:
    checkpoint_path = resolve_warm_start_checkpoint(args)
    if checkpoint_path is None:
        return None
    if not checkpoint_path.exists():
        raise FileNotFoundError("Warm-start checkpoint does not exist: {path}".format(path=checkpoint_path))
    state_dict = unwrap_state_dict(torch.load(checkpoint_path, map_location="cpu"))
    validate_warm_start_projection(state_dict, learn.model, args, checkpoint_path)
    try:
        learn.model.load_state_dict(state_dict, strict=True)
    except RuntimeError as exc:
        raise RuntimeError("Warm-start checkpoint is incompatible with this model: {path}\n{error}".format(path=checkpoint_path, error=exc)) from exc
    save_json(
        args.artifact_dir / "warm_start.json",
        {
            "checkpoint": str(checkpoint_path.resolve()),
            "parent_run_dir": "" if args.warm_start_run_dir is None else str(args.warm_start_run_dir),
            "n_aef_bands": args.n_aef_bands,
            "n_proj_bands": args.n_proj_bands,
        },
    )
    print({"warm_start_checkpoint": str(checkpoint_path), "warm_start_strict": True})
    return checkpoint_path


def save_json(path: Path, payload: dict) -> None:
    path.write_text(json.dumps(payload, indent=2))


def inspect_bottleneck(model: nn.Module, artifact_dir: Path) -> dict:
    weight = model.aef_proj.weight.detach().cpu().squeeze(-1).squeeze(-1)
    abs_weight = weight.abs()
    top_k = 5
    report = {
        "projection_shape": list(weight.shape),
        "weight_norm_per_proj_band": [float(v) for v in weight.norm(dim=1).tolist()],
        "weight_norm_per_aef_band": [float(v) for v in weight.norm(dim=0).tolist()],
    }
    for idx in range(weight.shape[0]):
        top_vals, top_idx = abs_weight[idx].topk(top_k)
        report["proj_band_{idx}_top{top}_aef_bands".format(idx=idx, top=top_k)] = top_idx.tolist()
        report["proj_band_{idx}_top{top}_weights".format(idx=idx, top=top_k)] = [
            round(v, 6) for v in top_vals.tolist()
        ]

    np.save(artifact_dir / "aef_bottleneck_weights.npy", weight.numpy())
    save_json(artifact_dir / "aef_bottleneck_report.json", report)
    print(
        {
            "bottleneck_weight_file": str(artifact_dir / "aef_bottleneck_weights.npy"),
            "bottleneck_report_file": str(artifact_dir / "aef_bottleneck_report.json"),
        }
    )
    return report


def get_tta_transforms(mode_count: int):
    transforms = [
        ("identity", lambda x: x),
        ("flip_horizontal", lambda x: torch.flip(x, dims=(-1,))),
        ("flip_vertical", lambda x: torch.flip(x, dims=(-2,))),
        ("flip_both", lambda x: torch.flip(x, dims=(-2, -1))),
    ]
    return transforms[: max(1, min(mode_count, len(transforms)))]


def is_cuda_oom(error: RuntimeError) -> bool:
    message = str(error).lower()
    return "out of memory" in message or "cuda error: out of memory" in message


def sliding_positions(size: int, tile_size: int, overlap: int) -> List[int]:
    if size <= tile_size:
        return [0]
    stride = max(1, tile_size - overlap)
    positions = list(range(0, size - tile_size + 1, stride))
    if positions[-1] != size - tile_size:
        positions.append(size - tile_size)
    return positions


def predict_record_logits_windowed(
    model: nn.Module,
    record: TripletRecord,
    band_mean: torch.Tensor,
    band_std: torch.Tensor,
    tile_size: int,
    overlap: int,
    device: torch.device,
    use_tta: bool = False,
    tta_mode_count: int = 4,
) -> torch.Tensor:
    """Blend window predictions on CPU while only one input patch resides on the GPU."""
    height, width = triplet_shape(record)
    y_positions = sliding_positions(height, tile_size, overlap)
    x_positions = sliding_positions(width, tile_size, overlap)
    logits_sum = None
    weight_sum = torch.zeros((height, width), dtype=torch.float32)
    transforms = get_tta_transforms(tta_mode_count) if use_tta else [("identity", lambda x: x)]

    for top in y_positions:
        for left in x_positions:
            window_height = min(tile_size, height - top)
            window_width = min(tile_size, width - left)
            image, _ = load_triplet_window(
                record,
                top=top,
                left=left,
                height=window_height,
                width=window_width,
            )
            image = normalize_image(image, band_mean, band_std)
            xb = image.unsqueeze(0).to(device, non_blocking=False)
            patch_logits_sum = None
            for _, transform in transforms:
                augmented_x = transform(xb)
                augmented_logits = model(augmented_x).float()
                restored_logits = transform(augmented_logits)
                patch_logits_sum = (
                    restored_logits if patch_logits_sum is None else patch_logits_sum + restored_logits
                )
            patch_logits = (patch_logits_sum / len(transforms))[0].detach().cpu()
            if logits_sum is None:
                logits_sum = torch.zeros(
                    (patch_logits.shape[0], height, width),
                    dtype=torch.float32,
                )
            bottom = top + window_height
            right = left + window_width
            logits_sum[:, top:bottom, left:right] += patch_logits[:, :window_height, :window_width]
            weight_sum[top:bottom, left:right] += 1.0
            del image, xb, patch_logits, patch_logits_sum

    if logits_sum is None:
        raise RuntimeError("No inference windows were produced for {name}".format(name=record.name))
    return logits_sum / weight_sum.clamp_min(1.0).unsqueeze(0)


def windowed_validation_context(learn: Learner) -> Tuple[List[TripletRecord], torch.Tensor, torch.Tensor]:
    dataset = learn.dls.valid.dataset
    if not hasattr(dataset, "records") or not hasattr(dataset, "band_mean") or not hasattr(dataset, "band_std"):
        raise TypeError("Validation dataset does not expose records and normalization statistics.")
    return list(dataset.records), dataset.band_mean, dataset.band_std


def predict_logits_tiled(model: nn.Module, xb: torch.Tensor, tile_size: int, overlap: int) -> torch.Tensor:
    batch_size, _, height, width = xb.shape
    if batch_size != 1:
        raise ValueError("Tiled inference expects batch size 1, got {bs}".format(bs=batch_size))

    y_positions = sliding_positions(height, tile_size, overlap)
    x_positions = sliding_positions(width, tile_size, overlap)
    logits_sum = None
    weight_sum = torch.zeros((1, 1, height, width), device=xb.device, dtype=torch.float32)

    for top in y_positions:
        for left in x_positions:
            bottom = min(top + tile_size, height)
            right = min(left + tile_size, width)
            tile = xb[:, :, top:bottom, left:right]
            logits = model(tile).float()
            if logits_sum is None:
                logits_sum = torch.zeros(
                    (1, logits.shape[1], height, width), device=xb.device, dtype=torch.float32
                )
            logits_sum[:, :, top:bottom, left:right] += logits
            weight_sum[:, :, top:bottom, left:right] += 1.0

    return logits_sum / weight_sum.clamp_min(1.0)


def predict_logits_with_fallback(
    model: nn.Module,
    xb: torch.Tensor,
    tile_size: int,
    overlap: int,
    force_tiled: bool = False,
) -> Tuple[torch.Tensor, bool]:
    if force_tiled:
        return predict_logits_tiled(model, xb, tile_size=tile_size, overlap=overlap), True
    try:
        return model(xb).float(), False
    except RuntimeError as error:
        if not is_cuda_oom(error):
            raise
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        print({"eval_fallback": "switching to tiled inference"})
        return predict_logits_tiled(model, xb, tile_size=tile_size, overlap=overlap), True


def predict_water_probabilities(
    model: nn.Module,
    xb: torch.Tensor,
    tile_size: int,
    overlap: int,
    use_tta: bool = False,
    tta_mode_count: int = 4,
    force_tiled: bool = False,
) -> Tuple[torch.Tensor, bool]:
    if not use_tta:
        logits, used_tiled = predict_logits_with_fallback(
            model, xb, tile_size=tile_size, overlap=overlap, force_tiled=force_tiled
        )
        return F.softmax(logits, dim=1)[:, 1], used_tiled

    transforms = get_tta_transforms(tta_mode_count)
    probs_sum = None
    used_tiled_any = force_tiled
    for _, transform in transforms:
        augmented_x = transform(xb)
        logits, used_tiled = predict_logits_with_fallback(
            model, augmented_x, tile_size=tile_size, overlap=overlap, force_tiled=used_tiled_any
        )
        augmented_probs = F.softmax(logits, dim=1)[:, 1:2]
        restored_probs = transform(augmented_probs)
        probs_sum = restored_probs if probs_sum is None else probs_sum + restored_probs
        used_tiled_any = used_tiled_any or used_tiled
    return (probs_sum / len(transforms))[:, 0], used_tiled_any


def metrics_from_binary_confusion(tp: int, fp: int, tn: int, fn: int) -> dict:
    precision = tp / max(tp + fp, 1)
    recall = tp / max(tp + fn, 1)
    water_iou = tp / max(tp + fp + fn, 1)
    dice = (2 * tp) / max((2 * tp) + fp + fn, 1)
    accuracy = (tp + tn) / max(tp + tn + fp + fn, 1)
    other_iou = tn / max(tn + fp + fn, 1)
    other_dice = (2 * tn) / max((2 * tn) + fp + fn, 1)
    mean_iou = (other_iou + water_iou) / 2.0
    mean_dice = (other_dice + dice) / 2.0
    return {
        "precision": float(precision),
        "recall": float(recall),
        "water_iou": float(water_iou),
        "dice": float(mean_dice),
        "accuracy": float(accuracy),
        "jaccard": float(mean_iou),
    }


def validate_model(
    learn: Learner,
    tile_size: int,
    overlap: int,
    force_tiled: bool = False,
    max_valid_batches: Optional[int] = None,
) -> dict:
    device = next(learn.model.parameters()).device
    loss_func = learn.loss_func
    records, band_mean, band_std = windowed_validation_context(learn)
    total_loss = 0.0
    batch_count = 0
    tp = fp = tn = fn = 0
    use_tiled = False

    was_training = learn.model.training
    learn.model.eval()
    with torch.inference_mode():
        for batch_idx, record in enumerate(records):
            if max_valid_batches is not None and batch_idx >= max_valid_batches:
                break
            logits = predict_record_logits_windowed(
                learn.model,
                record,
                band_mean=band_mean,
                band_std=band_std,
                tile_size=tile_size,
                overlap=overlap,
                device=device,
            )
            yb = open_mask_binary(record.label_path).long()
            batch_loss = loss_func(
                logits.unsqueeze(0).to(device),
                yb.unsqueeze(0).to(device),
            ).detach().float().item()
            pred = logits.argmax(dim=0)
            observed = yb != IGNORE_INDEX
            total_loss += batch_loss
            batch_count += 1
            tp += int(((pred == 1) & (yb == 1) & observed).sum().item())
            fp += int(((pred == 1) & (yb != 1) & observed).sum().item())
            tn += int(((pred != 1) & (yb != 1) & observed).sum().item())
            fn += int(((pred != 1) & (yb == 1) & observed).sum().item())
            use_tiled = True
            del logits, yb, pred, observed
    learn.model.train(was_training)

    metric_summary = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn)
    metric_summary.update(
        {
            "loss": float(total_loss / max(batch_count, 1)),
            "used_tiled_inference": use_tiled,
            "tp": tp,
            "fp": fp,
            "tn": tn,
            "fn": fn,
        }
    )
    return metric_summary


def evaluate_threshold_sweep(
    learn: Learner,
    thresholds: Sequence[float],
    tile_size: int,
    overlap: int,
    use_tta: bool = False,
    tta_mode_count: int = 4,
    force_tiled: bool = False,
    max_valid_batches: Optional[int] = None,
) -> pd.DataFrame:
    device = next(learn.model.parameters()).device
    records, band_mean, band_std = windowed_validation_context(learn)
    counts = {name: torch.zeros(len(thresholds), dtype=torch.long) for name in ["tp", "fp", "tn", "fn"]}
    eval_mode = "tta" if use_tta else "plain"
    if force_tiled:
        eval_mode = "{mode}_tiled".format(mode=eval_mode)
    use_tiled = force_tiled

    was_training = learn.model.training
    learn.model.eval()
    with torch.inference_mode():
        for batch_idx, record in enumerate(records):
            if max_valid_batches is not None and batch_idx >= max_valid_batches:
                break
            logits = predict_record_logits_windowed(
                learn.model,
                record,
                band_mean=band_mean,
                band_std=band_std,
                tile_size=tile_size,
                overlap=overlap,
                device=device,
                use_tta=use_tta,
                tta_mode_count=tta_mode_count,
            )
            probs = F.softmax(logits, dim=0)[1]
            labels = open_mask_binary(record.label_path)
            observed = labels != IGNORE_INDEX
            targ = labels == 1
            for threshold_idx, threshold in enumerate(thresholds):
                pred = probs >= float(threshold)
                counts["tp"][threshold_idx] += (pred & targ & observed).sum()
                counts["fp"][threshold_idx] += (pred & ~targ & observed).sum()
                counts["tn"][threshold_idx] += ((~pred) & ~targ & observed).sum()
                counts["fn"][threshold_idx] += ((~pred) & targ & observed).sum()
            use_tiled = True
            del logits, probs, targ, labels, observed
    learn.model.train(was_training)

    records: List[dict] = []
    for idx, threshold in enumerate(thresholds):
        tp = int(counts["tp"][idx].item())
        fp = int(counts["fp"][idx].item())
        tn = int(counts["tn"][idx].item())
        fn = int(counts["fn"][idx].item())
        metrics = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn)
        records.append(
            {
                "eval_mode": eval_mode,
                "threshold": float(threshold),
                "used_tiled_inference": use_tiled,
                **metrics,
                "tp": tp,
                "fp": fp,
                "tn": tn,
                "fn": fn,
            }
        )
    return pd.DataFrame(records)


def evaluate_tiles_at_threshold(
    learn: Learner,
    valid_records: Sequence[TripletRecord],
    threshold: float,
    tile_size: int,
    overlap: int,
    use_tta: bool = False,
    tta_mode_count: int = 4,
    force_tiled: bool = False,
    max_valid_batches: Optional[int] = None,
) -> pd.DataFrame:
    device = next(learn.model.parameters()).device
    _, band_mean, band_std = windowed_validation_context(learn)
    rows: List[dict] = []
    eval_mode = "tta" if use_tta else "plain"
    if force_tiled:
        eval_mode = "{mode}_tiled".format(mode=eval_mode)

    was_training = learn.model.training
    learn.model.eval()
    with torch.inference_mode():
        for batch_idx, record in enumerate(valid_records):
            if max_valid_batches is not None and batch_idx >= max_valid_batches:
                break
            logits = predict_record_logits_windowed(
                learn.model,
                record,
                band_mean=band_mean,
                band_std=band_std,
                tile_size=tile_size,
                overlap=overlap,
                device=device,
                use_tta=use_tta,
                tta_mode_count=tta_mode_count,
            )
            probs = F.softmax(logits, dim=0)[1]
            yb = open_mask_binary(record.label_path).long()
            pred = probs >= threshold
            observed = yb != IGNORE_INDEX
            targ = yb == 1
            tp = int((pred & targ & observed).sum().item())
            fp = int((pred & ~targ & observed).sum().item())
            tn = int(((~pred) & ~targ & observed).sum().item())
            fn = int(((~pred) & targ & observed).sum().item())
            metrics = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn)
            rows.append(
                {
                    "tile_name": record.name,
                    "source_id": record.source_id,
                    "eval_mode": eval_mode,
                    "threshold": float(threshold),
                    "used_tiled_inference": True,
                    "water_pixels": int(targ.sum().item()),
                    **metrics,
                    "tp": tp,
                    "fp": fp,
                    "tn": tn,
                    "fn": fn,
                }
            )
            del logits, probs, yb, pred, targ, observed
    learn.model.train(was_training)
    return pd.DataFrame(rows)


def summarize_threshold_results(threshold_df: pd.DataFrame) -> dict:
    summary: dict = {}
    if threshold_df.empty:
        return summary

    for eval_mode in threshold_df["eval_mode"].unique():
        mode_df = threshold_df[threshold_df["eval_mode"] == eval_mode].sort_values(
            ["water_iou", "dice", "threshold"],
            ascending=[False, False, True],
        )
        best = mode_df.iloc[0]
        summary.update(
            {
                "{mode}_best_threshold".format(mode=eval_mode): round(float(best["threshold"]), 4),
                "{mode}_best_water_iou".format(mode=eval_mode): round(float(best["water_iou"]), 6),
                "{mode}_best_dice".format(mode=eval_mode): round(float(best["dice"]), 6),
                "{mode}_best_precision".format(mode=eval_mode): round(float(best["precision"]), 6),
                "{mode}_best_recall".format(mode=eval_mode): round(float(best["recall"]), 6),
            }
        )

    best_row = threshold_df.sort_values(["water_iou", "dice", "threshold"], ascending=[False, False, True]).iloc[0]
    summary.update(
        {
            "best_eval_mode": str(best_row["eval_mode"]),
            "best_eval_threshold": round(float(best_row["threshold"]), 4),
            "best_eval_water_iou": round(float(best_row["water_iou"]), 6),
            "best_eval_dice": round(float(best_row["dice"]), 6),
            "best_eval_precision": round(float(best_row["precision"]), 6),
            "best_eval_recall": round(float(best_row["recall"]), 6),
        }
    )
    return summary


def summarize_worst_tiles(tile_df: pd.DataFrame, worst_tile_count: int) -> pd.DataFrame:
    if tile_df.empty:
        return tile_df
    return tile_df.sort_values(
        ["water_iou", "precision", "recall", "tile_name"],
        ascending=[True, True, True, True],
    ).head(worst_tile_count)


def summarize_source_metrics(tile_df: pd.DataFrame) -> Dict[str, dict]:
    if tile_df.empty or "source_id" not in tile_df.columns:
        return {}
    summaries: Dict[str, dict] = {}
    for source_id, source_df in tile_df.groupby("source_id", dropna=False):
        source_name = str(source_id) if str(source_id) else "unspecified"
        tp = int(source_df["tp"].sum())
        fp = int(source_df["fp"].sum())
        tn = int(source_df["tn"].sum())
        fn = int(source_df["fn"].sum())
        metrics = metrics_from_binary_confusion(tp=tp, fp=fp, tn=tn, fn=fn)
        metrics["water_f1"] = float((2 * tp) / max((2 * tp) + fp + fn, 1))
        summaries[source_name] = {
            "tile_count": int(len(source_df)),
            **{key: round(float(value), 6) for key, value in metrics.items()},
            "tp": tp,
            "fp": fp,
            "tn": tn,
            "fn": fn,
        }
    return summaries


def run_training_schedule(
    learn: Learner,
    args: argparse.Namespace,
    schedule: Sequence[Tuple[str, int, object]],
    epoch_timer: Optional[EpochTimerCallback] = None,
) -> pd.DataFrame:
    training_history = []
    for phase_name, epochs, lr_max in schedule:
        apply_phase_freeze(learn, phase_name)
        training_history.append(
            run_phase(
                learn=learn,
                args=args,
                phase=phase_name,
                epochs=epochs,
                lr_max=lr_max,
                epoch_timer=epoch_timer,
            )
        )
        if epochs > 0:
            learn.load(args.best_model_name, with_opt=False)
    save_json(args.artifact_dir / "global_best.json", getattr(learn, "global_best_where", {}))   # full_coverage_v1
    return pd.concat(training_history, ignore_index=True) if training_history else pd.DataFrame(columns=["phase"])


def effective_epoch_count(args: argparse.Namespace) -> int:
    return int(args.freeze_epochs + args.partial_unfreeze_epochs + args.finetune_epochs + args.finetune_tail_epochs)


def probe_batch_size(
    args: argparse.Namespace,
    train_records: Sequence[TripletRecord],
    valid_records: Sequence[TripletRecord],
    band_mean: np.ndarray,
    band_std: np.ndarray,
    device: torch.device,
) -> dict:
    if device.type != "cuda":
        return {
            "selected_batch_size": args.batch_size,
            "probe_results": [],
            "probe_skipped": "CUDA unavailable",
        }

    results: List[dict] = []
    selected_batch_size = None
    train_probe_records = list(train_records[: max(max(args.probe_batch_sizes), 16)])
    valid_probe_records = list(valid_records[:1])

    for candidate in args.probe_batch_sizes:
        candidate_args = copy_module.deepcopy(args)
        candidate_args.batch_size = candidate
        candidate_args.num_workers = min(candidate_args.num_workers, 2)
        try:
            dls = build_dataloaders(
                candidate_args,
                train_records=train_probe_records,
                valid_records=valid_probe_records,
                band_mean=band_mean,
                band_std=band_std,
                device=device,
            )
            learn = build_learner(candidate_args, dls, valid_records=valid_probe_records)
            learn.model.train()
            optimizer = torch.optim.AdamW(learn.model.parameters(), lr=args.warmup_lr)
            xb, yb = dls.one_batch()
            optimizer.zero_grad(set_to_none=True)
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
                torch.cuda.reset_peak_memory_stats(device)
            start = time.perf_counter()
            if args.use_fp16:
                scaler = torch.cuda.amp.GradScaler(enabled=True)
                with torch.cuda.amp.autocast(enabled=True):
                    pred = learn.model(xb)
                    loss = learn.loss_func(pred, yb)
                scaler.scale(loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                pred = learn.model(xb)
                loss = learn.loss_func(pred, yb)
                loss.backward()
                optimizer.step()
            torch.cuda.synchronize(device)
            elapsed = time.perf_counter() - start
            peak_mem_gb = torch.cuda.max_memory_allocated(device) / (1024**3)
            result = {
                "batch_size": candidate,
                "success": True,
                "loss": float(loss.detach().float().item()),
                "step_seconds": round(elapsed, 4),
                "max_memory_gb": round(peak_mem_gb, 4),
            }
            results.append(result)
            selected_batch_size = candidate
            del pred, loss, xb, yb, learn, dls, optimizer
            gc.collect()
            torch.cuda.empty_cache()
            break
        except RuntimeError as error:
            if not is_cuda_oom(error):
                raise
            results.append({"batch_size": candidate, "success": False, "error": "cuda_oom"})
            gc.collect()
            torch.cuda.empty_cache()

    if selected_batch_size is None:
        raise RuntimeError("Batch size probe failed for all candidates: {candidates}".format(candidates=args.probe_batch_sizes))

    probe_summary = {"selected_batch_size": selected_batch_size, "probe_results": results}
    print(probe_summary)
    return probe_summary


def maybe_override_batch_size_from_summary(args: argparse.Namespace) -> None:
    if args.benchmark_summary_path is None:
        return
    if not args.benchmark_summary_path.exists():
        raise FileNotFoundError(
            "Benchmark summary path does not exist: {path}".format(path=args.benchmark_summary_path)
        )
    summary = json.loads(args.benchmark_summary_path.read_text())
    selected = summary.get("selected_batch_size")
    if not selected:
        raise ValueError(
            "Benchmark summary at {path} is missing selected_batch_size".format(path=args.benchmark_summary_path)
        )
    args.batch_size = int(selected)
    print(
        {
            "batch_size_override": args.batch_size,
            "benchmark_summary_path": str(args.benchmark_summary_path),
        }
    )


def build_audit_summary(
    args: argparse.Namespace,
    train_records: Sequence[TripletRecord],
    valid_records: Sequence[TripletRecord],
    band_mean: np.ndarray,
    band_std: np.ndarray,
    dls: DataLoaders,
) -> dict:
    sample_image, sample_mask = load_center_triplet_window(train_records[0], args.crop_size)
    batch_images, batch_masks = dls.one_batch()
    return {
        "paired_tiles": len(train_records) + len(valid_records),
        "train_tiles": len(train_records),
        "valid_tiles": len(valid_records),
        "sample_image_shape": tuple(sample_image.shape),
        "sample_mask_shape": tuple(sample_mask.shape),
        "sample_mask_values": sorted(torch.unique(sample_mask).tolist()),
        "band_mean_s1": band_mean[:3].round(6).tolist(),
        "band_std_s1": band_std[:3].round(6).tolist(),
        "band_mean_aef_first5": band_mean[3:8].round(6).tolist() if using_aef(args) else [],
        "band_std_aef_first5": band_std[3:8].round(6).tolist() if using_aef(args) else [],
        "train_batch_images": tuple(batch_images.shape),
        "train_batch_masks": tuple(batch_masks.shape),
        "model_family": args.model_family,
        "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None,
        "loss": args.loss,
        "loss_aux_weight": args.loss_aux_weight,
        "optimizer": args.optimizer,
        "schedule_policy": args.schedule_policy,
        "tuning_strategy": args.tuning_strategy,
        "augmentation_rotate_max_deg": getattr(dls.train.dataset, "rotate_max_deg", None),
        "augmentation_scale_jitter": [
            getattr(dls.train.dataset, "scale_jitter_min", None),
            getattr(dls.train.dataset, "scale_jitter_max", None),
        ],
        "augmentation_s1_gain_jitter": getattr(dls.train.dataset, "s1_gain_jitter", None),
        "augmentation_s1_bias_jitter": getattr(dls.train.dataset, "s1_bias_jitter", None),
        "augmentation_gaussian_noise_std": getattr(dls.train.dataset, "gaussian_noise_std", None),
        "positive_crop_p": getattr(dls.train.dataset, "positive_crop_p", None),
        "min_water_fraction": getattr(dls.train.dataset, "min_water_fraction", None),
        "positive_crop_attempts": getattr(dls.train.dataset, "positive_crop_attempts", None),
        "target_source_id": args.target_source_id,
        "target_source_fraction": args.target_source_fraction,
        "use_permanent_prior": using_prior(args),
        "prior_dropout_p": args.prior_dropout_p,
        "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0),
    }


def run_benchmark_mode(
    args: argparse.Namespace,
    device: torch.device,
    train_records: Sequence[TripletRecord],
    valid_records: Sequence[TripletRecord],
    band_mean: np.ndarray,
    band_std: np.ndarray,
) -> dict:
    planned_full_epochs = effective_epoch_count(args)
    probe_summary = probe_batch_size(args, train_records, valid_records, band_mean, band_std, device)
    args.batch_size = int(probe_summary["selected_batch_size"])
    if args.max_benchmark_batch_size is not None:
        args.batch_size = min(args.batch_size, args.max_benchmark_batch_size)
        probe_summary["selected_batch_size"] = args.batch_size
    train_records = subset_records(train_records, args.max_train_batches, args.batch_size)
    valid_records = subset_records(valid_records, args.max_valid_batches, 1)
    dls = build_dataloaders(args, train_records, valid_records, band_mean, band_std, device)
    print(build_audit_summary(args, train_records, valid_records, band_mean, band_std, dls))

    learn = build_learner(args, dls, valid_records=valid_records)
    maybe_warm_start(learn, args)
    epoch_timer = EpochTimerCallback()
    history_df = run_training_schedule(
        learn=learn,
        args=args,
        schedule=build_schedule(args),
        epoch_timer=epoch_timer,
    )
    history_path = args.artifact_dir / args.history_csv
    history_df.to_csv(history_path, index=False)

    validation = validate_model(
        learn,
        tile_size=args.crop_size,
        overlap=args.tile_overlap,
        force_tiled=args.force_tiled_validation,
        max_valid_batches=args.max_valid_batches,
    )
    seconds_per_epoch = sum(epoch_timer.epoch_times) / max(len(epoch_timer.epoch_times), 1)
    samples_per_second = (
        (len(train_records) / seconds_per_epoch) if seconds_per_epoch > 0 else None
    )
    benchmark_summary = {
        "mode": "benchmark",
        "model_family": args.model_family,
        "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None,
        "selected_batch_size": args.batch_size,
        "probe_results": probe_summary["probe_results"],
        "epoch_times_seconds": [round(v, 4) for v in epoch_timer.epoch_times],
        "avg_seconds_per_epoch": round(seconds_per_epoch, 4),
        "samples_per_second": None if samples_per_second is None else round(samples_per_second, 4),
        "projected_full_train_seconds": round(seconds_per_epoch * planned_full_epochs, 1),
        "projected_full_train_hours": round((seconds_per_epoch * planned_full_epochs) / 3600.0, 3),
        "train_records_used": len(train_records),
        "valid_records_used": len(valid_records),
        "validation": validation,
        "history_csv": str(history_path),
        "loss": args.loss,
        "selection_metric": args.selection_metric,
        "tuning_strategy": args.tuning_strategy,
        "force_tiled_validation": bool(args.force_tiled_validation),
        "split_seed": args.seed if args.split_seed is None else args.split_seed,
        "train_subset_seed": args.seed if args.train_subset_seed is None else args.train_subset_seed,
        "max_train_records": args.max_train_records,
        "target_source_id": args.target_source_id,
        "target_source_fraction": args.target_source_fraction,
        "use_permanent_prior": using_prior(args),
        "prior_dropout_p": args.prior_dropout_p,
        "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0),
    }
    save_json(args.artifact_dir / args.benchmark_summary_json, benchmark_summary)
    print(benchmark_summary)
    return benchmark_summary


def run_train_mode(
    args: argparse.Namespace,
    device: torch.device,
    train_records: Sequence[TripletRecord],
    valid_records: Sequence[TripletRecord],
    band_mean: np.ndarray,
    band_std: np.ndarray,
) -> dict:
    maybe_override_batch_size_from_summary(args)
    train_records = subset_records(train_records, args.max_train_batches, args.batch_size)
    valid_records = subset_records(valid_records, args.max_valid_batches, 1)
    dls = build_dataloaders(args, train_records, valid_records, band_mean, band_std, device)
    print(build_audit_summary(args, train_records, valid_records, band_mean, band_std, dls))

    learn = build_learner(args, dls, valid_records=valid_records)
    maybe_warm_start(learn, args)
    history_df = run_training_schedule(learn=learn, args=args, schedule=build_schedule(args))
    history_path = args.artifact_dir / args.history_csv
    history_df.to_csv(history_path, index=False)
    print(history_df.tail())

    if using_aef(args):
        bottleneck_report = inspect_bottleneck(learn.model, args.artifact_dir)
        print({"bottleneck_weight_norms": bottleneck_report["weight_norm_per_proj_band"]})

    validation = validate_model(
        learn,
        tile_size=args.crop_size,
        overlap=args.tile_overlap,
        force_tiled=args.force_tiled_validation,
        max_valid_batches=args.max_valid_batches,
    )
    threshold_results = [
        evaluate_threshold_sweep(
            learn,
            thresholds=args.eval_thresholds,
            tile_size=args.crop_size,
            overlap=args.tile_overlap,
            use_tta=False,
            force_tiled=args.force_tiled_validation,
            max_valid_batches=args.max_valid_batches,
        )
    ]
    if args.use_tta_eval:
        threshold_results.append(
            evaluate_threshold_sweep(
                learn,
                thresholds=args.eval_thresholds,
                tile_size=args.crop_size,
                overlap=args.tile_overlap,
                use_tta=True,
                tta_mode_count=args.tta_eval_mode_count,
                force_tiled=args.force_tiled_validation,
                max_valid_batches=args.max_valid_batches,
            )
        )
    threshold_df = pd.concat(threshold_results, ignore_index=True)
    threshold_path = args.artifact_dir / args.threshold_sweep_csv
    threshold_df.to_csv(threshold_path, index=False)
    print(
        threshold_df.sort_values(
            ["eval_mode", "water_iou", "precision", "threshold"],
            ascending=[True, False, False, True],
        ).groupby("eval_mode").head(5)
    )

    best_eval_row = threshold_df.sort_values(
        ["water_iou", "dice", "threshold"],
        ascending=[False, False, True],
    ).iloc[0]
    best_eval_mode = str(best_eval_row["eval_mode"])
    tile_metrics_df = evaluate_tiles_at_threshold(
        learn,
        valid_records=valid_records,
        threshold=float(best_eval_row["threshold"]),
        tile_size=args.crop_size,
        overlap=args.tile_overlap,
        use_tta=best_eval_mode.startswith("tta"),
        tta_mode_count=args.tta_eval_mode_count,
        force_tiled=best_eval_mode.endswith("_tiled"),
        max_valid_batches=args.max_valid_batches,
    )
    tile_metrics_path = args.artifact_dir / args.tile_metrics_csv
    tile_metrics_df.to_csv(tile_metrics_path, index=False)
    worst_tiles_df = summarize_worst_tiles(tile_metrics_df, worst_tile_count=args.worst_tile_count)
    worst_tiles_path = args.artifact_dir / args.worst_tiles_csv
    worst_tiles_df.to_csv(worst_tiles_path, index=False)

    summary = {
        "model_family": args.model_family,
        "segformer_variant": args.segformer_variant if args.model_family == "segformer" else None,
        "final_valid_loss": round(float(validation["loss"]), 6),
        "final_valid_jaccard": round(float(validation["jaccard"]), 6),
        "final_valid_dice": round(float(validation["dice"]), 6),
        "final_valid_water_iou": round(float(validation["water_iou"]), 6),
        "final_valid_precision": round(float(validation["precision"]), 6),
        "final_valid_recall": round(float(validation["recall"]), 6),
        "used_tiled_validation": bool(validation["used_tiled_inference"]),
        "best_model_path": str(args.artifact_dir / "models" / "{name}.pth".format(name=args.best_model_name)),
        "history_csv": str(history_path),
        "threshold_sweep_csv": str(threshold_path),
        "tile_metrics_csv": str(tile_metrics_path),
        "worst_tiles_csv": str(worst_tiles_path),
        "loss": args.loss,
        "selection_metric": args.selection_metric,
        "tuning_strategy": args.tuning_strategy,
        "force_tiled_validation": bool(args.force_tiled_validation),
        "split_seed": args.seed if args.split_seed is None else args.split_seed,
        "train_subset_seed": args.seed if args.train_subset_seed is None else args.train_subset_seed,
        "max_train_records": args.max_train_records,
        "target_source_id": args.target_source_id,
        "target_source_fraction": args.target_source_fraction,
        "use_permanent_prior": using_prior(args),
        "prior_dropout_p": args.prior_dropout_p,
        "aef_dropout_p": getattr(args, "aef_dropout_p", 0.0),
        "train_records_used": len(train_records),
        "valid_records_used": len(valid_records),
        "source_metrics_at_best_threshold": summarize_source_metrics(tile_metrics_df),
    }
    if using_aef(args):
        summary["bottleneck_proj_shape"] = "{aef}->{proj}".format(
            aef=args.n_aef_bands, proj=args.n_proj_bands
        )
    else:
        summary["input_bands"] = args.n_s1_bands
    summary.update(summarize_threshold_results(threshold_df))
    if not worst_tiles_df.empty:
        summary["worst_tile_examples"] = worst_tiles_df[
            ["tile_name", "water_iou", "precision", "recall"]
        ].round(6).to_dict(orient="records")
    save_json(args.artifact_dir / args.eval_summary_json, summary)
    print(summary)
    return summary


def main() -> None:
    args = parse_args()
    device = setup_environment(args)

    if args.triplet_manifest is not None:
        triplets = collect_triplets_from_manifest(
            args.triplet_manifest,
            args.triplet_sample_id,
            expect_aef=using_aef(args),
            expect_prior=using_prior(args),
        )
    else:
        triplets = collect_triplets(
            args.input_dir,
            args.aef_dir if using_aef(args) else None,
            args.label_dir,
        )
    split_seed = args.seed if args.split_seed is None else args.split_seed
    train_subset_seed = args.seed if args.train_subset_seed is None else args.train_subset_seed
    if args.fixed_split_manifest is not None:
        full_train_records, valid_records = split_triplets_from_manifest(
            triplets, args.fixed_split_manifest
        )
        split_protocol = "fixed_manifest"
    else:
        full_train_records, valid_records = split_triplets(
            triplets, valid_pct=args.valid_pct, seed=split_seed
        )
        split_protocol = "seeded"
    train_records = limit_train_records(
        full_train_records,
        max_records=args.max_train_records,
        subset_seed=train_subset_seed,
    )
    print(
        {
            "split_seed": split_seed,
            "split_protocol": split_protocol,
            "fixed_split_manifest": None if args.fixed_split_manifest is None else str(args.fixed_split_manifest),
            "train_subset_seed": train_subset_seed,
            "full_train_tiles": len(full_train_records),
            "train_tiles_after_record_cap": len(train_records),
            "valid_tiles": len(valid_records),
            "max_train_records": args.max_train_records,
        }
    )
    write_split_manifest(args.artifact_dir, train_records, valid_records)
    band_mean, band_std = get_band_stats(
        train_records,
        args.artifact_dir,
        args.stats_sample_size,
        source_stats_path=resolve_band_stats_source(args),
        expected_bands=args.n_s1_bands + args.n_aef_bands,
    )

    if args.mode == "benchmark":
        run_benchmark_mode(args, device, train_records, valid_records, band_mean, band_std)
    else:
        run_train_mode(args, device, train_records, valid_records, band_mean, band_std)


if __name__ == "__main__":
    main()