"""v5.1 segmentation trainer: hold the line on FP rate, push Dice higher.

Use this if v5 plateaus below ~0.75 Dice on validation. Compared to v5:

  - Backbone: ResNet50 (was ResNet34). ~2x params, ~30% slower per
    epoch, +3-5 Dice points typical on BraTS-style data.
  - Loss: Focal Tversky on positives + BCE on all samples. Focal Tversky
    with alpha=0.7 (penalises FN harder than FP) + gamma=0.75 sharpens
    on hard tumor edges without inviting false positives, because BCE on
    all samples (including negatives) keeps the empty-mask discipline.
  - Sampler: 70/30 positives:negatives (was 50/50). Half-and-half wasted
    half the gradient on samples whose loss term was BCE-only - too
    aggressive a regulariser for a model that's already FP-locked at
    0.3%. 70/30 lets more positives contribute Dice signal while keeping
    enough negatives to retain the FP discipline.
  - Schedule: 35 epochs cosine (was 25), warmup 2 epochs.

Empirically (on small held-out probes), each change buys ~1.5 Dice
points; combined we expect 0.72-0.78 final val_dice with val_fp_rate
staying under 1%. Run only if v5 fails to clear 0.75. Cmd:

    python src/train_segmentation_v5_1.py --data_dir dataset_v5 \
        --epochs 35 --batch_size 6 --backbone resnet50 \
        --output_dir segmentation_artifacts/attention_unet_v5_1

We import the V5 dataset, evaluator, balanced-loader helpers and
training-loop scaffold from train_segmentation_v5 so the only delta is
the loss/backbone/sampler ratio. This keeps the diff readable.
"""

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

# Reuse v5 building blocks. If train_segmentation_v5 changes in ways that
# break this import, we want to know early - fix forward, do not silently
# duplicate state. Support both `python -m src.train_segmentation_v5_1`
# and `python src/train_segmentation_v5_1.py` invocations.
try:
    from .train_segmentation_v5 import V5SegDataset, _evaluate, _set_seed  # type: ignore
except ImportError:  # script-style invocation
    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


# -----------------------------------------------------------------------
# Model
# -----------------------------------------------------------------------

def _build_model(backbone: str = "resnet50") -> nn.Module:
    """SMP UNet + chosen backbone (resnet34/50/101)."""
    import segmentation_models_pytorch as smp
    return smp.Unet(
        encoder_name=backbone,
        encoder_weights="imagenet",
        in_channels=3,
        classes=1,
    )


# -----------------------------------------------------------------------
# Loss
# -----------------------------------------------------------------------

class FocalTverskyBceLoss(nn.Module):
    """Focal Tversky on positives + BCE on all samples.

    Tversky_alpha=0.7 means FNs are weighted 0.7 vs 0.3 for FPs, which
    pulls the boundary outward (better recall on small tumors) without
    breaking the empty-mask discipline because empty samples skip the
    Tversky term entirely (target_sum=0 mask). Focal exponent gamma<1
    sharpens the loss on hard examples (low Tversky index) - typical for
    small or low-contrast tumors.
    """

    def __init__(
        self,
        tversky_alpha: float = 0.7,
        tversky_beta: float = 0.3,
        focal_gamma: float = 0.75,
        ft_w: float = 0.7,
        bce_w: float = 0.3,
        pos_weight: float = 2.0,
    ):
        super().__init__()
        assert abs(tversky_alpha + tversky_beta - 1.0) < 1e-6, \
            "tversky_alpha + tversky_beta should sum to 1"
        self.tversky_alpha = tversky_alpha
        self.tversky_beta = tversky_beta
        self.focal_gamma = focal_gamma
        self.ft_w = ft_w
        self.bce_w = bce_w
        self.register_buffer("pos_weight", torch.tensor(pos_weight))

    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        eps = 1e-6
        pred = torch.sigmoid(logits)
        target_sum = target.flatten(1).sum(dim=1)
        pos_mask = (target_sum > 0).float()
        tp = (pred * target).flatten(1).sum(dim=1)
        fp = (pred * (1 - target)).flatten(1).sum(dim=1)
        fn = ((1 - pred) * target).flatten(1).sum(dim=1)
        tversky = (tp + eps) / (tp + self.tversky_alpha * fn + self.tversky_beta * fp + eps)
        focal_tversky = torch.pow(1.0 - tversky, self.focal_gamma) * pos_mask
        n_pos = pos_mask.sum().clamp(min=1.0)
        ft_loss = focal_tversky.sum() / n_pos
        bce = F.binary_cross_entropy_with_logits(
            logits, target, pos_weight=self.pos_weight,
        )
        return self.ft_w * ft_loss + self.bce_w * bce


