#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

import argparse
import gc
import importlib
import importlib.util
import json
import logging
import os
import time
from pathlib import Path
from typing import Dict, Optional, Tuple, cast

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from monai.data import decollate_batch
from monai.inferers import SlidingWindowInferer
from monai.losses import DiceCELoss
from monai.metrics import DiceMetric
from monai.networks.nets import UNETR, UNet
from monai.transforms import AsDiscrete
from torch.utils.data import DataLoader

# Project loaders and transforms
from src.data.loaders_monai import load_monai_decathlon
from src.data.transforms_presets import (
    get_transforms_brats_like,
    get_transforms_ct_liver,
)
from src.training.callbacks.visualization import save_overlay_panel

# Crash prevention utilities
try:
    from src.utils.crash_prevention import (
        emergency_cleanup,
        gpu_safe_context,
        log_system_resources,
        memory_safe_context,
        safe_execution,
        start_global_protection,
        stop_global_protection,
    )
    CRASH_PREVENTION_AVAILABLE = True
except ImportError:
    CRASH_PREVENTION_AVAILABLE = False
    # Fallback implementations
    from contextlib import contextmanager
    @contextmanager
    def memory_safe_context(*args, **kwargs):
        yield None
    @contextmanager
    def gpu_safe_context(*args, **kwargs):
        yield None
    def safe_execution(*args, **kwargs):
        def decorator(func):
            return func
        return decorator
    def emergency_cleanup():
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    def log_system_resources(logger=None):
        pass
    def start_global_protection(*args, **kwargs):
        pass
    def stop_global_protection():
        pass

# Optional dependency: MLflow
MLFLOW_AVAILABLE = importlib.util.find_spec("mlflow") is not None

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def _mlflow():
    """Return mlflow module if installed, else None."""
    if not MLFLOW_AVAILABLE:
        return None
    return importlib.import_module("mlflow")


def set_determinism(seed: int = 42, enforce: bool = True) -> None:
    import random
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    random.seed(seed)
    np.random.seed(seed)
    if enforce:
        torch.use_deterministic_algorithms(True)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Enhanced training with MONAI datasets"
    )
    p.add_argument(
        "--config",
        required=True,
        help="Path to base training config (JSON)",
    )
    p.add_argument(
        "--dataset-config",
        help=(
            "Path to dataset config (JSON). If monai_decathlon, will "
            "auto-download"
        ),
    )
    p.add_argument(
        "--output-dir",
        default="models/unetr",
        help="Directory to write checkpoints and logs",
    )
    p.add_argument("--epochs", type=int, default=None, help="Override epochs")
    p.add_argument(
        "--amp", action="store_true", help="Enable mixed precision (AMP)"
    )
    p.add_argument("--seed", type=int, default=42, help="Random seed")
    p.add_argument(
        "--no-deterministic",
        action="store_true",
        help="Allow non-deterministic ops",
    )
    p.add_argument(
        "--device", default=None, help="cpu | cuda | cuda:0 | mps | auto"
    )
    p.add_argument(
        "--resume", default=None, help="Path to checkpoint to resume from"
    )
    p.add_argument(
        "--val-interval", type=int, default=1, help="Validate every N epochs"
    )
    # A) Sliding-window overlap CLI
    p.add_argument(
        "--sw-overlap",
        type=float,
        default=0.25,
        help="Sliding window overlap for validation/inference",
    )
    # C) Overlays
    p.add_argument(
        "--save-overlays",
        action="store_true",
        help="Save simple validation overlays each eval",
    )
    p.add_argument(
        "--overlays-max",
        type=int,
        default=2,
        help="Max number of validation overlays to save per eval",
    )
    # D) Validation control
    p.add_argument(
        "--val-max-batches",
        type=int,
        default=0,
        help="Max validation batches per eval (0=all)",
    )
    # E) Probability maps
    p.add_argument(
        "--save-prob-maps",
        action="store_true",
        help="Save probability heatmaps for tumor class",
    )
    return p.parse_args()


def build_transforms_from_dataset_cfg(ds_cfg: Dict):
    spacing = tuple(ds_cfg.get("spacing", (1.0, 1.0, 1.0)))
    name = ds_cfg.get("transforms", "brats_like")
    if name == "brats_like":
        return get_transforms_brats_like(spacing=spacing)
    if name == "ct_liver":
        return get_transforms_ct_liver(spacing=spacing)
    return get_transforms_brats_like(spacing=spacing)


