src/train_segmentation_v10.py

"""v9 brain-2D trainer: integrates all v9 research heads with crash-safe training.

Architecture summary (see src/research/_v10_universal_hyperbolic/v10_model.py for details)
---------------------------------------------------------------
  Geometric prior (SDF) -> SMP UNet encoder (ConvNeXt-Tiny, 4-channel)
    -> bottleneck -> latent_dim=256
    -> hyperbolic projection (learnable curvature)
    -> causal SCM head: (anatomy, tumor, scanner) + NOTEARS DAG
    -> recompose (anatomy + tumor) -> SMP UNet decoder
    -> mask logits

  AND on a parallel branch:
    -> SCM with z_tumor=0 -> counterfactual healthy decoder
    -> counterfactual image + residual

Multi-task loss:
  L = L_seg (Tversky + Dice + BCE)         <- segmentation
    + lambda_o * (L_o_at + L_o_as + L_o_ts) <- SCM orthogonality
    + lambda_dag * L_dag                    <- NOTEARS acyclicity
    + lambda_forbid * L_dag_forbidden       <- biological priors
    + lambda_cf * L_cf_recon                <- counterfactual reconstruction
    + lambda_hyp * L_hyp_reg                <- hyperbolic embedding reg

All loss weights default to safe values (small SCM regularizers) so the
segmentation task dominates initially. Trained with the same crash-safe
checkpointing + AMP + RAM cache infrastructure as v7.

Designed to run on the same dataset_v8 as v7/v8. No multi-organ scope.
Multi-organ extension is v10.

Usage (Colab or local):
  python src/train_segmentation_v10.py --data_dir dataset_v8 \
    --output_dir segmentation_artifacts/attention_unet_v9 \
    --epochs 100 --batch_size 32 --image_size 384 \
    --amp --cache_in_ram --resume auto
"""

from __future__ import annotations

import argparse
import math
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, WeightedRandomSampler

try:
    from .train_segmentation_v5 import V5SegDataset, _evaluate, _set_seed  # type: ignore
    from .train_segmentation_v7 import V7SegDataset, TverskyDiceBceLoss, _evaluate_with_micro, \
        _save_checkpoint, _atomic_save  # type: ignore
    from .research._v10_universal_hyperbolic.v10_model import V10Model  # type: ignore
except ImportError:
    import sys as _sys
    from pathlib import Path as _Path
    _sys.path.insert(0, str(_Path(__file__).resolve().parents[1]))
    from src.train_segmentation_v5 import V5SegDataset, _evaluate, _set_seed  # type: ignore
    from src.train_segmentation_v7 import V7SegDataset, TverskyDiceBceLoss, _evaluate_with_micro, \
        _save_checkpoint, _atomic_save  # type: ignore
    from src.research._v10_universal_hyperbolic.v10_model import V10Model  # type: ignore


# -----------------------------------------------------------------------
# Balanced sampler (matches v7)
# -----------------------------------------------------------------------

def _make_balanced_loader(ds, batch_size: int, num_workers: int,
                          prefetch_factor: int = 4) -> DataLoader:
    extra = {}
    if num_workers > 0:
        extra['persistent_workers'] = True
        extra['prefetch_factor'] = prefetch_factor
    flags = ds.has_tumor_flags()
    n_pos = sum(1 for f in flags if f)
    n_neg = sum(1 for f in flags if not f)
    if n_pos == 0 or n_neg == 0:
        return DataLoader(ds, batch_size=batch_size, shuffle=True,
                          num_workers=num_workers, pin_memory=True,
                          drop_last=True, **extra)
    w_pos = 0.5 / n_pos
    w_neg = 0.5 / n_neg
    weights = [w_pos if f else w_neg for f in flags]
    sampler = WeightedRandomSampler(weights, num_samples=len(ds), replacement=True)
    return DataLoader(ds, batch_size=batch_size, sampler=sampler,
                      num_workers=num_workers, pin_memory=True,
                      drop_last=True, **extra)


# -----------------------------------------------------------------------
# v9 evaluation (extends v7 evaluator with SCM monitoring)
# -----------------------------------------------------------------------