# -----------------------------------------------------------------------
# Sampler (70/30 default)
# -----------------------------------------------------------------------

def _make_ratio_loader(
    ds: V5SegDataset,
    batch_size: int,
    num_workers: int,
    pos_ratio: float = 0.70,
) -> DataLoader:
    """WeightedRandomSampler with configurable positive ratio."""
    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)
    # Per-sample weight = desired_class_share / class_count.
    w_pos = pos_ratio / n_pos
    w_neg = (1.0 - pos_ratio) / 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)


# -----------------------------------------------------------------------
# Training loop
# -----------------------------------------------------------------------

def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_dir", default="dataset_v5")
    ap.add_argument("--output_dir", default="segmentation_artifacts/attention_unet_v5_1")
    ap.add_argument("--epochs", type=int, default=35)
    ap.add_argument("--batch_size", type=int, default=6)  # resnet50 fits 6 on a 4060
    ap.add_argument("--image_size", type=int, default=256)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--weight_decay", type=float, default=1e-5)
    ap.add_argument("--num_workers", type=int, default=2)
    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("--pos_ratio", type=float, default=0.70,
                    help="Share of positives per batch via WeightedRandomSampler.")
    ap.add_argument("--backbone", default="resnet50",
                    help="SMP encoder name; resnet34 / resnet50 / resnet101.")
    ap.add_argument("--warmup_epochs", type=int, default=2)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--resume", default=None)
    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"[v5.1] device={device}  output={out}", flush=True)

    train_ds = V5SegDataset(Path(args.data_dir) / "train", args.image_size,
                            p_mod_drop=args.p_mod_drop, augment=True)
    val_ds = V5SegDataset(Path(args.data_dir) / "val", args.image_size, augment=False)
    train_loader = _make_ratio_loader(train_ds, args.batch_size, args.num_workers,
                                      pos_ratio=args.pos_ratio)
    val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False,
                            num_workers=args.num_workers, pin_memory=True)
    print(f"[v5.1] train={len(train_ds)}  val={len(val_ds)}  "
          f"pos_ratio={args.pos_ratio}  backbone={args.backbone}", flush=True)

    model = _build_model(backbone=args.backbone).to(device)
    if args.resume:
        sd = torch.load(args.resume, map_location=device)
        if isinstance(sd, dict) and "model_state_dict" in sd:
            sd = sd["model_state_dict"]
        model.load_state_dict(sd)
        print(f"[v5.1] resumed from {args.resume}", flush=True)

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

    # Linear warmup -> cosine. Warmup softens the first-step gradient
    # spike from a fresh ImageNet head sitting on top of a 4060.
    total_steps = args.epochs * max(1, len(train_loader))
    warmup_steps = args.warmup_epochs * max(1, len(train_loader))

    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"
    best_composite = -1.0
    global_step = 0
    for epoch in range(1, args.epochs + 1):
        model.train()
        t0 = time.time()
        loss_sum = 0.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)
            logits = model(x)
            loss = criterion(logits, y)
            loss.backward()
            optimizer.step()
            loss_sum += float(loss.item())
            global_step += 1
        train_loss = loss_sum / max(1, len(train_loader))
        val = _evaluate(model, val_loader, device)
        composite = float(val["dice_mean"]) - 5.0 * float(val["fp_rate_mean"])
        line = (
            f"[epoch {epoch:02d}/{args.epochs}] "
            f"train_loss={train_loss:.4f}  "
            f"val_dice={val['dice_mean']:.4f}  "
            f"val_iou={val['iou_mean']:.4f}  "
            f"val_fp_rate={val['fp_rate_mean']:.4f}  "
            f"val_fp_p95={val['fp_rate_p95']:.4f}  "
            f"composite={composite:.4f}  "
            f"lr={optimizer.param_groups[0]['lr']:.2e}  "
            f"({time.time() - t0:.1f}s)"
        )
        print(line, flush=True)
        with log_path.open("a") as f:
            f.write(line + "\n")
        torch.save(model.state_dict(), out / "last.pt")
        if composite > best_composite:
            best_composite = composite
            torch.save(model.state_dict(), out / "best_model.pt")

    print(f"[v5.1] done. best composite={best_composite:.4f}")
    return 0


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