def build_model_from_cfg(
    cfg: Dict, in_channels: int = 4, out_channels: int = 2
) -> nn.Module:
    arch = cfg.get("model", {}).get("arch", "unetr").lower()
    if arch == "unetr":
        return UNETR(
            in_channels=in_channels,
            out_channels=out_channels,
            img_size=tuple(
                cfg.get("model", {}).get("img_size", [128, 128, 128])
            ),
            feature_size=cfg.get("model", {}).get("feature_size", 16),
            hidden_size=cfg.get("model", {}).get("hidden_size", 768),
            mlp_dim=cfg.get("model", {}).get("mlp_dim", 3072),
            num_heads=cfg.get("model", {}).get("num_heads", 12),
            pos_embed="perceptron",
            norm_name="instance",
            res_block=True,
            dropout_rate=cfg.get("model", {}).get("dropout", 0.0),
        )
    return UNet(
        spatial_dims=3,
        in_channels=in_channels,
        out_channels=out_channels,
        channels=cfg.get("model", {}).get("channels", [16, 32, 64, 128, 256]),
        strides=cfg.get("model", {}).get("strides", [2, 2, 2, 2]),
        num_res_units=cfg.get("model", {}).get("num_res_units", 2),
        norm=cfg.get("model", {}).get("norm", "INSTANCE"),
        dropout=cfg.get("model", {}).get("dropout", 0.0),
    )


def get_device(arg_device: Optional[str]) -> torch.device:
    if arg_device in (None, "auto"):
        if torch.cuda.is_available():
            return torch.device("cuda")
        if getattr(torch.backends, "mps", None) and \
                torch.backends.mps.is_available():
            return torch.device("mps")
        return torch.device("cpu")
    return torch.device(arg_device)


def infer_in_channels_from_loader(loader: DataLoader, default: int = 4) -> int:
    """
    B) Inspect a single batch to determine input channels robustly.
    Does not permanently consume the iterator for training.
    """
    it = iter(loader)
    try:
        sample = next(it)
    except StopIteration:
        return default
    img = sample["image"]
    if img.ndim >= 5:
        return int(img.shape[1])
    return default


