training/train_cold_tiled_mainline.py

from __future__ import annotations

import argparse
import csv
import json
import math
import random
import time
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset

from eval_metrics import metric_bundle
from train_utils import set_seed


def load_json(path: Path) -> Dict[str, object]:
    return json.loads(path.read_text(encoding="utf-8"))


def read_rows(path: Path) -> List[Dict[str, str]]:
    with path.open("r", encoding="utf-8", newline="") as fh:
        return list(csv.DictReader(fh))


def write_json(path: Path, data: Dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, indent=2), encoding="utf-8")


def integral_image(mask: np.ndarray) -> np.ndarray:
    return np.pad(mask.astype(np.int32), ((1, 0), (1, 0)), mode="constant").cumsum(0).cumsum(1)


def rect_sum(ii: np.ndarray, top: int, left: int, height: int, width: int) -> int:
    bottom = top + height
    right = left + width
    return int(ii[bottom, right] - ii[top, right] - ii[bottom, left] + ii[top, left])


def centered_tile_origin(y: int, x: int, tile_h: int, tile_w: int, h: int, w: int) -> Tuple[int, int]:
    top = min(max(y - tile_h // 2, 0), h - tile_h)
    left = min(max(x - tile_w // 2, 0), w - tile_w)
    return int(top), int(left)


def containing_tile_origin(y: int, x: int, tile_h: int, tile_w: int, h: int, w: int, rng: random.Random) -> Tuple[int, int]:
    max_top = max(h - tile_h, 0)
    max_left = max(w - tile_w, 0)
    top_min = max(0, y - tile_h + 1)
    top_max = min(y, max_top)
    left_min = max(0, x - tile_w + 1)
    left_max = min(x, max_left)
    if top_min > top_max or left_min > left_max:
        return centered_tile_origin(y, x, tile_h, tile_w, h, w)
    return int(rng.randint(top_min, top_max)), int(rng.randint(left_min, left_max))


def positive_tile_origins(
    mask: np.ndarray,
    tile_h: int,
    tile_w: int,
    max_tiles: int,
    rng: random.Random,
    placement: str = "center",
) -> List[Tuple[int, int]]:
    ys, xs = np.where(mask > 0.5)
    if ys.size == 0:
        return []
    origin_fn = containing_tile_origin if placement == "random_containing" else centered_tile_origin
    origins = []
    seen = set()
    indices = list(range(int(ys.size)))
    rng.shuffle(indices)
    for i in indices:
        y = int(ys[i])
        x = int(xs[i])
        if placement == "random_containing":
            origin = origin_fn(y, x, tile_h, tile_w, mask.shape[0], mask.shape[1], rng)
        else:
            origin = origin_fn(y, x, tile_h, tile_w, mask.shape[0], mask.shape[1])
        if origin in seen:
            continue
        seen.add(origin)
        origins.append(origin)
        if max_tiles > 0 and len(origins) >= max_tiles:
            break
    rng.shuffle(origins)
    if max_tiles > 0:
        origins = origins[:max_tiles]
    return origins


def negative_tile_origins(
    mask: np.ndarray,
    tile_h: int,
    tile_w: int,
    desired: int,
    rng: random.Random,
    forbidden: Iterable[Tuple[int, int]] = (),
) -> List[Tuple[int, int]]:
    h, w = mask.shape
    ii = integral_image(mask)
    forbidden_set = set(forbidden)
    out: List[Tuple[int, int]] = []
    seen = set(forbidden_set)
    max_top = max(h - tile_h, 0)
    max_left = max(w - tile_w, 0)
    max_attempts = max(2000, desired * 50)
    attempts = 0
    while len(out) < desired and attempts < max_attempts:
        attempts += 1
        top = rng.randint(0, max_top)
        left = rng.randint(0, max_left)
        origin = (top, left)
        if origin in seen:
            continue
        if rect_sum(ii, top, left, tile_h, tile_w) != 0:
            continue
        seen.add(origin)
        out.append(origin)
    return out


def make_norm(norm_type: str, num_channels: int, norm_groups: int) -> nn.Module:
    if norm_type == "batch":
        return nn.BatchNorm2d(num_channels)
    if norm_type == "group":
        groups = max(1, min(int(norm_groups), num_channels))
        while num_channels % groups != 0 and groups > 1:
            groups -= 1
        return nn.GroupNorm(groups, num_channels)
    if norm_type == "instance":
        return nn.InstanceNorm2d(num_channels, affine=True)
    if norm_type in {"none", "identity"}:
        return nn.Identity()
    raise ValueError(f"Unsupported norm_type: {norm_type}")


class ConvBlock(nn.Module):
    def __init__(self, in_ch: int, out_ch: int, norm_type: str, norm_groups: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),
            make_norm(norm_type, out_ch, norm_groups),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False),
            make_norm(norm_type, out_ch, norm_groups),
            nn.ReLU(inplace=True),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class UNetSmallFlex(nn.Module):
    def __init__(
        self,
        in_ch: int,
        base: int = 32,
        dropout: float = 0.1,
        norm_type: str = "group",
        norm_groups: int = 8,
        prior_prob: float | None = None,
        use_aux_spatial_head: bool = False,
        aux_prior_prob: float | None = None,
    ):
        super().__init__()
        self.enc1 = ConvBlock(in_ch, base, norm_type, norm_groups)
        self.enc2 = ConvBlock(base, base * 2, norm_type, norm_groups)
        self.enc3 = ConvBlock(base * 2, base * 4, norm_type, norm_groups)
        self.enc4 = ConvBlock(base * 4, base * 8, norm_type, norm_groups)
        self.pool = nn.MaxPool2d(2)
        self.bottleneck = ConvBlock(base * 8, base * 16, norm_type, norm_groups)
        self.up4 = nn.ConvTranspose2d(base * 16, base * 8, 2, stride=2)
        self.dec4 = ConvBlock(base * 16, base * 8, norm_type, norm_groups)
        self.up3 = nn.ConvTranspose2d(base * 8, base * 4, 2, stride=2)
        self.dec3 = ConvBlock(base * 8, base * 4, norm_type, norm_groups)
        self.up2 = nn.ConvTranspose2d(base * 4, base * 2, 2, stride=2)
        self.dec2 = ConvBlock(base * 4, base * 2, norm_type, norm_groups)
        self.up1 = nn.ConvTranspose2d(base * 2, base, 2, stride=2)
        self.dec1 = ConvBlock(base * 2, base, norm_type, norm_groups)
        self.drop = nn.Dropout2d(p=dropout)
        self.head = nn.Conv2d(base, 1, kernel_size=1)
        self.use_aux_spatial_head = bool(use_aux_spatial_head)
        self.aux_head = nn.Conv2d(base, 1, kernel_size=1) if self.use_aux_spatial_head else None
        if prior_prob is not None:
            prior_prob = float(min(max(prior_prob, 1e-6), 1.0 - 1e-6))
            nn.init.constant_(self.head.bias, math.log(prior_prob / (1.0 - prior_prob)))
        if self.aux_head is not None and aux_prior_prob is not None:
            aux_prior_prob = float(min(max(aux_prior_prob, 1e-6), 1.0 - 1e-6))
            nn.init.constant_(self.aux_head.bias, math.log(aux_prior_prob / (1.0 - aux_prior_prob)))

    @staticmethod
    def _match_hw(x: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
        diff_y = ref.size(2) - x.size(2)
        diff_x = ref.size(3) - x.size(3)
        if diff_y > 0 or diff_x > 0:
            x = F.pad(x, [diff_x // 2, diff_x - diff_x // 2, diff_y // 2, diff_y - diff_y // 2])
        if diff_y < 0:
            y0 = (-diff_y) // 2
            x = x[:, :, y0 : y0 + ref.size(2), :]
        if diff_x < 0:
            x0 = (-diff_x) // 2
            x = x[:, :, :, x0 : x0 + ref.size(3)]
        return x

    def forward(self, x: torch.Tensor, return_aux: bool = False):
        e1 = self.enc1(x)
        e2 = self.enc2(self.pool(e1))
        e3 = self.enc3(self.pool(e2))
        e4 = self.enc4(self.pool(e3))
        b = self.bottleneck(self.pool(e4))
        d4 = self.dec4(torch.cat([self._match_hw(self.up4(b), e4), e4], dim=1))
        d3 = self.dec3(torch.cat([self._match_hw(self.up3(d4), e3), e3], dim=1))
        d2 = self.dec2(torch.cat([self._match_hw(self.up2(d3), e2), e2], dim=1))
        d1 = self.dec1(torch.cat([self._match_hw(self.up1(d2), e1), e1], dim=1))
        features = self.drop(d1)
        logits = self.head(features)
        if return_aux and self.aux_head is not None:
            return logits, self.aux_head(features)
        return logits


def normalization_config(config: Dict[str, object]) -> Dict[str, object]:
    raw = config.get("input_normalization", {})
    if raw is True:
        return {"enabled": True}
    if isinstance(raw, dict):
        return raw
    return {"enabled": False}


def compute_input_normalization_stats(
    rows: List[Dict[str, str]],
    continuous_channel_indices: List[int],
    eps: float,
) -> Dict[str, object]:
    count: Dict[int, int] = {int(idx): 0 for idx in continuous_channel_indices}
    sum_x: Dict[int, float] = {int(idx): 0.0 for idx in continuous_channel_indices}
    sum_x2: Dict[int, float] = {int(idx): 0.0 for idx in continuous_channel_indices}
    channel_names: List[str] | None = None
    static_names: List[str] | None = None
    for row in rows:
        sample = np.load(row["sample_path"], allow_pickle=True)
        weather = ColdFeatureStore._sanitize(sample["weather"].astype(np.float32))
        firewx = ColdFeatureStore._sanitize(sample["firewx"].astype(np.float32))
        extra = []
        if "firewx_valid" in sample:
            extra.append(sample["firewx_valid"].astype(np.float32))
        static_npz = np.load(row["static_path"], allow_pickle=True)
        static = ColdFeatureStore._sanitize(static_npz["static"].astype(np.float32))
        static_parts = []
        if "static_valid" in static_npz:
            static_parts.append(static_npz["static_valid"].astype(np.float32))
        static_parts.append(static)
        if channel_names is None:
            weather_names = [str(v) for v in sample.get("weather_names", np.array([], dtype=object)).tolist()]
            firewx_names = [str(v) for v in sample.get("firewx_names", np.array([], dtype=object)).tolist()]
            extra_names = ["firewx_valid"] if "firewx_valid" in sample else []
            static_names = [str(v) for v in static_npz.get("static_names", np.array([], dtype=object)).tolist()]
            static_valid_names = ["static_valid"] if "static_valid" in static_npz else []
            channel_names = weather_names + firewx_names + extra_names + static_valid_names + static_names
        x = np.concatenate([weather, firewx, *extra, *static_parts], axis=0).astype(np.float32)
        for idx in continuous_channel_indices:
            arr = x[int(idx)].astype(np.float64, copy=False).ravel()
            count[int(idx)] += int(arr.size)
            sum_x[int(idx)] += float(arr.sum())
            sum_x2[int(idx)] += float(np.square(arr).sum())
    channels = []
    for idx in continuous_channel_indices:
        n = max(int(count[int(idx)]), 1)
        mean = float(sum_x[int(idx)] / n)
        variance = max(float(sum_x2[int(idx)] / n - mean * mean), 0.0)
        std = float(max(math.sqrt(variance), eps))
        channels.append(
            {
                "index": int(idx),
                "name": channel_names[int(idx)] if channel_names and int(idx) < len(channel_names) else str(idx),
                "mean": mean,
                "std": std,
                "count": int(count[int(idx)]),
            }
        )
    return {
        "enabled": True,
        "method": "per-channel z-score",
        "stats_source": "train split full maps",
        "continuous_channel_indices": [int(idx) for idx in continuous_channel_indices],
        "eps": float(eps),
        "channels": channels,
        "channel_names": channel_names or [],
        "static_names": static_names or [],
    }


def load_or_compute_input_normalization_stats(
    config: Dict[str, object],
    train_rows: List[Dict[str, str]],
    metric_dir: Path,
) -> Dict[str, object] | None:
    norm_cfg = normalization_config(config)
    if not bool(norm_cfg.get("enabled", False)):
        return None
    stats_path_value = str(norm_cfg.get("stats_path", "")).strip()
    if stats_path_value:
        stats = load_json(Path(stats_path_value))
    else:
        stats = compute_input_normalization_stats(
            rows=train_rows,
            continuous_channel_indices=[int(v) for v in norm_cfg.get("continuous_channel_indices", list(range(10)))],
            eps=float(norm_cfg.get("eps", 1e-6)),
        )
    write_json(metric_dir / "input_normalization_stats.json", stats)
    return stats


class ColdFeatureStore:
    def __init__(self, rows: List[Dict[str, str]], normalization_stats: Dict[str, object] | None = None):
        self.cache: Dict[str, Dict[str, np.ndarray]] = {}
        self.normalization_stats = normalization_stats
        static_path = Path(rows[0]["static_path"])
        static_npz = np.load(static_path, allow_pickle=True)
        static = self._sanitize(static_npz["static"].astype(np.float32))
        static_valid = static_npz["static_valid"].astype(np.float32) if "static_valid" in static_npz else None
        static_parts = []
        if static_valid is not None:
            static_parts.append(static_valid.astype(np.float32))
        static_parts.append(static)
        static_x = np.concatenate(static_parts, axis=0).astype(np.float32)
        for row in rows:
            sample = np.load(row["sample_path"], allow_pickle=True)
            weather = self._sanitize(sample["weather"].astype(np.float32))
            firewx = self._sanitize(sample["firewx"].astype(np.float32))
            extra = []
            if "firewx_valid" in sample:
                extra.append(sample["firewx_valid"].astype(np.float32))
            x = np.concatenate([weather, firewx, *extra, static_x], axis=0).astype(np.float32)
            x = self._apply_normalization(x)
            y = np.nan_to_num(sample["y_occ"].astype(np.float32), nan=0.0, posinf=0.0, neginf=0.0)
            self.cache[str(row["sample_id"])] = {"x": x, "y": y}

    @staticmethod
    def _sanitize(x: np.ndarray) -> np.ndarray:
        x = np.where(x <= -9000.0, np.nan, x)
        x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)
        return x.astype(np.float32, copy=False)

    def _apply_normalization(self, x: np.ndarray) -> np.ndarray:
        if not self.normalization_stats:
            return x
        for item in self.normalization_stats.get("channels", []):
            idx = int(item["index"])
            if idx < 0 or idx >= x.shape[0]:
                continue
            mean = float(item["mean"])
            std = max(float(item["std"]), float(self.normalization_stats.get("eps", 1e-6)))
            x[idx] = (x[idx] - mean) / std
        return x.astype(np.float32, copy=False)

    def get(self, sample_id: str) -> Dict[str, np.ndarray]:
        return self.cache[sample_id]


def pooled_positive_rate(rows: List[Dict[str, str]], store: ColdFeatureStore, radius: int) -> float:
    if radius <= 0:
        return float(full_map_stats(rows, store)["positive_rate"])
    pos = 0.0
    total = 0.0
    for row in rows:
        sample = store.get(str(row["sample_id"]))
        y = torch.from_numpy(sample["y"].astype(np.float32)).unsqueeze(0)
        pooled = F.max_pool2d(y, kernel_size=radius * 2 + 1, stride=1, padding=radius)
        pos += float((pooled > 0.5).sum().item())
        total += float(pooled.numel())
    return float(pos / total) if total > 0 else 0.0


def transform_spatial_aux_target(target: torch.Tensor, radius: int) -> torch.Tensor:
    if radius <= 0:
        return target
    return F.max_pool2d(target, kernel_size=radius * 2 + 1, stride=1, padding=radius)


def transform_train_target(target: torch.Tensor, config: Dict[str, object]) -> torch.Tensor:
    mode = str(config.get("train_target_mode", "hard"))
    radius = int(config.get("train_target_radius", 0))
    if mode == "hard" or radius <= 0:
        return target
    kernel = radius * 2 + 1
    if mode == "dilate_max":
        return F.max_pool2d(target, kernel_size=kernel, stride=1, padding=radius)
    if mode == "soft_pool":
        return F.avg_pool2d(target, kernel_size=kernel, stride=1, padding=radius)
    raise ValueError(f"Unsupported train_target_mode: {mode}")


def maybe_unpack_logits(output):
    if isinstance(output, tuple):
        return output
    return output, None


class FocalBCEWithLogitsLoss(nn.Module):
    def __init__(self, gamma: float = 2.0, alpha: float = 0.75):
        super().__init__()
        self.gamma = float(gamma)
        self.alpha = float(alpha)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        bce = F.binary_cross_entropy_with_logits(logits, targets, reduction="none")
        probs = torch.sigmoid(logits)
        pt = torch.where(targets > 0.5, probs, 1.0 - probs)
        alpha_t = torch.where(
            targets > 0.5,
            torch.full_like(targets, self.alpha),
            torch.full_like(targets, 1.0 - self.alpha),
        )
        loss = alpha_t * ((1.0 - pt) ** self.gamma) * bce
        return loss.mean()


class OHEMBCEWithLogitsLoss(nn.Module):
    def __init__(self, pos_weight: float, neg_pos_ratio: float = 8.0, min_negatives: int = 64):
        super().__init__()
        self.pos_weight = float(pos_weight)
        self.neg_pos_ratio = float(neg_pos_ratio)
        self.min_negatives = int(min_negatives)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        losses = F.binary_cross_entropy_with_logits(
            logits,
            targets,
            reduction="none",
            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
        )
        batch_losses: List[torch.Tensor] = []
        batch_size = int(logits.shape[0])
        for i in range(batch_size):
            loss_i = losses[i].reshape(-1)
            tgt_i = targets[i].reshape(-1) > 0.5
            pos_loss = loss_i[tgt_i]
            neg_loss = loss_i[~tgt_i]
            num_pos = int(pos_loss.numel())
            keep_neg = max(self.min_negatives, int(self.neg_pos_ratio * max(num_pos, 1)))
            keep_neg = min(keep_neg, int(neg_loss.numel()))
            if keep_neg > 0:
                neg_loss = torch.topk(neg_loss, k=keep_neg, largest=True).values
            else:
                neg_loss = neg_loss[:0]
            denom = max(num_pos + int(neg_loss.numel()), 1)
            batch_losses.append((pos_loss.sum() + neg_loss.sum()) / denom)
        return torch.stack(batch_losses).mean()


class SoftFbetaBCEWithLogitsLoss(nn.Module):
    def __init__(self, pos_weight: float, beta: float = 2.0, bce_weight: float = 0.4, smooth: float = 1.0):
        super().__init__()
        self.pos_weight = float(pos_weight)
        self.beta = float(beta)
        self.bce_weight = float(bce_weight)
        self.smooth = float(smooth)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        targets = targets.clamp(0.0, 1.0)
        bce = F.binary_cross_entropy_with_logits(
            logits,
            targets,
            reduction="mean",
            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
        )
        probs = torch.sigmoid(logits.float())
        targets_f = targets.float()
        dims = tuple(range(1, probs.ndim))
        tp = (probs * targets_f).sum(dim=dims)
        fp = (probs * (1.0 - targets_f)).sum(dim=dims)
        fn = ((1.0 - probs) * targets_f).sum(dim=dims)
        beta_sq = self.beta * self.beta
        fbeta = ((1.0 + beta_sq) * tp + self.smooth) / ((1.0 + beta_sq) * tp + beta_sq * fn + fp + self.smooth)
        soft_loss = 1.0 - fbeta.mean()
        return self.bce_weight * bce + (1.0 - self.bce_weight) * soft_loss


class TverskyBCEWithLogitsLoss(nn.Module):
    def __init__(
        self,
        pos_weight: float,
        alpha: float = 0.3,
        beta: float = 0.7,
        bce_weight: float = 0.4,
        smooth: float = 1.0,
    ):
        super().__init__()
        self.pos_weight = float(pos_weight)
        self.alpha = float(alpha)
        self.beta = float(beta)
        self.bce_weight = float(bce_weight)
        self.smooth = float(smooth)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        targets = targets.clamp(0.0, 1.0)
        bce = F.binary_cross_entropy_with_logits(
            logits,
            targets,
            reduction="mean",
            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
        )
        probs = torch.sigmoid(logits.float())
        targets_f = targets.float()
        dims = tuple(range(1, probs.ndim))
        tp = (probs * targets_f).sum(dim=dims)
        fp = (probs * (1.0 - targets_f)).sum(dim=dims)
        fn = ((1.0 - probs) * targets_f).sum(dim=dims)
        tversky = (tp + self.smooth) / (tp + self.alpha * fp + self.beta * fn + self.smooth)
        return self.bce_weight * bce + (1.0 - self.bce_weight) * (1.0 - tversky.mean())


class FSSBCEWithLogitsLoss(nn.Module):
    def __init__(
        self,
        pos_weight: float,
        radii: Iterable[int] = (1, 2),
        bce_weight: float = 0.5,
        eps: float = 1e-6,
    ):
        super().__init__()
        self.pos_weight = float(pos_weight)
        self.radii = [int(v) for v in radii if int(v) > 0]
        if not self.radii:
            self.radii = [1]
        self.bce_weight = float(bce_weight)
        self.eps = float(eps)

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        targets = targets.clamp(0.0, 1.0)
        bce = F.binary_cross_entropy_with_logits(
            logits,
            targets,
            reduction="mean",
            pos_weight=torch.tensor([self.pos_weight], device=logits.device, dtype=logits.dtype),
        )
        probs = torch.sigmoid(logits.float())
        targets_f = targets.float()
        fss_losses: List[torch.Tensor] = []
        for radius in self.radii:
            kernel = radius * 2 + 1
            pred_frac = F.avg_pool2d(probs, kernel_size=kernel, stride=1, padding=radius)
            target_frac = F.avg_pool2d(targets_f, kernel_size=kernel, stride=1, padding=radius)
            mse = torch.mean((pred_frac - target_frac) ** 2)
            reference = torch.mean(pred_frac**2 + target_frac**2)
            fss = 1.0 - mse / (reference + self.eps)
            fss_losses.append(1.0 - fss)
        fss_loss = torch.stack(fss_losses).mean()
        return self.bce_weight * bce + (1.0 - self.bce_weight) * fss_loss


def build_loss(config: Dict[str, object], pos_weight: float, device: torch.device) -> nn.Module:
    loss_type = str(config.get("loss_type", "bce"))
    if loss_type == "bce":
        return nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weight], dtype=torch.float32, device=device))
    if loss_type == "focal_bce":
        return FocalBCEWithLogitsLoss(
            gamma=float(config.get("focal_gamma", 2.0)),
            alpha=float(config.get("focal_alpha", 0.75)),
        )
    if loss_type == "ohem_bce":
        return OHEMBCEWithLogitsLoss(
            pos_weight=pos_weight,
            neg_pos_ratio=float(config.get("ohem_neg_pos_ratio", 8.0)),
            min_negatives=int(config.get("ohem_min_negatives", 64)),
        )
    if loss_type == "soft_fbeta_bce":
        return SoftFbetaBCEWithLogitsLoss(
            pos_weight=pos_weight,
            beta=float(config.get("soft_fbeta_beta", 2.0)),
            bce_weight=float(config.get("soft_metric_bce_weight", 0.4)),
            smooth=float(config.get("soft_metric_smooth", 1.0)),
        )
    if loss_type == "tversky_bce":
        return TverskyBCEWithLogitsLoss(
            pos_weight=pos_weight,
            alpha=float(config.get("tversky_alpha", 0.3)),
            beta=float(config.get("tversky_beta", 0.7)),
            bce_weight=float(config.get("soft_metric_bce_weight", 0.4)),
            smooth=float(config.get("soft_metric_smooth", 1.0)),
        )
    if loss_type == "fss_bce":
        return FSSBCEWithLogitsLoss(
            pos_weight=pos_weight,
            radii=[int(v) for v in config.get("fss_loss_radii", [1, 2])],
            bce_weight=float(config.get("soft_metric_bce_weight", 0.5)),
        )
    raise ValueError(f"Unsupported loss_type: {loss_type}")


def full_map_stats(rows: List[Dict[str, str]], store: ColdFeatureStore) -> Dict[str, float]:
    pos = 0.0
    total = 0.0
    for row in rows:
        sample = store.get(str(row["sample_id"]))
        y = sample["y"]
        pos += float((y > 0.5).sum())
        total += float(y.size)
    neg = max(total - pos, 0.0)
    raw_pos_weight = float(max(neg / max(pos, 1.0), 1.0)) if total > 0 else 1.0
    return {
        "source": "full_map",
        "positive_cells": pos,
        "total_cells": total,
        "positive_rate": float(pos / total) if total > 0 else 0.0,
        "raw_pos_weight": raw_pos_weight,
    }


def tile_stats(tile_rows: List[Dict[str, object]], store: ColdFeatureStore) -> Dict[str, float]:
    pos = 0.0
    total = 0.0
    for row in tile_rows:
        sample = store.get(str(row["sample_id"]))
        top = int(row["top"])
        left = int(row["left"])
        tile_size = int(row["tile_size"])
        target = sample["y"][0, top : top + tile_size, left : left + tile_size]
        pos += float((target > 0.5).sum())
        total += float(target.size)
    neg = max(total - pos, 0.0)
    raw_pos_weight = float(max(neg / max(pos, 1.0), 1.0)) if total > 0 else 1.0
    return {
        "source": "tiles",
        "positive_cells": pos,
        "total_cells": total,
        "positive_rate": float(pos / total) if total > 0 else 0.0,
        "raw_pos_weight": raw_pos_weight,
    }


def select_stats(source: str, train_rows: List[Dict[str, str]], tile_rows: List[Dict[str, object]], store: ColdFeatureStore) -> Dict[str, float]:
    if source == "full_map":
        return full_map_stats(train_rows, store)
    if source == "tiles":
        return tile_stats(tile_rows, store)
    raise ValueError(f"Unsupported stats source: {source}")


def build_train_tiles(
    train_rows: List[Dict[str, str]],
    store: ColdFeatureStore,
    tile_size: int,
    max_positive_tiles_per_sample: int,
    min_negative_tiles_per_sample: int,
    neg_pos_ratio: float,
    positive_tile_placement: str,
    rng: random.Random,
) -> List[Dict[str, object]]:
    tile_rows: List[Dict[str, object]] = []
    for row in train_rows:
        sample_id = str(row["sample_id"])
        mask = store.get(sample_id)["y"][0]
        pos_origins = positive_tile_origins(
            mask=mask,
            tile_h=tile_size,
            tile_w=tile_size,
            max_tiles=max_positive_tiles_per_sample,
            rng=rng,
            placement=positive_tile_placement,
        )
        neg_count = max(min_negative_tiles_per_sample, int(math.ceil(len(pos_origins) * neg_pos_ratio)))
        neg_origins = negative_tile_origins(
            mask=mask,
            tile_h=tile_size,
            tile_w=tile_size,
            desired=neg_count,
            rng=rng,
            forbidden=pos_origins,
        )
        for top, left in pos_origins:
            tile_rows.append(
                {
                    "sample_id": sample_id,
                    "tile_type": "positive",
                    "top": int(top),
                    "left": int(left),
                    "tile_size": int(tile_size),
                }
            )
        for top, left in neg_origins:
            tile_rows.append(
                {
                    "sample_id": sample_id,
                    "tile_type": "negative",
                    "top": int(top),
                    "left": int(left),
                    "tile_size": int(tile_size),
                }
            )
    return tile_rows


class TrainTileDataset(Dataset):
    def __init__(self, tile_rows: List[Dict[str, object]], store: ColdFeatureStore, augment_flip: bool, seed: int):
        self.rows = tile_rows
        self.store = store
        self.augment_flip = bool(augment_flip)
        self.rng = random.Random(seed)

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, idx: int):
        row = self.rows[idx]
        sample = self.store.get(str(row["sample_id"]))
        top = int(row["top"])
        left = int(row["left"])
        tile_size = int(row["tile_size"])
        x = sample["x"][:, top : top + tile_size, left : left + tile_size]
        y = sample["y"][:, top : top + tile_size, left : left + tile_size]
        if self.augment_flip:
            if self.rng.random() < 0.5:
                x = x[:, :, ::-1].copy()
                y = y[:, :, ::-1].copy()
            if self.rng.random() < 0.5:
                x = x[:, ::-1, :].copy()
                y = y[:, ::-1, :].copy()
        return {"x": torch.from_numpy(x), "y": torch.from_numpy(y)}


class FullMapDataset(Dataset):
    def __init__(self, rows: List[Dict[str, str]], store: ColdFeatureStore):
        self.rows = rows
        self.store = store

    def __len__(self) -> int:
        return len(self.rows)

    def __getitem__(self, idx: int):
        row = self.rows[idx]
        sample = self.store.get(str(row["sample_id"]))
        return {"x": torch.from_numpy(sample["x"]), "y": torch.from_numpy(sample["y"]), "sample_id": str(row["sample_id"])}


def evaluate(
    model: nn.Module,
    loader: DataLoader,
    device: torch.device,
    thresholds: List[float],
    topk_area_fractions: List[float],
    fss_radii: List[int],
    n_bins: int,
    reference_positive_rate: float,
    criterion: nn.Module,
    amp: bool,
) -> Dict[str, object]:
    model.eval()
    total_loss = 0.0
    total_items = 0
    all_prob_maps: List[np.ndarray] = []
    all_target_maps: List[np.ndarray] = []
    with torch.no_grad():
        for batch in loader:
            x = torch.nan_to_num(batch["x"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
            y = torch.nan_to_num(batch["y"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=amp and device.type == "cuda"):
                logits, _ = maybe_unpack_logits(model(x))
                loss = criterion(logits, y)
            prob = torch.sigmoid(logits.float()).detach().cpu().numpy()[:, 0, :, :]
            target = y.float().detach().cpu().numpy()[:, 0, :, :]
            all_prob_maps.append(prob)
            all_target_maps.append(target)
            total_loss += float(loss.item()) * x.size(0)
            total_items += int(x.size(0))
    prob_maps = np.concatenate(all_prob_maps, axis=0)
    target_maps = np.concatenate(all_target_maps, axis=0)
    metrics = metric_bundle(
        prob_maps=prob_maps,
        target_maps=target_maps,
        thresholds=thresholds,
        topk_fractions=topk_area_fractions,
        fss_radii=fss_radii,
        n_bins=n_bins,
        reference_positive_rate=reference_positive_rate,
    )
    metrics.update({"loss": total_loss / max(total_items, 1), "num_samples": int(prob_maps.shape[0])})
    return metrics


def save_checkpoint(path: Path, model: nn.Module, optimizer: torch.optim.Optimizer, epoch: int, metrics: Dict[str, float], config: Dict[str, object]) -> None:
    torch.save(
        {
            "epoch": epoch,
            "model": model.state_dict(),
            "optimizer": optimizer.state_dict(),
            "metrics": metrics,
            "config": config,
        },
        path,
    )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--run-name", type=str, required=True)
    args = parser.parse_args()

    config = load_json(args.config)
    set_seed(int(config.get("seed", 7)))
    rng = random.Random(int(config.get("seed", 7)))
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True

    index_root = Path(config["index_root"])
    run_root = Path(config["run_root"])
    ckpt_dir = run_root / "checkpoints" / args.run_name
    metric_dir = run_root / "metrics" / args.run_name
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    metric_dir.mkdir(parents=True, exist_ok=True)

    train_rows = read_rows(index_root / "splits" / "train.csv")
    val_rows = read_rows(index_root / "splits" / "val.csv")
    test_rows = read_rows(index_root / "splits" / "test.csv")
    normalization_stats = load_or_compute_input_normalization_stats(config, train_rows, metric_dir)
    store = ColdFeatureStore(train_rows + val_rows + test_rows, normalization_stats=normalization_stats)

    tile_rows = build_train_tiles(
        train_rows=train_rows,
        store=store,
        tile_size=int(config.get("tile_size", 16)),
        max_positive_tiles_per_sample=int(config.get("max_positive_tiles_per_sample", 64)),
        min_negative_tiles_per_sample=int(config.get("min_negative_tiles_per_sample", 4)),
        neg_pos_ratio=float(config.get("negative_to_positive_ratio", 2.0)),
        positive_tile_placement=str(config.get("positive_tile_placement", "center")),
        rng=rng,
    )
    train_ds = TrainTileDataset(tile_rows=tile_rows, store=store, augment_flip=bool(config.get("augment_flip", True)), seed=int(config.get("seed", 7)))
    val_ds = FullMapDataset(val_rows, store)
    test_ds = FullMapDataset(test_rows, store)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_loader = DataLoader(train_ds, batch_size=int(config.get("batch_size", 64)), shuffle=True, num_workers=0, pin_memory=device.type == "cuda")
    val_loader = DataLoader(val_ds, batch_size=int(config.get("eval_batch_size", 8)), shuffle=False, num_workers=0, pin_memory=device.type == "cuda")
    test_loader = DataLoader(test_ds, batch_size=int(config.get("eval_batch_size", 8)), shuffle=False, num_workers=0, pin_memory=device.type == "cuda")

    positive_rate_source = str(config.get("positive_rate_source", "full_map"))
    pos_weight_source = str(config.get("pos_weight_source", "full_map"))
    metric_reference_positive_rate = full_map_stats(train_rows, store)["positive_rate"]
    positive_rate_stats = select_stats(positive_rate_source, train_rows, tile_rows, store)
    pos_weight_stats = select_stats(pos_weight_source, train_rows, tile_rows, store)
    train_positive_rate = float(positive_rate_stats["positive_rate"])
    raw_pos_weight = float(pos_weight_stats["raw_pos_weight"])
    pos_weight_cap = float(config.get("pos_weight_cap", 300.0))
    pos_weight = float(min(max(raw_pos_weight, 1.0), pos_weight_cap))
    use_aux_spatial_head = bool(config.get("use_aux_spatial_head", False))
    aux_spatial_radius = int(config.get("aux_spatial_radius", 1))
    aux_positive_rate = pooled_positive_rate(train_rows, store, aux_spatial_radius) if use_aux_spatial_head else None

    in_ch = int(store.get(str(train_rows[0]["sample_id"]))["x"].shape[0])
    model = UNetSmallFlex(
        in_ch=in_ch,
        base=int(config.get("base_channels", 32)),
        dropout=float(config.get("dropout", 0.1)),
        norm_type=str(config.get("norm_type", "group")),
        norm_groups=int(config.get("norm_groups", 8)),
        prior_prob=train_positive_rate if bool(config.get("init_head_bias_from_positive_rate", True)) else None,
        use_aux_spatial_head=use_aux_spatial_head,
        aux_prior_prob=aux_positive_rate if bool(config.get("init_head_bias_from_positive_rate", True)) else None,
    ).to(device)
    init_checkpoint = config.get("init_checkpoint")
    if init_checkpoint:
        checkpoint = torch.load(str(init_checkpoint), map_location="cpu")
        state = checkpoint.get("model", checkpoint)
        if not isinstance(state, dict):
            raise RuntimeError(f"Unexpected init_checkpoint format: {init_checkpoint}")
        model.load_state_dict(state, strict=False)
    criterion = build_loss(config=config, pos_weight=pos_weight, device=device)
    aux_criterion = nn.BCEWithLogitsLoss()
    optimizer = torch.optim.AdamW(model.parameters(), lr=float(config.get("learning_rate", 3e-4)), weight_decay=float(config.get("weight_decay", 1e-4)))
    scaler = torch.amp.GradScaler("cuda", enabled=bool(config.get("amp", True)) and device.type == "cuda")

    thresholds = [float(v) for v in config.get("metric_thresholds", [0.1, 0.2, 0.3, 0.5])]
    topk_area_fractions = [float(v) for v in config.get("topk_area_fractions", [0.01, 0.05, 0.1])]
    fss_radii = [int(v) for v in config.get("fss_radii", [1, 2, 4, 8])]
    n_bins = int(config.get("reliability_bins", 10))
    default_threshold_key = f"{float(config.get('threshold', 0.3)):.4f}"

    summary_seed = {
        "index_root": str(index_root),
        "tile_size": int(config.get("tile_size", 16)),
        "num_train_tiles": len(tile_rows),
        "positive_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "positive")),
        "negative_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "negative")),
        "input_channels": in_ch,
        "positive_tile_placement": str(config.get("positive_tile_placement", "center")),
        "train_positive_rate": train_positive_rate,
        "metric_reference_positive_rate": metric_reference_positive_rate,
        "pos_weight_source": pos_weight_source,
        "positive_rate_source": positive_rate_source,
        "pos_weight": pos_weight,
        "loss_type": str(config.get("loss_type", "bce")),
        "train_target_mode": str(config.get("train_target_mode", "hard")),
        "train_target_radius": int(config.get("train_target_radius", 0)),
        "input_normalization": normalization_stats if normalization_stats else {"enabled": False},
    }
    (metric_dir / "tile_summary.json").write_text(json.dumps(summary_seed, indent=2), encoding="utf-8")

    best_pr = -1.0
    best_state = None
    history: List[Dict[str, float]] = []
    aux_spatial_loss_weight = float(config.get("aux_spatial_loss_weight", 0.0))
    for epoch in range(1, int(config.get("epochs", 30)) + 1):
        model.train()
        epoch_start = time.time()
        train_loss = 0.0
        train_main_loss = 0.0
        train_aux_spatial_loss = 0.0
        train_items = 0
        for batch in train_loader:
            x = torch.nan_to_num(batch["x"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
            y = torch.nan_to_num(batch["y"], nan=0.0, posinf=0.0, neginf=0.0).to(device, non_blocking=True)
            train_target = transform_train_target(y.float(), config)
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=bool(config.get("amp", True)) and device.type == "cuda"):
                logits, aux_logits = maybe_unpack_logits(model(x, return_aux=use_aux_spatial_head))
                main_loss = criterion(logits.float(), train_target.float())
            aux_loss = torch.zeros((), device=device, dtype=torch.float32)
            if aux_logits is not None and aux_spatial_loss_weight > 0.0:
                aux_target = transform_spatial_aux_target(y.float(), aux_spatial_radius)
                aux_loss = aux_criterion(aux_logits.float(), aux_target)
            loss = main_loss + aux_spatial_loss_weight * aux_loss
            if not torch.isfinite(loss):
                raise RuntimeError("Non-finite tiled cold-start loss detected.")
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
            train_loss += float(loss.item()) * x.size(0)
            train_main_loss += float(main_loss.item()) * x.size(0)
            train_aux_spatial_loss += float(aux_loss.item()) * x.size(0)
            train_items += int(x.size(0))

        metrics = evaluate(
            model=model,
            loader=val_loader,
            device=device,
            thresholds=thresholds,
            topk_area_fractions=topk_area_fractions,
            fss_radii=fss_radii,
            n_bins=n_bins,
            reference_positive_rate=metric_reference_positive_rate,
            criterion=criterion,
            amp=bool(config.get("amp", True)),
        )
        threshold_metrics = metrics["threshold_metrics"][default_threshold_key]
        record = {
            "epoch": epoch,
            "train_loss": train_loss / max(train_items, 1),
            "train_main_loss": train_main_loss / max(train_items, 1),
            "train_aux_spatial_loss": train_aux_spatial_loss / max(train_items, 1),
            "val_loss": metrics["loss"],
            "val_pr_auc": metrics["pr_auc"],
            "val_auroc": metrics["auroc"],
            "val_brier": metrics["brier"],
            "val_brier_skill_score": metrics["brier_skill_score"],
            "val_ece": metrics["ece"],
            "val_positive_rate": metrics["positive_rate"],
            "val_precision": threshold_metrics["precision"],
            "val_positive_recall": threshold_metrics["recall"],
            "val_far": threshold_metrics["far"],
            "val_csi": threshold_metrics["csi"],
            "val_f1": threshold_metrics["f1"],
            "val_f2": threshold_metrics["f2"],
            "val_frequency_bias": threshold_metrics["frequency_bias"],
            "val_threshold_metrics": metrics["threshold_metrics"],
            "val_topk_area_metrics": metrics["topk_area_metrics"],
            "val_fss": metrics["fss"],
            "minutes": (time.time() - epoch_start) / 60.0,
        }
        history.append(record)
        (metric_dir / "history.json").write_text(json.dumps(history, indent=2), encoding="utf-8")
        save_checkpoint(ckpt_dir / "latest.pt", model, optimizer, epoch, record, config)
        if record["val_pr_auc"] > best_pr:
            best_pr = record["val_pr_auc"]
            best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()}
            torch.save({"model": best_state, "epoch": epoch, "metrics": record, "config": config}, ckpt_dir / "best_firms_prauc.pt")
        print(json.dumps(record), flush=True)

    if best_state is None:
        best_state = {k: v.detach().cpu() for k, v in model.state_dict().items()}
    model.load_state_dict(best_state, strict=True)
    val_best = evaluate(
        model=model,
        loader=val_loader,
        device=device,
        thresholds=thresholds,
        topk_area_fractions=topk_area_fractions,
        fss_radii=fss_radii,
        n_bins=n_bins,
        reference_positive_rate=metric_reference_positive_rate,
        criterion=criterion,
        amp=bool(config.get("amp", True)),
    )
    test_best = evaluate(
        model=model,
        loader=test_loader,
        device=device,
        thresholds=thresholds,
        topk_area_fractions=topk_area_fractions,
        fss_radii=fss_radii,
        n_bins=n_bins,
        reference_positive_rate=metric_reference_positive_rate,
        criterion=criterion,
        amp=bool(config.get("amp", True)),
    )
    summary = {
        "run_name": args.run_name,
        "device": str(device),
        "index_root": str(index_root),
        "input_channels": in_ch,
        "tile_size": int(config.get("tile_size", 16)),
        "num_train_tiles": len(tile_rows),
        "positive_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "positive")),
        "negative_train_tiles": int(sum(1 for row in tile_rows if row["tile_type"] == "negative")),
        "positive_tile_placement": str(config.get("positive_tile_placement", "center")),
        "train_positive_rate": train_positive_rate,
        "metric_reference_positive_rate": metric_reference_positive_rate,
        "pos_weight": pos_weight,
        "use_aux_spatial_head": use_aux_spatial_head,
        "aux_spatial_radius": aux_spatial_radius,
        "aux_spatial_loss_weight": aux_spatial_loss_weight,
        "aux_positive_rate": aux_positive_rate,
        "init_checkpoint": str(init_checkpoint) if init_checkpoint else "",
        "input_normalization": normalization_stats if normalization_stats else {"enabled": False},
        "best_val_pr_auc": float(val_best["pr_auc"]),
        "best_val_auroc": float(val_best["auroc"]),
        "best_test_pr_auc": float(test_best["pr_auc"]),
        "best_test_auroc": float(test_best["auroc"]),
        "best_test_brier": float(test_best["brier"]),
        "best_test_ece": float(test_best["ece"]),
        "best_test_topk_area_metrics": test_best["topk_area_metrics"],
        "best_test_threshold_metrics": test_best["threshold_metrics"],
        "best_test_fss": test_best["fss"],
        "epochs": int(config.get("epochs", 30)),
    }
    (metric_dir / "run_summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
    print(json.dumps(summary, indent=2), flush=True)


if __name__ == "__main__":
    main()