@torch.no_grad()
def _evaluate_v9(model, loader, device, threshold: float = 0.5, amp: bool = False) -> dict:
    """Evaluate v9 model: standard seg metrics + SCM disentanglement health."""
    model.eval()
    use_amp = amp and device.type == "cuda" and torch.cuda.is_bf16_supported()
    dices, ious, fp_rates = [], [], []
    tp_total = fp_total = fn_total = 0
    ortho_at_sum = ortho_as_sum = ortho_ts_sum = 0.0
    dag_sum = 0.0
    cf_recon_loss_sum = 0.0
    n_batches = 0
    n_pos = n_neg = 0

    for x, y, _ in loader:
        x, y = x.to(device), y.to(device)
        if use_amp:
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                out = model(x, return_counterfactual=True)
        else:
            out = model(x, return_counterfactual=True)
        logits = out["mask_logits"]
        p = torch.sigmoid(logits)
        m = (p >= threshold).float()
        # Per-sample seg metrics
        for i in range(x.size(0)):
            yi, mi = y[i], m[i]
            if yi.sum() > 0:
                inter = (mi * yi).sum().item()
                pred_sum = mi.sum().item()
                tgt_sum = yi.sum().item()
                dices.append((2 * inter + 1e-6) / (pred_sum + tgt_sum + 1e-6))
                ious.append((inter + 1e-6) / (pred_sum + tgt_sum - inter + 1e-6))
                n_pos += 1
            else:
                fp_rates.append(mi.mean().item())
                n_neg += 1
            tp_total += float((mi * yi).sum().item())
            fp_total += float((mi * (1 - yi)).sum().item())
            fn_total += float(((1 - mi) * yi).sum().item())

        # SCM metrics (averaged across batch)
        aux = out["aux_losses"]
        ortho_at_sum += float(aux["ortho_at"].item())
        ortho_as_sum += float(aux["ortho_as"].item())
        ortho_ts_sum += float(aux["ortho_ts"].item())
        dag_sum += float(aux["dag"].item())

        # Counterfactual reconstruction loss (only meaningful for empty-mask negatives)
        if out["x_counterfactual"] is not None:
            # For healthy scans (y empty), x_cf should equal x.
            healthy_mask = (y.sum(dim=(1, 2, 3)) == 0).float()
            if healthy_mask.any():
                per_sample_loss = (out["x_counterfactual"] - x).abs().mean(dim=(1, 2, 3))
                cf_recon_loss_sum += float((per_sample_loss * healthy_mask).sum() / healthy_mask.sum())

        n_batches += 1

    import numpy as _np
    micro_dice = (2 * tp_total + 1e-6) / (2 * tp_total + fp_total + fn_total + 1e-6)
    return {
        "n_positive": n_pos,
        "n_negative": n_neg,
        "dice_mean": float(_np.mean(dices)) if dices else 0.0,
        "iou_mean": float(_np.mean(ious)) if ious else 0.0,
        "micro_dice": float(micro_dice),
        "fp_rate_mean": float(_np.mean(fp_rates)) if fp_rates else 0.0,
        "fp_rate_p95": float(_np.percentile(fp_rates, 95)) if fp_rates else 0.0,
        "ortho_at": ortho_at_sum / max(1, n_batches),
        "ortho_as": ortho_as_sum / max(1, n_batches),
        "ortho_ts": ortho_ts_sum / max(1, n_batches),
        "dag_h": dag_sum / max(1, n_batches),
        "cf_recon_loss": cf_recon_loss_sum / max(1, n_batches),
    }


# -----------------------------------------------------------------------
# v9 save (full state, atomic)
# -----------------------------------------------------------------------

def _save_v9_checkpoint(out: Path, name: str, *, model, optimizer, epoch: int,
                        global_step: int, best_micro: float, best_composite: float,
                        args) -> None:
    payload = {
        "model_state_dict": model.state_dict(),
        "optimizer_state_dict": optimizer.state_dict(),
        "epoch": int(epoch),
        "global_step": int(global_step),
        "best_micro": float(best_micro),
        "best_composite": float(best_composite),
        "args": vars(args),
        "schema_version": "v9_brain2d_1",
    }
    _atomic_save(payload, out / name)


