utils/unified_trainer.py

from __future__ import annotations

import csv
import json
import time
from copy import deepcopy
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from datasets.cd_dataset import CDDataset
from utils.config_loader import load_dataset_config, load_model_config
from utils.dataset_cache import apply_dataloader_cli_overrides, dataloader_kwargs, dataloader_policy_lines, dataset_runtime_summary, print_dataloader_policy
from utils.gpu_utils import print_gpu_diagnostics, resolve_gpu
from utils.metrics import BinaryMetrics, BoundaryMetrics, normalize_binary_prediction
from utils.model_adapters import BaseModelAdapter, get_model_adapter
from utils.profiling import GpuProfiler, ProfilingUnavailable, count_flops, count_parameters
from utils.qualitative import (
    denormalize,
    manifest_ids,
    rank_for_sample,
    safe_sample_id,
    save_binary_prediction,
    save_probability_map,
    save_visual_panel,
    select_or_load_manifest,
)
from utils.results_writer import append_to_comparison_table, save_metrics


ROOT = Path(__file__).resolve().parents[1]


def _json_dump(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, sort_keys=True)


def _append_jsonl(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as f:
        f.write(json.dumps(payload, sort_keys=True) + "\n")


def _threshold_sweep_enabled(dataset_cfg: dict, split: str) -> bool:
    eval_cfg = dataset_cfg.get("eval", {})
    if split == "val":
        return bool(eval_cfg.get("sweep_val_threshold", True))
    if split == "test":
        return bool(eval_cfg.get("sweep_test_threshold", False))
    return bool(eval_cfg.get("sweep_threshold", True))


def _thresholds(dataset_cfg: dict, split: str) -> list[float]:
    eval_cfg = dataset_cfg.get("eval", {})
    if not _threshold_sweep_enabled(dataset_cfg, split):
        return [float(eval_cfg.get("threshold", 0.5))]
    if "thresholds" in eval_cfg:
        return [float(x) for x in eval_cfg["thresholds"]]
    start = float(eval_cfg.get("threshold_min", 0.05))
    stop = float(eval_cfg.get("threshold_max", 0.95))
    step = float(eval_cfg.get("threshold_step", 0.05))
    values = []
    current = start
    while current <= stop + 1e-9:
        values.append(round(current, 4))
        current += step
    return values


def _cfg_with_threshold(dataset_cfg: dict, threshold: float) -> dict:
    cfg = deepcopy(dataset_cfg)
    cfg.setdefault("eval", {})["threshold"] = float(threshold)
    return cfg


def _last_tensor(raw_output: Any) -> torch.Tensor | None:
    value = raw_output[-1] if isinstance(raw_output, (list, tuple)) and raw_output else raw_output
    return value if torch.is_tensor(value) else None


def _has_nonfinite_tensor(value: Any) -> bool:
    if torch.is_tensor(value):
        return not bool(torch.isfinite(value.detach()).all().item())
    if isinstance(value, (list, tuple)):
        return any(_has_nonfinite_tensor(item) for item in value)
    if isinstance(value, dict):
        return any(_has_nonfinite_tensor(item) for item in value.values())
    return False


def _batch_stats(raw_output: Any, normalized, mask: torch.Tensor) -> dict[str, float | None]:
    binary = normalized.binary.float()
    target = (mask > 0).float()
    score = normalized.score
    raw_tensor = _last_tensor(raw_output)
    return {
        "gt_positive_ratio": float(target.mean().item()),
        "pred_positive_ratio": float(binary.mean().item()),
        "mean_prob": float(score.mean().item()) if score is not None else None,
        "logit_mean": float(raw_tensor.detach().float().mean().item()) if raw_tensor is not None else None,
        "logit_std": float(raw_tensor.detach().float().std().item()) if raw_tensor is not None and raw_tensor.numel() > 1 else None,
    }


def _loader(dataset_cfg: dict, split: str, batch_size: int, shuffle: bool, drop_last: bool = False) -> DataLoader:
    ds = CDDataset(dataset_cfg["data_root"], split, cfg=dataset_cfg, return_format="tuple")
    return DataLoader(
        ds,
        batch_size=batch_size,
        shuffle=shuffle,
        **dataloader_kwargs(dataset_cfg, torch.cuda.is_available()),
        drop_last=drop_last,
    )


def _write_trajectory(metrics_dir: Path, trajectory: list[dict]) -> None:
    _json_dump(metrics_dir / "trajectory.json", trajectory)
    if not trajectory:
        return
    columns = [
        "epoch",
        "train_loss",
        "val_threshold",
        "val_f1",
        "val_iou",
        "val_miou",
        "val_precision",
        "val_recall",
        "val_oa",
        "val_bf1",
        "val_gt_positive_ratio",
        "val_pred_positive_ratio",
        "val_mean_prob",
        "diagnostic_test_f1",
        "diagnostic_test_iou",
        "val_test_f1_gap",
        "val_test_iou_gap",
        "best_val_f1_so_far",
        "best_val_epoch_or_iter_so_far",
        "checkpoint_path",
    ]
    with (metrics_dir / "trajectory.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=columns, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(trajectory)


def _evaluate_split(
    *,
    model_name: str,
    model: torch.nn.Module,
    adapter: BaseModelAdapter,
    dataset_cfg: dict,
    loader: DataLoader,
    device: torch.device,
    threshold: float,
    split: str,
    out_dir: Path,
    save_outputs: bool,
    checkpoint_path: Path | None = None,
) -> dict:
    cfg = _cfg_with_threshold(dataset_cfg, threshold)
    boundary = BoundaryMetrics(tolerance=int(dataset_cfg.get("eval", {}).get("boundary_tolerance", 2)))
    metrics = BinaryMetrics(threshold=threshold)
    pred_dir = out_dir / "predictions" / split
    prob_dir = out_dir / "predictions" / f"{split}_prob"
    visual_dir = out_dir / "visuals" / "selected_20"
    manifest = select_or_load_manifest(dataset_cfg) if save_outputs and split == "test" else {"samples": []}
    selected = manifest_ids(manifest)
    mean_a = dataset_cfg.get("mean_a", [0.485, 0.456, 0.406])
    std_a = dataset_cfg.get("std_a", [0.229, 0.224, 0.225])
    mean_b = dataset_cfg.get("mean_b", mean_a)
    std_b = dataset_cfg.get("std_b", std_a)

    n_samples = 0
    model_time = 0.0
    end_to_end_start = time.perf_counter()
    stat_sums = {"gt_positive_ratio": 0.0, "pred_positive_ratio": 0.0, "mean_prob": 0.0}
    stat_counts = {"mean_prob": 0}

    model.eval()
    with torch.inference_mode(), GpuProfiler(device=device, required=False) as gpu_profiler:
        for batch in loader:
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            start = time.perf_counter()
            raw = adapter.forward(model, batch, device)
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            elapsed = time.perf_counter() - start
            a, b, mask, names = batch
            normalized = adapter.normalize_output(raw, batch, cfg)
            metrics.update(normalized.metric_tensor, mask)
            boundary.update(normalized.binary, mask)
            n_samples += int(mask.shape[0])
            model_time += elapsed
            stats = _batch_stats(raw, normalized, mask)
            stat_sums["gt_positive_ratio"] += float(stats["gt_positive_ratio"] or 0.0) * int(mask.shape[0])
            stat_sums["pred_positive_ratio"] += float(stats["pred_positive_ratio"] or 0.0) * int(mask.shape[0])
            if stats["mean_prob"] is not None:
                stat_sums["mean_prob"] += float(stats["mean_prob"]) * int(mask.shape[0])
                stat_counts["mean_prob"] += int(mask.shape[0])

            if save_outputs:
                for i, sample_id in enumerate(names):
                    clean_id = safe_sample_id(str(sample_id))
                    pred_i = normalized.binary[i]
                    save_binary_prediction(pred_i, pred_dir / f"{clean_id}_pred.png")
                    prob_i = normalized.score[i] if normalized.score is not None else None
                    if prob_i is not None:
                        save_probability_map(prob_i, prob_dir / f"{clean_id}_prob.png")
                    if split == "test" and str(sample_id) in selected:
                        rank = rank_for_sample(manifest, str(sample_id))
                        save_visual_panel(
                            denormalize(a[i].detach().cpu(), mean_a, std_a),
                            denormalize(b[i].detach().cpu(), mean_b, std_b),
                            mask[i],
                            pred_i,
                            visual_dir / f"{rank:02d}_{clean_id}_panel.png",
                            prob=prob_i,
                        )

    elapsed_total = time.perf_counter() - end_to_end_start
    result = metrics.compute()
    result.update(boundary.compute())
    result.update(gpu_profiler.summary())
    result.update({
        "model": model_name,
        "dataset": dataset_cfg["name"],
        "split": split,
        "threshold": threshold,
        "threshold_mode": adapter.get_threshold_mode(),
        "gt_positive_ratio": stat_sums["gt_positive_ratio"] / max(n_samples, 1),
        "pred_positive_ratio": stat_sums["pred_positive_ratio"] / max(n_samples, 1),
        "mean_prob": stat_sums["mean_prob"] / stat_counts["mean_prob"] if stat_counts["mean_prob"] else None,
        "fps": n_samples / model_time if model_time > 0 else None,
        "fps_model_only": n_samples / model_time if model_time > 0 else None,
        "fps_end_to_end": n_samples / elapsed_total if elapsed_total > 0 else None,
        "sample_count": n_samples,
        "test_sample_count": n_samples if split == "test" else None,
        "checkpoint": str(checkpoint_path) if checkpoint_path else None,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "status": "complete",
    })
    return result


def _metrics_from_cached_predictions(
    *,
    model_name: str,
    dataset_cfg: dict,
    cached: list[tuple[torch.Tensor, torch.Tensor]],
    threshold: float,
    split: str,
    n_samples: int,
    model_time: float,
    elapsed_total: float,
    gpu_summary: dict[str, object],
) -> dict:
    metrics = BinaryMetrics(threshold=threshold)
    boundary = BoundaryMetrics(tolerance=int(dataset_cfg.get("eval", {}).get("boundary_tolerance", 2)))
    stat_sums = {"gt_positive_ratio": 0.0, "pred_positive_ratio": 0.0, "mean_prob": 0.0}
    stat_counts = {"mean_prob": 0}

    for metric_tensor, mask in cached:
        binary, score = normalize_binary_prediction(metric_tensor, threshold=threshold)
        metrics.update(metric_tensor, mask)
        boundary.update(binary, mask)
        batch_size = int(mask.shape[0])
        target = (mask > 0).float()
        stat_sums["gt_positive_ratio"] += float(target.mean().item()) * batch_size
        stat_sums["pred_positive_ratio"] += float(binary.float().mean().item()) * batch_size
        if score is not None:
            stat_sums["mean_prob"] += float(score.float().mean().item()) * batch_size
            stat_counts["mean_prob"] += batch_size

    result = metrics.compute()
    result.update(boundary.compute())
    result.update(gpu_summary)
    result.update({
        "model": model_name,
        "dataset": dataset_cfg["name"],
        "split": split,
        "threshold": threshold,
        "threshold_mode": "threshold",
        "gt_positive_ratio": stat_sums["gt_positive_ratio"] / max(n_samples, 1),
        "pred_positive_ratio": stat_sums["pred_positive_ratio"] / max(n_samples, 1),
        "mean_prob": stat_sums["mean_prob"] / stat_counts["mean_prob"] if stat_counts["mean_prob"] else None,
        "fps": n_samples / model_time if model_time > 0 else None,
        "fps_model_only": n_samples / model_time if model_time > 0 else None,
        "fps_end_to_end": n_samples / elapsed_total if elapsed_total > 0 else None,
        "sample_count": n_samples,
        "test_sample_count": n_samples if split == "test" else None,
        "checkpoint": None,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "status": "complete",
        "sweep_cache": True,
    })
    return result


def _threshold_sweep(
    *,
    model_name: str,
    model: torch.nn.Module,
    adapter: BaseModelAdapter,
    dataset_cfg: dict,
    loader: DataLoader,
    device: torch.device,
    split: str,
    out_dir: Path,
) -> dict:
    if adapter.get_threshold_mode() == "argmax":
        threshold = adapter.get_threshold(dataset_cfg)
        metrics = _evaluate_split(
            model_name=model_name,
            model=model,
            adapter=adapter,
            dataset_cfg=dataset_cfg,
            loader=loader,
            device=device,
            threshold=threshold,
            split=split,
            out_dir=out_dir,
            save_outputs=False,
        )
        return {
            "threshold_mode": "argmax",
            "selected_threshold": threshold,
            "best_f1": metrics["f1"],
            "results": [metrics],
        }
    cached: list[tuple[torch.Tensor, torch.Tensor]] = []
    n_samples = 0
    model_time = 0.0
    end_to_end_start = time.perf_counter()
    model.eval()
    with torch.inference_mode(), GpuProfiler(device=device, required=False) as gpu_profiler:
        for batch in loader:
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            start = time.perf_counter()
            raw = adapter.forward(model, batch, device)
            if device.type == "cuda":
                torch.cuda.synchronize(device)
            model_time += time.perf_counter() - start
            _a, _b, mask, _names = batch
            normalized = adapter.normalize_output(raw, batch, dataset_cfg)
            cached.append((normalized.metric_tensor.detach().cpu(), mask.detach().cpu()))
            n_samples += int(mask.shape[0])
    elapsed_total = time.perf_counter() - end_to_end_start
    gpu_summary = gpu_profiler.summary()

    results = []
    best = None
    for threshold in _thresholds(dataset_cfg, split):
        metrics = _metrics_from_cached_predictions(
            model_name=model_name,
            dataset_cfg=dataset_cfg,
            threshold=threshold,
            split=split,
            cached=cached,
            n_samples=n_samples,
            model_time=model_time,
            elapsed_total=elapsed_total,
            gpu_summary=gpu_summary,
        )
        results.append(metrics)
        if best is None or metrics["f1"] > best["f1"]:
            best = metrics
    assert best is not None
    return {
        "threshold_mode": "threshold",
        "selected_threshold": float(best["threshold"]),
        "best_f1": best["f1"],
        "results": results,
    }


def train_with_adapter(model_name: str, args) -> int:
    dataset_cfg = load_dataset_config(args.dataset)
    apply_dataloader_cli_overrides(dataset_cfg, args)
    model_cfg = load_model_config(model_name)
    if args.epochs is not None:
        model_cfg["num_epochs"] = int(args.epochs)
    if args.lr is not None:
        model_cfg["lr"] = float(args.lr)
    batch_size = int(args.batch_size or dataset_cfg.get("batch_size", 8))
    adapter = get_model_adapter(model_name)
    if not adapter.supports_unified_training:
        raise RuntimeError(f"{model_name} does not support unified training: {adapter.notes_or_failure_reason}")

    if args.dry_run:
        print(f"[DRY-RUN] unified model={model_name} dataset={dataset_cfg['name']} root={dataset_cfg['data_root']}")
        print(f"[DATASET] {dataset_runtime_summary(dataset_cfg)}")
        print_dataloader_policy(dataset_cfg, torch.cuda.is_available())
        return 0

    gpu_resolution = resolve_gpu(args.gpu)
    print_gpu_diagnostics(gpu_resolution)
    device = torch.device(gpu_resolution.local_device)

    out_dir = Path(args.output_dir) if args.output_dir else ROOT / "results" / model_name / dataset_cfg["name"]
    if not out_dir.is_absolute():
        out_dir = ROOT / out_dir
    ckpt_dir = out_dir / "checkpoints"
    log_dir = out_dir / "logs"
    metrics_dir = out_dir / "metrics"
    for path in (ckpt_dir, log_dir, metrics_dir, out_dir / "predictions" / "test", out_dir / "visuals" / "selected_20"):
        path.mkdir(parents=True, exist_ok=True)

    train_loader = _loader(dataset_cfg, "train", batch_size, shuffle=True, drop_last=True)
    val_loader = _loader(dataset_cfg, "val", batch_size, shuffle=False)
    test_loader = _loader(dataset_cfg, "test", batch_size, shuffle=False)
    print_dataloader_policy(dataset_cfg, torch.cuda.is_available())
    model = adapter.build_model(model_cfg, dataset_cfg, device)
    optimizer = adapter.build_optimizer(model, model_cfg)
    scheduler = adapter.build_scheduler(optimizer, model_cfg)
    amp_enabled = device.type == "cuda" and bool(getattr(adapter, "supports_amp_training", True))
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
    epochs = int(model_cfg.get("num_epochs", 200))
    best_path = ckpt_dir / "best_model.pth"
    latest_path = ckpt_dir / "latest.pth"
    best_f1 = -1.0
    best_epoch = 0
    selected_threshold = adapter.get_threshold(dataset_cfg)
    trajectory: list[dict] = []
    train_history: list[dict] = []
    val_history: list[dict] = []
    start_epoch = 1

    if args.resume and latest_path.exists():
        checkpoint = torch.load(latest_path, map_location=device)
        adapter.load_checkpoint(model, latest_path, device)
        if isinstance(checkpoint, dict) and "optimizer_state_dict" in checkpoint:
            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        if scheduler is not None and isinstance(checkpoint, dict) and "scheduler_state_dict" in checkpoint:
            scheduler.load_state_dict(checkpoint["scheduler_state_dict"])
        start_epoch = int(checkpoint.get("epoch", 0)) + 1
        best_f1 = float(checkpoint.get("best_val_f1", best_f1))
        best_epoch = int(checkpoint.get("best_epoch", 0))
        selected_threshold = float(checkpoint.get("selected_threshold", selected_threshold))

    if args.eval_only:
        adapter.load_checkpoint(model, best_path, device)
        test_metrics = _evaluate_split(
            model_name=model_name,
            model=model,
            adapter=adapter,
            dataset_cfg=dataset_cfg,
            loader=test_loader,
            device=device,
            threshold=selected_threshold,
            split="test",
            out_dir=out_dir,
            save_outputs=True,
            checkpoint_path=best_path,
        )
        save_metrics(model_name, dataset_cfg["name"], "test", test_metrics)
        _json_dump(metrics_dir / "test_metrics.json", test_metrics)
        append_to_comparison_table()
        return 0

    train_log = log_dir / "train.log"
    with train_log.open("w", encoding="utf-8") as log:
        log.write(f"# Unified Mamba-CD-style training: {model_name}/{dataset_cfg['name']}\n")
        log.write(f"# Dataset: {dataset_runtime_summary(dataset_cfg)}\n")
        for line in dataloader_policy_lines(dataset_cfg, torch.cuda.is_available()):
            log.write(line + "\n")
        if dataset_cfg.get("io_warning"):
            print(f"[DATASET-WARNING] {dataset_cfg['io_warning']}")
            log.write(f"# Warning: {dataset_cfg['io_warning']}\n")
        for epoch in range(start_epoch, epochs + 1):
            model.train()
            total_loss = 0.0
            component_sums: dict[str, float] = {}
            batch_count = 0
            progress = tqdm(train_loader, desc=f"{model_name} epoch {epoch}/{epochs}", dynamic_ncols=True)
            for batch_idx, batch in enumerate(progress, start=1):
                optimizer.zero_grad(set_to_none=True)
                with torch.amp.autocast("cuda", enabled=amp_enabled):
                    raw = adapter.forward(model, batch, device)
                if _has_nonfinite_tensor(raw):
                    raise RuntimeError(
                        f"{model_name}/{dataset_cfg['name']} produced non-finite model output "
                        f"at epoch {epoch}, batch {batch_idx}."
                    )
                with torch.amp.autocast("cuda", enabled=False):
                    loss_dict = adapter.compute_loss(raw, batch, model_cfg, dataset_cfg, device)
                    loss = loss_dict["loss"]
                if not bool(torch.isfinite(loss.detach()).all().item()):
                    components = {
                        key: float(value.detach().float().cpu().item())
                        for key, value in loss_dict.items()
                        if torch.is_tensor(value) and value.numel() == 1
                    }
                    raise RuntimeError(
                        f"{model_name}/{dataset_cfg['name']} produced non-finite loss "
                        f"at epoch {epoch}, batch {batch_idx}: {components}"
                    )
                scaler.scale(loss).backward()
                clip = model_cfg.get("grad_clip", model_cfg.get("gradient_clip", None))
                if clip is not None:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), float(clip))
                scaler.step(optimizer)
                scaler.update()
                total_loss += float(loss.detach().cpu().item())
                batch_count += 1
                for key, value in loss_dict.items():
                    component_sums[key] = component_sums.get(key, 0.0) + float(value.detach().cpu().item())
                lr = optimizer.param_groups[0]["lr"]
                gpu_mem = torch.cuda.max_memory_allocated(device) / (1024.0 ** 3) if device.type == "cuda" else 0.0
                progress.set_postfix(loss=f"{loss.item():.4f}", lr=f"{lr:.2e}", best_f1=f"{best_f1:.4f}", gpu_gb=f"{gpu_mem:.2f}")
            if scheduler is not None:
                scheduler.step()

            train_loss = total_loss / max(batch_count, 1)
            train_row = {
                "epoch": epoch,
                "train_loss": train_loss,
                "lr": optimizer.param_groups[0]["lr"],
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }
            for key, value in component_sums.items():
                train_row[key] = value / max(batch_count, 1)
            train_history.append(train_row)

            sweep = _threshold_sweep(
                model_name=model_name,
                model=model,
                adapter=adapter,
                dataset_cfg=dataset_cfg,
                loader=val_loader,
                device=device,
                split="val",
                out_dir=out_dir,
            )
            selected_threshold = float(sweep["selected_threshold"])
            _json_dump(metrics_dir / f"val_threshold_sweep_epoch_{epoch}.json", sweep)
            val_metrics = dict(max(sweep["results"], key=lambda row: row.get("f1", -1.0)))
            val_metrics["checkpoint"] = str(latest_path)
            is_best = val_metrics["f1"] > best_f1
            if is_best:
                best_f1 = float(val_metrics["f1"])
                best_epoch = epoch
            metadata = {
                "model": model_name,
                "dataset": dataset_cfg["name"],
                "epoch": epoch,
                "best_epoch": best_epoch,
                "best_val_f1": best_f1,
                "selected_threshold": selected_threshold,
                "threshold_mode": adapter.get_threshold_mode(),
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }
            adapter.save_checkpoint(model, optimizer, scheduler, latest_path, metadata)
            if is_best:
                adapter.save_checkpoint(model, optimizer, scheduler, best_path, metadata)
            val_payload = dict(val_metrics)
            val_payload.update({"epoch": epoch, "best_epoch": best_epoch, "selected_threshold": selected_threshold})
            val_history.append(val_payload)
            save_metrics(model_name, dataset_cfg["name"], "val", val_payload)
            _json_dump(metrics_dir / "val_history.json", val_history)
            _json_dump(metrics_dir / "train_history.json", train_history)
            trajectory_row = {
                "epoch": epoch,
                "train_loss": train_loss,
                "val_threshold": selected_threshold,
                "val_f1": val_payload["f1"],
                "val_iou": val_payload["iou"],
                "val_miou": val_payload["miou"],
                "val_precision": val_payload["precision"],
                "val_recall": val_payload["recall"],
                "val_oa": val_payload["oa"],
                "val_bf1": val_payload.get("bf1"),
                "val_gt_positive_ratio": val_payload.get("gt_positive_ratio"),
                "val_pred_positive_ratio": val_payload.get("pred_positive_ratio"),
                "val_mean_prob": val_payload.get("mean_prob"),
                "diagnostic_test_f1": None,
                "diagnostic_test_iou": None,
                "val_test_f1_gap": None,
                "val_test_iou_gap": None,
                "best_val_f1_so_far": best_f1,
                "best_val_epoch_or_iter_so_far": best_epoch,
                "checkpoint_path": str(best_path if is_best else latest_path),
            }
            trajectory.append(trajectory_row)
            _write_trajectory(metrics_dir, trajectory)
            warning = ""
            if val_payload["recall"] == 0:
                warning += " zero_recall"
            if val_payload.get("pred_positive_ratio", 0) == 0:
                warning += " all_background_prediction"
            line = (
                f"[Epoch {epoch:03d}/{epochs:03d}] train_loss={train_loss:.4f} "
                f"val_F1={val_payload['f1']:.4f} val_IoU={val_payload['iou']:.4f} "
                f"val_Prec={val_payload['precision']:.4f} val_Rec={val_payload['recall']:.4f} "
                f"threshold={selected_threshold} best_F1={best_f1:.4f}{' BEST' if is_best else ''}{warning}"
            )
            print(line, flush=True)
            log.write(line + "\n")
            log.flush()

    adapter.load_checkpoint(model, best_path, device)
    checkpoint = torch.load(best_path, map_location=device)
    selected_threshold = float(checkpoint.get("selected_threshold", selected_threshold)) if isinstance(checkpoint, dict) else selected_threshold
    test_metrics = _evaluate_split(
        model_name=model_name,
        model=model,
        adapter=adapter,
        dataset_cfg=dataset_cfg,
        loader=test_loader,
        device=device,
        threshold=selected_threshold,
        split="test",
        out_dir=out_dir,
        save_outputs=True,
        checkpoint_path=best_path,
    )
    test_metrics.update(count_parameters(model))
    try:
        if not adapter.supports_flops:
            raise ProfilingUnavailable(f"{model_name} adapter does not support FLOPs.")
        test_metrics.update(count_flops(model, lambda: adapter.get_dummy_inputs(dataset_cfg, device), device))
    except ProfilingUnavailable as exc:
        test_metrics.update({"flops": None, "flops_g": None, "flops_error": str(exc)})
    test_metrics.update({
        "best_epoch": best_epoch,
        "selected_threshold": selected_threshold,
        "checkpoint_path": str(best_path),
        "official_selection_metric": "validation_f1",
    })
    save_metrics(model_name, dataset_cfg["name"], "test", test_metrics)
    _json_dump(metrics_dir / "test_metrics.json", test_metrics)

    if _threshold_sweep_enabled(dataset_cfg, "test"):
        test_sweep = _threshold_sweep(
            model_name=model_name,
            model=model,
            adapter=adapter,
            dataset_cfg=dataset_cfg,
            loader=test_loader,
            device=device,
            split="test",
            out_dir=out_dir,
        )
        test_sweep.update({
            "diagnostic_only": True,
            "warning": "Test threshold sweep is diagnostic only and was not used for model selection.",
            "official_validation_selected_threshold": selected_threshold,
            "official_test_metrics": test_metrics,
        })
    else:
        test_sweep = {
            "diagnostic_only": True,
            "skipped": True,
            "reason": "Set eval.sweep_test_threshold: true in the dataset config to enable diagnostic test threshold sweeps.",
            "official_validation_selected_threshold": selected_threshold,
            "official_test_metrics": test_metrics,
        }
    _json_dump(metrics_dir / "test_threshold_sweep.json", test_sweep)

    best_val = max(val_history, key=lambda row: row.get("f1", -1.0)) if val_history else {}
    overfit = {
        "best_validation_epoch": best_epoch,
        "best_validation_f1": best_f1,
        "official_test_f1_at_validation_best_checkpoint": test_metrics.get("f1"),
        "validation_test_f1_gap": (best_val.get("f1", 0.0) - test_metrics.get("f1", 0.0)) if best_val else None,
        "validation_test_iou_gap": (best_val.get("iou", 0.0) - test_metrics.get("iou", 0.0)) if best_val else None,
        "gap_warning": bool(best_val and (best_val.get("f1", 0.0) - test_metrics.get("f1", 0.0)) > 0.1),
        "diagnostic_test_during_training_enabled": False,
        "note": "Test metrics were not used for checkpoint selection.",
    }
    _json_dump(metrics_dir / "overfit_diagnostics.json", overfit)
    if trajectory:
        trajectory[-1]["val_test_f1_gap"] = overfit["validation_test_f1_gap"]
        trajectory[-1]["val_test_iou_gap"] = overfit["validation_test_iou_gap"]
        _write_trajectory(metrics_dir, trajectory)
    append_to_comparison_table()
    _append_jsonl(ROOT / "results" / "training_log.jsonl", {
        "model": model_name,
        "dataset": dataset_cfg["name"],
        "status": "complete",
        "best_epoch": best_epoch,
        "best_val_f1": best_f1,
        "test_f1": test_metrics.get("f1"),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "trainer": "unified",
    })
    return 0