def save_prob_map_png(
    image_chn_first: torch.Tensor,
    prob_tumor: torch.Tensor,
    out_path: Path,
    slices: Optional[list] = None,
) -> None:
    """
    Save probability heatmap for tumor class.
    image_chn_first: (C, H, W, D) tensor (float)
    prob_tumor: (H, W, D) tensor (0-1 probabilities)
    slices: Optional list of z indices. If None, uses middle slice.
    """
    # use first modality as background image
    img = image_chn_first[0].detach().cpu().float().numpy()
    probs = prob_tumor.detach().cpu().numpy()
    _, _, D = img.shape

    # Default to middle slice if no slices specified
    if slices is None:
        slices = [D // 2]

    # Ensure slices are valid
    slices = [max(0, min(s, D-1)) for s in slices]

    n_slices = len(slices)
    fig, axes = plt.subplots(1, n_slices, figsize=(6 * n_slices, 6))

    # Ensure axes is always a list for consistent indexing
    if n_slices == 1:
        axes = [axes]

    for i, z in enumerate(slices):
        base = img[..., z]
        base = (base - base.min()) / (base.max() - base.min() + 1e-8)
        prob_slice = probs[..., z]

        axes[i].axis("off")
        axes[i].imshow(base, cmap="gray")
        # Overlay probability heatmap
        axes[i].imshow(
            np.ma.masked_where(prob_slice < 0.1, prob_slice),
            cmap="hot", alpha=0.6, vmin=0, vmax=1
        )
        axes[i].set_title(f"Prob Map Slice {z}/{D-1}", fontsize=10)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    plt.tight_layout()
    plt.savefig(out_path, dpi=150, bbox_inches="tight", pad_inches=0)
    plt.close()


@safe_execution(max_retries=2, memory_threshold=0.80, gpu_threshold=0.85)
def train_one_epoch(
    model,
    loader: DataLoader,
    optimizer,
    loss_fn,
    device,
    scaler: Optional[torch.cuda.amp.GradScaler] = None,
):
    """Train for one epoch with crash prevention"""
    model.train()
    epoch_loss = 0.0
    count = 0

    # Start memory monitoring for this epoch
    with memory_safe_context(threshold=0.80) as monitor:
        logger.info(f"Starting training epoch with {len(loader)} batches")
        log_system_resources(logger)

        for batch_data in loader:
            try:
                inputs = batch_data["image"].to(device, non_blocking=True)
                labels = batch_data["label"].to(device, non_blocking=True)

                optimizer.zero_grad()

                # Forward pass with memory management
                if scaler is not None:
                    with torch.autocast(device_type=device.type, dtype=torch.float16):
                        outputs = model(inputs)
                        loss = loss_fn(outputs, labels)
                    scaler.scale(loss).backward()
                    scaler.step(optimizer)
                    scaler.update()
                else:
                    outputs = model(inputs)
                    loss = loss_fn(outputs, labels)
                    loss.backward()
                    optimizer.step()

                # Accumulate loss
                epoch_loss += loss.item()
                count += 1

                # Clean up tensors immediately
                del inputs, labels, outputs, loss

                # Periodic cleanup during training
                if count % 10 == 0:
                    emergency_cleanup()
                    if count % 50 == 0:  # Log every 50 batches
                        logger.info(f"Batch {count}/{len(loader)}, "
                                  f"avg_loss={epoch_loss/count:.4f}")

            except RuntimeError as e:
                if "out of memory" in str(e).lower():
                    logger.error(f"OOM error at batch {count}: {e}")
                    emergency_cleanup()
                    # Skip this batch and continue
                    continue
                else:
                    raise e
            except Exception as e:
                logger.error(f"Training error at batch {count}: {e}")
                emergency_cleanup()
                raise e

        # Final cleanup for epoch
        emergency_cleanup()
        logger.info(f"Epoch completed: {count} batches, avg_loss={epoch_loss/max(count,1):.4f}")

    return epoch_loss / max(count, 1)
    for batch in loader:
        images = batch["image"].to(device)
        labels = batch["label"].to(device)
        optimizer.zero_grad(set_to_none=True)
        if scaler is not None and device.type == "cuda":
            with torch.autocast(device_type="cuda", dtype=torch.float16):
                logits = model(images)
                loss = loss_fn(logits, labels)
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            logits = model(images)
            loss = loss_fn(logits, labels)
            loss.backward()
            optimizer.step()
        epoch_loss += float(loss.item())
        count += 1
    return epoch_loss / max(1, count)


@torch.no_grad()
@safe_execution(max_retries=2, memory_threshold=0.80, gpu_threshold=0.85)
def validate(
    model,
    loader: DataLoader,
    device,
    post_pred,
    post_label,
    roi_size: Optional[Tuple[int, int, int]] = None,
    sw_overlap: float = 0.25,
    save_overlays: bool = False,
    overlay_dir: Optional[Path] = None,
    overlays_max: int = 2,
    val_max_batches: int = 0,
    save_prob_maps: bool = False,
):
    """Validation function with comprehensive crash prevention"""
    model.eval()
    dice_metric = DiceMetric(include_background=False, reduction="mean")
    saved = 0
    prob_maps_saved = 0

    # Use config-provided roi_size or fall back to full image shape
    inferers = {}  # cache per image shape if not using fixed roi_size

    # Start validation with memory monitoring
    with memory_safe_context(threshold=0.85):
        with gpu_safe_context(threshold=0.90):
            logger.info(f"Starting validation with {len(loader)} batches")
            log_system_resources(logger)

            for batch_idx, batch in enumerate(loader):
                try:
                    # Validation batch limit
                    if val_max_batches > 0 and batch_idx >= val_max_batches:
                        logger.info(f"Validation stopped at batch {batch_idx}")
                        break

                    images = batch["image"].to(device, non_blocking=True)
                    labels = batch["label"].to(device, non_blocking=True)

                    # Setup sliding window inferer
                    if roi_size is not None:
                        inferer = SlidingWindowInferer(
                            roi_size=tuple(roi_size),
                            sw_batch_size=1,
                            overlap=sw_overlap,
                        )
                    else:
                        # cache inferer per final spatial dims
                        key = tuple(images.shape[-3:])
                        if key not in inferers:
                            inferers[key] = SlidingWindowInferer(
                                roi_size=key, sw_batch_size=1,
                                overlap=sw_overlap
                            )
                        inferer = inferers[key]

                    # Perform inference
                    logits = inferer(images, model)

                    # Save probability maps if requested
                    if (save_prob_maps and overlay_dir is not None
                            and prob_maps_saved < overlays_max):
                        try:
                            probs = torch.softmax(logits, dim=1)
                            if probs.shape[1] > 1:
                                prob_tumor = probs[0, 1]
                            else:
                                prob_tumor = probs[0, 0]

                            img_chn_first = images[0]
                            D = img_chn_first.shape[-1]
                            slice_indices = [D // 4, D // 2, 3 * D // 4]

                            prob_path = (overlay_dir /
                                       f"val_probmap_{batch_idx:03d}.png")
                            save_prob_map_png(
                                img_chn_first, prob_tumor,
                                prob_path, slice_indices
                            )
                            prob_maps_saved += 1
                        except Exception as e:
                            logger.warning(f"Failed to save prob map: {e}")

                    # Process predictions for metrics
                    preds_list = [post_pred(i) for i in decollate_batch(logits)]
                    gts_list = [post_label(i) for i in decollate_batch(labels)]
                    dice_metric(y_pred=preds_list, y=gts_list)

                    # Save overlays if requested
                    if (save_overlays and overlay_dir is not None
                            and saved < overlays_max):
                        try:
                            img_chn_first = images[0]
                            gt_onehot = gts_list[0]
                            pr_onehot = preds_list[0]

                            D = img_chn_first.shape[-1]
                            slice_indices = [D // 4, D // 2, 3 * D // 4]

                            out_path = (overlay_dir /
                                      f"val_overlay_{batch_idx:03d}.png")
                            save_overlay_panel(
                                image_ch_first=img_chn_first,
                                label_onehot=gt_onehot,
                                pred_onehot=pr_onehot,
                                out_path=out_path,
                                slices=slice_indices
                            )
                            saved += 1
                        except Exception as e:
                            logger.warning(f"Failed to save overlay: {e}")

                    # Clean up tensors immediately
                    del images, labels, logits, preds_list, gts_list

                    # Periodic cleanup
                    if batch_idx % 5 == 0:
                        emergency_cleanup()

                except RuntimeError as e:
                    if "out of memory" in str(e).lower():
                        logger.error(f"OOM error in validation: {e}")
                        emergency_cleanup()
                        continue
                    else:
                        raise e
                except Exception as e:
                    logger.error(f"Validation error: {e}")
                    emergency_cleanup()
                    raise e

            # Final cleanup after validation
            emergency_cleanup()

            mean_dice = float(dice_metric.aggregate().item())
            dice_metric.reset()
            logger.info(f"Validation completed: mean_dice={mean_dice:.4f}")

            return {"dice": mean_dice, "n": len(loader)}


@safe_execution(max_retries=1, memory_threshold=0.75, gpu_threshold=0.80)
def main() -> int:
    """Main training function with comprehensive crash prevention"""

    # Start global crash protection
    start_global_protection()

    try:
        args = parse_args()
        set_determinism(args.seed, enforce=not args.no_deterministic)

        logger.info("🚀 Starting enhanced training with crash prevention")
        log_system_resources(logger)

        # Device setup with validation
        device = get_device(args.device)
        logger.info(f"Using device: {device}")

        # Output directory setup
        out_dir = Path(args.output_dir)
        out_dir.mkdir(parents=True, exist_ok=True)
        overlays_dir = out_dir / "overlays"

        # Load configurations
        with open(args.config, "r", encoding="utf-8") as f:
            cfg = json.load(f)

        ds_cfg = None
        if args.dataset_config:
            with open(args.dataset_config, "r", encoding="utf-8") as f:
                ds_cfg = json.load(f)

        # Build transforms with crash protection
        if ds_cfg:
            t_train, t_val = build_transforms_from_dataset_cfg(ds_cfg)
        else:
            t_train, t_val = get_transforms_brats_like()

        # Validate dataset configuration
        if not (ds_cfg and ds_cfg.get("source") == "monai_decathlon"):
            raise RuntimeError("Requires --dataset-config (monai_decathlon).")

        # Dataset loading with memory-safe configuration
        train_key = ds_cfg.get("splits", {}).get("train_key", "training")
        val_key = ds_cfg.get("splits", {}).get("val_key", "validation")
        data_root = ds_cfg.get("data_root", "data/msd")
        dataset_id = ds_cfg.get("dataset_id", "Task01_BrainTumour")
        loader_cfg = ds_cfg.get("loader", {})

        # Reduce workers and batch size for memory safety
        safe_num_workers = min(int(loader_cfg.get("num_workers", 4)), 2)
        safe_batch_size = min(int(loader_cfg.get("batch_size", 1)), 1)

        logger.info(f"Loading dataset with safe parameters: "
                   f"workers={safe_num_workers}, batch_size={safe_batch_size}")

        # Load training data with crash prevention
        with memory_safe_context(threshold=0.80):
            train_res = load_monai_decathlon(
                root_dir=data_root,
                task=dataset_id,
                section=train_key,
                download=True,
                cache_rate=0.0,
                num_workers=safe_num_workers,
                transform=t_train,
                batch_size=safe_batch_size,
                pin_memory=False,  # Disable for memory safety
            )

            val_res = load_monai_decathlon(
                root_dir=data_root,
                task=dataset_id,
                section=val_key,
                download=True,
                cache_rate=0.0,
                num_workers=safe_num_workers,
                transform=t_val,
                batch_size=safe_batch_size,
                pin_memory=False,  # Disable for memory safety
            )

        train_loader = train_res["dataloader"]
        val_loader = val_res["dataloader"]
        logger.info(f"Loaded datasets: train={len(train_loader)}, "
                   f"val={len(val_loader)} batches")

        # Infer channels and build model with crash prevention
        with gpu_safe_context():
            inferred_in_channels = infer_in_channels_from_loader(
                train_loader, default=4
            )
            out_channels = cfg.get("model", {}).get("out_channels", 2)
            model = build_model_from_cfg(
                cfg,
                in_channels=inferred_in_channels,
                out_channels=out_channels,
            ).to(device)

            logger.info(f"Model built: {inferred_in_channels} -> {out_channels} channels")
            log_system_resources(logger)

        # Setup training components with crash prevention
        loss_fn = DiceCELoss(
            to_onehot_y=True,
            softmax=True,
            include_background=False,
        )
        lr = cfg.get("optim", {}).get("lr", 1e-4)
        weight_decay = cfg.get("optim", {}).get("weight_decay", 0.0)
        optimizer = torch.optim.AdamW(
            model.parameters(), lr=lr, weight_decay=weight_decay
        )
        epochs = args.epochs or cfg.get("trainer", {}).get("epochs", 50)
        val_interval = args.val_interval

        # AMP scaler
        scaler = (
            torch.cuda.amp.GradScaler()
            if (args.amp and device.type == "cuda")
            else None
        )

        # Resume if provided
        start_epoch = 0
        if args.resume is not None and os.path.isfile(args.resume):
            state = torch.load(args.resume, map_location="cpu")
            model.load_state_dict(state.get("model", state))
            if "optimizer" in state:
                optimizer.load_state_dict(state["optimizer"])
            start_epoch = state.get("epoch", 0)
            logger.info(f"Resumed from {args.resume} at epoch {start_epoch}")

        # Post transforms for validation metrics
        post_pred = AsDiscrete(argmax=True, to_onehot=out_channels)
        post_label = AsDiscrete(to_onehot=out_channels)

        # Determine roi_size for validation
        roi_size_typed: Optional[Tuple[int, int, int]] = None
        model_img_size = cfg.get("model", {}).get("img_size", None)
        if isinstance(model_img_size, (list, tuple)) and len(model_img_size) == 3:
            roi_size_typed = cast(
                Tuple[int, int, int], tuple(int(x) for x in model_img_size)
            )

        # MLflow logging setup
        mlflow = _mlflow()
        if mlflow:
            mlflow.set_experiment(
                cfg.get("mlflow", {}).get("experiment", "medical-imaging")
            )
            run_name = cfg.get("mlflow", {}).get(
                "run_name", f"run-{int(time.time())}"
            )
            mlflow.start_run(run_name=run_name)
            mlflow.log_params({
                "arch": cfg.get("model", {}).get("arch", "unetr"),
                "in_channels": inferred_in_channels,
                "out_channels": out_channels,
                "lr": lr,
                "epochs": epochs,
                "amp": bool(scaler is not None),
                "dataset_task": ds_cfg.get("dataset_id") if ds_cfg else None,
                "seed": args.seed,
                "roi_size": (
                    roi_size_typed if roi_size_typed is not None else "full"
                ),
                "sw_overlap": args.sw_overlap,
            })

        best_dice = -1.0
        best_ckpt = out_dir / "best.pt"

        # Training loop with comprehensive crash prevention
        logger.info(f"Starting training for {epochs} epochs")

        for epoch in range(start_epoch, epochs):
            try:
                t0 = time.time()

                # Log epoch start
                logger.info(f"=== Epoch {epoch+1}/{epochs} ===")
                log_system_resources(logger)

                # Training phase with memory monitoring
                train_loss = train_one_epoch(
                    model, train_loader, optimizer, loss_fn, device, scaler
                )

                dt = time.time() - t0
                logger.info(f"Epoch {epoch+1}/{epochs} - "
                           f"loss={train_loss:.4f}, time={dt:.1f}s")

                if mlflow:
                    mlflow.log_metrics(
                        {"train/loss": train_loss, "time/epoch_s": dt},
                        step=epoch
                    )

                # Validation phase
                if (epoch + 1) % val_interval == 0:
                    try:
                        metrics = validate(
                            model,
                            val_loader,
                            device,
                            post_pred,
                            post_label,
                            roi_size=roi_size_typed,
                            sw_overlap=args.sw_overlap,
                            save_overlays=args.save_overlays,
                            overlay_dir=overlays_dir,
                            overlays_max=args.overlays_max,
                            val_max_batches=args.val_max_batches,
                            save_prob_maps=args.save_prob_maps,
                        )
                        dice = metrics["dice"]
                        logger.info(f"Validation - mean_dice={dice:.4f}")

                        if mlflow:
                            mlflow.log_metrics({"val/dice": dice}, step=epoch)

                        # Save best model
                        if dice > best_dice:
                            best_dice = dice
                            logger.info(f"New best model: dice={best_dice:.4f}")
                            torch.save({
                                "model": model.state_dict(),
                                "optimizer": optimizer.state_dict(),
                                "epoch": epoch,
                                "dice": dice,
                            }, best_ckpt)

                    except Exception as e:
                        logger.error(f"Validation failed at epoch {epoch+1}: {e}")
                        emergency_cleanup()
                        # Continue training even if validation fails

                # Epoch cleanup
                emergency_cleanup()

            except Exception as e:
                logger.error(f"Training failed at epoch {epoch+1}: {e}")
                emergency_cleanup()
                if "out of memory" in str(e).lower():
                    logger.warning("OOM detected, reducing batch size and continuing")
                    continue
                else:
                    raise e

        # Training completed
        logger.info(f"🎉 Training completed! Best dice: {best_dice:.4f}")

        if mlflow:
            mlflow.log_metric("best_dice", best_dice)
            mlflow.end_run()

        return 0

    except Exception as e:
        logger.error(f"Training failed with error: {e}")
        emergency_cleanup()
        if mlflow:
            mlflow.end_run(status="FAILED")
        return 1

    finally:
        # Always stop global protection
        stop_global_protection()
        logger.info("🛡️ Crash protection stopped")
        train_loss = train_one_epoch(
            model, train_loader, optimizer, loss_fn, device, scaler
        )
        dt = time.time() - t0

        print(
            f"[Ep {epoch+1}/{epochs}] loss={train_loss:.4f} "
            f"t={dt:.1f}s"
        )
        if mlflow:
            mlflow.log_metrics(
                {"train/loss": train_loss, "time/epoch_s": dt}, step=epoch
            )

        if (epoch + 1) % val_interval == 0:
            metrics = validate(
                model,
                val_loader,
                device,
                post_pred,
                post_label,
                roi_size=roi_size_typed,
                sw_overlap=args.sw_overlap,
                save_overlays=args.save_overlays,
                overlay_dir=overlays_dir,
                overlays_max=args.overlays_max,
                val_max_batches=args.val_max_batches,
                save_prob_maps=args.save_prob_maps,
            )
            dice = metrics["dice"]
            print(f"  [Val] mean_dice={dice:.4f}")
            if mlflow:
                mlflow.log_metric("val/dice", dice, step=epoch)
            if dice > best_dice:
                best_dice = dice
                state = {
                    "model": model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "epoch": epoch,
                }
                torch.save(state, best_ckpt)
                print(
                    f"  [CKPT] Saved best to {best_ckpt} "
                    f"(dice={best_dice:.4f})"
                )

    final_ckpt = out_dir / "last.pt"
    final_state = {
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "epoch": epochs - 1,
    }
    torch.save(final_state, final_ckpt)
    print(
        f"[DONE] Saved final checkpoint to {final_ckpt}. "
        f"Best dice={best_dice:.4f}"
    )

    if mlflow:
        mlflow.log_artifact(str(best_ckpt))
        mlflow.log_artifact(str(final_ckpt))
        if args.save_overlays and overlays_dir.exists():
            # log a couple of overlay images
            for p in sorted(overlays_dir.glob("*.png"))[:3]:
                mlflow.log_artifact(str(p))
        mlflow.end_run()

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