# -----------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_dir", default="dataset_v8")
    ap.add_argument("--output_dir", default="segmentation_artifacts/attention_unet_v9")
    ap.add_argument("--epochs", type=int, default=100)
    ap.add_argument("--batch_size", type=int, default=16)
    ap.add_argument("--image_size", type=int, default=384)
    ap.add_argument("--lr", type=float, default=8e-5)
    ap.add_argument("--weight_decay", type=float, default=1e-5)
    ap.add_argument("--num_workers", type=int, default=4)
    ap.add_argument("--prefetch_factor", type=int, default=4)
    ap.add_argument("--p_mod_drop", type=float, default=0.3)
    ap.add_argument("--bce_pos_weight", type=float, default=2.0)
    ap.add_argument("--warmup_epochs", type=int, default=3)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--resume", default=None,
                    help="Path to checkpoint or 'auto' for last.pt in output_dir.")
    ap.add_argument("--checkpoint_every_steps", type=int, default=500)
    ap.add_argument("--amp", action="store_true",
                    help="bf16 mixed precision on A100/H100.")
    ap.add_argument("--cache_in_ram", action="store_true",
                    help="Preload dataset bytes in RAM (~900 MB for dataset_v8).")

    # v9-specific model hyperparams
    ap.add_argument("--latent_dim", type=int, default=256)
    ap.add_argument("--anatomy_dim", type=int, default=128)
    ap.add_argument("--tumor_dim", type=int, default=64)
    ap.add_argument("--scanner_dim", type=int, default=32)
    ap.add_argument("--no_counterfactual", action="store_true",
                    help="Disable counterfactual healthy decoder (lighter).")
    ap.add_argument("--no_geometric_prior", action="store_true",
                    help="Disable SDF geometric prior (3-channel input instead of 4).")
    ap.add_argument("--hyperbolic_curvature_init", type=float, default=1.0)

    # v9-specific loss weights (multi-task balance)
    ap.add_argument("--lambda_ortho", type=float, default=0.05,
                    help="Weight on SCM orthogonality losses (anatomy-tumor, anatomy-scanner, tumor-scanner).")
    ap.add_argument("--lambda_dag", type=float, default=0.01,
                    help="Weight on NOTEARS DAG-ness loss.")
    ap.add_argument("--lambda_forbidden", type=float, default=0.05,
                    help="Weight on forbidden-edge penalty (scanner->anatomy, tumor->anatomy).")
    ap.add_argument("--lambda_cf", type=float, default=0.10,
                    help="Weight on counterfactual reconstruction loss.")
    args = ap.parse_args()

    _set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    out = Path(args.output_dir)
    out.mkdir(parents=True, exist_ok=True)
    print(f"[v9] device={device}  output={out}", flush=True)

    train_ds = V7SegDataset(Path(args.data_dir) / "train", args.image_size,
                            p_mod_drop=args.p_mod_drop, augment=True,
                            cache_in_ram=args.cache_in_ram)
    val_ds = V7SegDataset(Path(args.data_dir) / "val", args.image_size, augment=False,
                          cache_in_ram=args.cache_in_ram)
    train_loader = _make_balanced_loader(train_ds, args.batch_size, args.num_workers,
                                          prefetch_factor=args.prefetch_factor)
    val_extra = {}
    if args.num_workers > 0:
        val_extra["persistent_workers"] = True
        val_extra["prefetch_factor"] = args.prefetch_factor
    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
                            num_workers=args.num_workers, pin_memory=True, **val_extra)
    amp_status = "bf16-AMP" if args.amp else "fp32"
    cache_status = "RAM-cached" if args.cache_in_ram else "disk-streamed"
    print(f"[v9] train={len(train_ds)}  val={len(val_ds)}  image_size={args.image_size}  "
          f"precision={amp_status}  data={cache_status}", flush=True)

    # Build model with all v9 heads
    model = V10Model(
        image_size=args.image_size,
        latent_dim=args.latent_dim,
        anatomy_dim=args.anatomy_dim,
        tumor_dim=args.tumor_dim,
        scanner_dim=args.scanner_dim,
        use_counterfactual=not args.no_counterfactual,
        use_geometric_prior=not args.no_geometric_prior,
        hyperbolic_curvature_init=args.hyperbolic_curvature_init,
    ).to(device)
    n_params = model.num_parameters()
    print(f"[v9] params={n_params/1e6:.1f}M  "
          f"(counterfactual={'on' if not args.no_counterfactual else 'off'}, "
          f"geometric_prior={'on' if not args.no_geometric_prior else 'off'})", flush=True)

    criterion = TverskyDiceBceLoss(pos_weight=args.bce_pos_weight).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr,
                                  weight_decay=args.weight_decay)

    total_steps = args.epochs * max(1, len(train_loader))
    warmup_steps = args.warmup_epochs * max(1, len(train_loader))

    # ---- Resume support ----
    start_epoch = 1
    best_micro = -1.0
    best_composite = -1.0
    global_step = 0
    resume_path = None
    if args.resume:
        if args.resume.lower() == "auto":
            cand = out / "last.pt"
            if cand.exists():
                resume_path = cand
            else:
                print(f"[v9] --resume auto: no last.pt in {out}, starting fresh", flush=True)
        else:
            resume_path = Path(args.resume)
    if resume_path is not None:
        try:
            ckpt = torch.load(resume_path, map_location=device, weights_only=False)
        except Exception as exc:
            print(f"[v9] CHECKPOINT CORRUPT at {resume_path}: {exc}", flush=True)
            ckpt = None
        if ckpt is not None and isinstance(ckpt, dict) and "model_state_dict" in ckpt:
            try:
                model.load_state_dict(ckpt["model_state_dict"])
                if "optimizer_state_dict" in ckpt:
                    try:
                        optimizer.load_state_dict(ckpt["optimizer_state_dict"])
                    except Exception as exc:
                        print(f"[v9] optimizer state restore failed: {exc} (continuing fresh)",
                              flush=True)
                start_epoch = int(ckpt.get("epoch", 0)) + 1
                global_step = int(ckpt.get("global_step", 0))
                best_micro = float(ckpt.get("best_micro", -1.0))
                best_composite = float(ckpt.get("best_composite", -1.0))
                print(f"[v9] resumed from {resume_path}  epoch={start_epoch}  "
                      f"step={global_step}  best_micro={best_micro:.4f}", flush=True)
            except RuntimeError as exc:
                print(f"[v9] model state mismatch (probably arch change): {exc}", flush=True)
                print(f"[v9] starting fresh from epoch 1", flush=True)

    if start_epoch > args.epochs:
        print(f"[v9] checkpoint already past target epoch {args.epochs}. Nothing to do.",
              flush=True)
        return 0

    def lr_at(step):
        if step < warmup_steps:
            return float(step + 1) / max(1, warmup_steps)
        progress = (step - warmup_steps) / max(1, total_steps - warmup_steps)
        return 0.5 * (1.0 + math.cos(math.pi * progress))

    log_path = out / "training.log"
    amp_enabled = bool(args.amp) and device.type == "cuda"
    if amp_enabled and not torch.cuda.is_bf16_supported():
        print("[v9] WARNING: --amp requested but no bf16 support, falling back to fp32",
              flush=True)
        amp_enabled = False

    for epoch in range(start_epoch, args.epochs + 1):
        model.train()
        t0 = time.time()
        seg_loss_sum = 0.0
        scm_loss_sum = 0.0
        cf_loss_sum = 0.0
        steps_in_epoch = 0
        for x, y, _ in train_loader:
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)
            for g in optimizer.param_groups:
                g["lr"] = args.lr * lr_at(global_step)
            optimizer.zero_grad(set_to_none=True)

            if amp_enabled:
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    out_dict = model(x, return_counterfactual=not args.no_counterfactual)
                    logits = out_dict["mask_logits"]
                    seg_loss = criterion(logits, y)
                    aux = out_dict["aux_losses"]
                    scm_loss = (args.lambda_ortho * (aux["ortho_at"] + aux["ortho_as"] + aux["ortho_ts"])
                                + args.lambda_dag * aux["dag"]
                                + args.lambda_forbidden * aux["dag_forbidden"])
                    cf_loss = torch.tensor(0.0, device=device)
                    if out_dict["x_counterfactual"] is not None and args.lambda_cf > 0:
                        cf_loss = model.cf_decoder.reconstruction_loss(
                            x, out_dict["x_counterfactual"], y,
                            lambda_outside=1.0, lambda_inside=0.5,
                        )
                    loss = seg_loss + scm_loss + args.lambda_cf * cf_loss
            else:
                out_dict = model(x, return_counterfactual=not args.no_counterfactual)
                logits = out_dict["mask_logits"]
                seg_loss = criterion(logits, y)
                aux = out_dict["aux_losses"]
                scm_loss = (args.lambda_ortho * (aux["ortho_at"] + aux["ortho_as"] + aux["ortho_ts"])
                            + args.lambda_dag * aux["dag"]
                            + args.lambda_forbidden * aux["dag_forbidden"])
                cf_loss = torch.tensor(0.0, device=device)
                if out_dict["x_counterfactual"] is not None and args.lambda_cf > 0:
                    cf_loss = model.cf_decoder.reconstruction_loss(
                        x, out_dict["x_counterfactual"], y,
                        lambda_outside=1.0, lambda_inside=0.5,
                    )
                loss = seg_loss + scm_loss + args.lambda_cf * cf_loss

            loss.backward()
            optimizer.step()

            seg_loss_sum += float(seg_loss.item())
            scm_loss_sum += float(scm_loss.item())
            cf_loss_sum += float(cf_loss.item()) if torch.is_tensor(cf_loss) else 0.0
            global_step += 1
            steps_in_epoch += 1

            if args.checkpoint_every_steps > 0 and global_step % args.checkpoint_every_steps == 0:
                try:
                    _save_v9_checkpoint(out, "last.pt",
                                        model=model, optimizer=optimizer,
                                        epoch=epoch - 1,
                                        global_step=global_step,
                                        best_micro=best_micro,
                                        best_composite=best_composite,
                                        args=args)
                except Exception as exc:
                    print(f"[v9] intra-epoch save failed: {exc}", flush=True)

        n = max(1, steps_in_epoch)
        train_seg = seg_loss_sum / n
        train_scm = scm_loss_sum / n
        train_cf = cf_loss_sum / n

        val = _evaluate_v9(model, val_loader, device, amp=amp_enabled)
        composite = float(val["dice_mean"]) - 5.0 * float(val["fp_rate_mean"])
        c_curvature = float(model.hyperbolic.c.detach().item())
        line = (
            f"[epoch {epoch:02d}/{args.epochs}] "
            f"seg={train_seg:.4f}  scm={train_scm:.4f}  cf={train_cf:.4f}  | "
            f"dice={val['dice_mean']:.4f}  micro={val['micro_dice']:.4f}  "
            f"fp={val['fp_rate_mean']:.4f}  comp={composite:.4f}  | "
            f"o_at={val['ortho_at']:.4f}  o_as={val['ortho_as']:.4f}  "
            f"o_ts={val['ortho_ts']:.4f}  dag_h={val['dag_h']:.4f}  "
            f"c={c_curvature:.3f}  cf_recon={val['cf_recon_loss']:.4f}  "
            f"lr={optimizer.param_groups[0]['lr']:.2e}  ({time.time() - t0:.1f}s)"
        )
        print(line, flush=True)
        with log_path.open("a") as f:
            f.write(line + "\n")

        # Save checkpoints (last + best by micro + best by composite)
        _save_v9_checkpoint(out, "last.pt",
                            model=model, optimizer=optimizer,
                            epoch=epoch, global_step=global_step,
                            best_micro=best_micro, best_composite=best_composite,
                            args=args)
        if val["micro_dice"] > best_micro:
            best_micro = val["micro_dice"]
            _save_v9_checkpoint(out, "best_micro.pt",
                                model=model, optimizer=optimizer,
                                epoch=epoch, global_step=global_step,
                                best_micro=best_micro, best_composite=best_composite,
                                args=args)
        if composite > best_composite:
            best_composite = composite
            _save_v9_checkpoint(out, "best_model.pt",
                                model=model, optimizer=optimizer,
                                epoch=epoch, global_step=global_step,
                                best_micro=best_micro, best_composite=best_composite,
                                args=args)

    print(f"[v9] done. best micro={best_micro:.4f}  best composite={best_composite:.4f}",
          flush=True)
    return 0


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