from __future__ import annotations

import argparse
import json
import random
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.utils.data import DataLoader

from .datasets import compute_class_histogram
from .datasets_instance import KymoInstanceDataset
from .models import UNetInstance


def default_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def moving_stats_from_batch(logits: torch.Tensor, target: torch.Tensor, *, ignore_label: int = 255) -> tuple[float, float, float]:
    pred = torch.argmax(logits, dim=1)
    valid = target != int(ignore_label)
    pred_m = torch.logical_and(pred == 1, valid)
    gt_m = torch.logical_and(target == 1, valid)
    tp = torch.logical_and(pred_m, gt_m).sum().item()
    fp = torch.logical_and(pred_m, ~gt_m).sum().item()
    fn = torch.logical_and(~pred_m, gt_m).sum().item()
    return float(tp), float(fp), float(fn)


def moving_metrics(tp: float, fp: float, fn: float) -> dict[str, float]:
    precision = tp / (tp + fp + 1e-9)
    recall = tp / (tp + fn + 1e-9)
    f1 = 2 * precision * recall / (precision + recall + 1e-9)
    iou = tp / (tp + fp + fn + 1e-9)
    return {
        "moving_precision": float(precision),
        "moving_recall": float(recall),
        "moving_f1": float(f1),
        "moving_iou": float(iou),
    }


def _resolve_manifest_default_is_real(mode: str, manifest_path: str) -> bool:
    txt = str(mode).strip().lower()
    if txt in {"1", "true", "yes", "y", "real"}:
        return True
    if txt in {"0", "false", "no", "n", "synth", "synthetic"}:
        return False
    path_txt = str(Path(manifest_path).as_posix()).lower()
    return "/real/" in path_txt or path_txt.endswith("/real/index.jsonl") or path_txt.endswith("/real/index_train.jsonl") or path_txt.endswith("/real/index_val.jsonl")


def _inverse_freq_weights(hist: np.ndarray) -> np.ndarray:
    hist = np.asarray(hist, dtype=np.float64)
    present = hist > 0
    if not np.any(present):
        return np.ones_like(hist, dtype=np.float64)
    inv = np.zeros_like(hist, dtype=np.float64)
    inv[present] = 1.0 / np.maximum(hist[present], 1e-12)
    z = float(np.sum(inv[present]))
    if z <= 0:
        return np.ones_like(hist, dtype=np.float64)
    out = np.zeros_like(hist, dtype=np.float64)
    out[present] = (inv[present] / z) * float(np.count_nonzero(present))
    return out


def make_class_weights(
    dataset,
    num_classes: int = 3,
    *,
    log_progress: bool = True,
    synth_only: bool = False,
) -> torch.Tensor:
    if bool(synth_only):
        samples = getattr(dataset, "samples", None)
        if isinstance(samples, list):
            idxs = [i for i, s in enumerate(samples) if not bool(getattr(s, "is_real", False))]
        else:
            idxs = list(range(len(dataset)))
    else:
        idxs = list(range(len(dataset)))

    n = len(idxs)
    progress_every = max(1, n // 20) if (log_progress and n > 0) else 0
    if log_progress:
        scope = "synth-only " if synth_only else ""
        print(f"[info] computing {scope}class histogram over {n} training sample(s)...", flush=True)

    hist = np.zeros((num_classes,), dtype=np.float64)
    for j, ds_idx in enumerate(idxs, start=1):
        item = dataset[int(ds_idx)]
        y = item["label"].numpy().astype(np.int64)
        vals, counts = np.unique(y, return_counts=True)
        for v, c in zip(vals, counts):
            if 0 <= int(v) < num_classes:
                hist[int(v)] += float(c)
        if progress_every > 0 and (j % progress_every == 0 or j == n):
            print(
                f"[info] class histogram: {j}/{n} ({(100.0 * j / max(1, n)):.1f}%)",
                flush=True,
            )

    if n == 0:
        if log_progress:
            print("[warn] synth-only histogram had 0 samples; falling back to full dataset histogram.", flush=True)
        hist = compute_class_histogram(dataset, num_classes=num_classes)

    w = _inverse_freq_weights(hist)
    if log_progress:
        print(f"[info] class histogram: {hist.tolist()}", flush=True)
        print(f"[info] class weights: {w.tolist()}", flush=True)
    return torch.tensor(w, dtype=torch.float32)


def real_loss_moving_vs_rest(
    logits: torch.Tensor,
    targets: torch.Tensor,
    *,
    ignore_label: int = 255,
    bootstrap_ignore_prob: float | None = None,
    bootstrap_max_frac: float = 0.2,
) -> torch.Tensor:
    """Supervise real labels as moving-vs-rest with optional bootstrap ignore on unlabeled pixels."""
    valid_mask = targets != int(ignore_label)
    targets01 = targets.clamp(min=0, max=1)
    moving = logits[:, 1]
    rest = torch.logsumexp(torch.stack([logits[:, 0], logits[:, 2]], dim=1), dim=1)
    two_logits = torch.stack([rest, moving], dim=1)

    if bootstrap_ignore_prob is not None and float(bootstrap_ignore_prob) > 0.0:
        with torch.no_grad():
            probs = torch.softmax(two_logits, dim=1)[:, 1]
            candidates = torch.logical_and(valid_mask, targets01 == 0)
            candidates = torch.logical_and(candidates, probs >= float(bootstrap_ignore_prob))
            if torch.any(candidates):
                max_frac = float(max(0.0, min(1.0, bootstrap_max_frac)))
                cand_idx = torch.nonzero(candidates, as_tuple=False)
                n_cand = int(cand_idx.shape[0])
                n_valid = int(torch.count_nonzero(valid_mask).item())
                cap = max(0, int(max_frac * max(1, n_valid)))
                if cap <= 0:
                    valid_mask[candidates] = False
                elif n_cand <= cap:
                    valid_mask[candidates] = False
                else:
                    sel = torch.randperm(n_cand, device=logits.device)[:cap]
                    chosen = cand_idx[sel]
                    valid_mask[chosen[:, 0], chosen[:, 1], chosen[:, 2]] = False

    loss_map = F.cross_entropy(two_logits, targets01, reduction="none")
    valid = valid_mask.float()
    denom = valid.sum()
    if float(denom.item()) <= 0.0:
        return logits.sum() * 0.0
    return (loss_map * valid).sum() / denom


def set_backbone_trainable(model: UNetInstance, trainable: bool) -> None:
    modules = [model.inc, model.down1, model.down2, model.down3, model.up1, model.up2, model.up3]
    for mod in modules:
        for p in mod.parameters():
            p.requires_grad = bool(trainable)


def load_pretrained_into_instance(model: UNetInstance, ckpt_path: str | Path) -> tuple[int, int]:
    """Load compatible weights from a seg/instance checkpoint into UNetInstance.

    Returns:
        (num_loaded_tensors, total_model_tensors)
    """

    load_kwargs = {"map_location": "cpu"}
    try:
        ckpt = torch.load(str(ckpt_path), weights_only=True, **load_kwargs)
    except TypeError:
        ckpt = torch.load(str(ckpt_path), **load_kwargs)

    state_dict = ckpt.get("model_state") if isinstance(ckpt, dict) and "model_state" in ckpt else ckpt
    if not isinstance(state_dict, dict):
        raise ValueError(f"Checkpoint does not contain a valid state_dict: {ckpt_path}")

    mapped: dict[str, torch.Tensor] = {}
    for key, value in state_dict.items():
        new_key = str(key)
        if new_key.startswith("model."):
            new_key = new_key[len("model.") :]
        if new_key.startswith("outc."):
            new_key = "seg_head." + new_key[len("outc.") :]
        mapped[new_key] = value

    model_state = model.state_dict()
    compatible: dict[str, torch.Tensor] = {}
    for key, value in mapped.items():
        if key in model_state and tuple(model_state[key].shape) == tuple(value.shape):
            compatible[key] = value

    model.load_state_dict(compatible, strict=False)
    return len(compatible), len(model_state)


def _maybe_crop_hw_instance(
    img: torch.Tensor,
    lbl: torch.Tensor,
    inst: torch.Tensor,
    *,
    max_h: int,
    max_w: int,
    random_crop: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    h = int(img.shape[-2])
    w = int(img.shape[-1])
    if max_h > 0 and h > max_h:
        top = int(np.random.randint(0, h - max_h + 1)) if random_crop else int((h - max_h) // 2)
        img = img[..., top : top + max_h, :]
        lbl = lbl[..., top : top + max_h, :]
        inst = inst[..., top : top + max_h, :]
        h = max_h
    if max_w > 0 and w > max_w:
        left = int(np.random.randint(0, w - max_w + 1)) if random_crop else int((w - max_w) // 2)
        img = img[..., :, left : left + max_w]
        lbl = lbl[..., :, left : left + max_w]
        inst = inst[..., :, left : left + max_w]
    return img, lbl, inst


def make_pad_collate(
    label_pad_value: int,
    *,
    max_h: int = 0,
    max_w: int = 0,
    random_crop: bool = False,
):
    """Pad variable-size (T, X) kymographs within a batch to max shape."""

    def _collate(batch: list[dict]) -> dict:
        cap_h = int(max_h)
        cap_w = int(max_w)
        batch_h = 0
        batch_w = 0
        for item in batch:
            img = item["image"]
            h0 = int(img.shape[-2])
            w0 = int(img.shape[-1])
            if cap_h > 0:
                h0 = min(h0, cap_h)
            if cap_w > 0:
                w0 = min(w0, cap_w)
            batch_h = max(batch_h, h0)
            batch_w = max(batch_w, w0)

        images = []
        labels = []
        instances = []
        has_instances = []
        is_real = []
        image_paths = []
        label_paths = []
        meta_paths = []

        for item in batch:
            img = item["image"]
            lbl = item["label"]
            inst = item["instance"]
            img, lbl, inst = _maybe_crop_hw_instance(
                img,
                lbl,
                inst,
                max_h=cap_h,
                max_w=cap_w,
                random_crop=bool(random_crop),
            )
            h = int(img.shape[-2])
            w = int(img.shape[-1])
            pad_h = batch_h - h
            pad_w = batch_w - w
            if pad_h > 0 or pad_w > 0:
                img = F.pad(img, (0, pad_w, 0, pad_h), mode="constant", value=0.0)
                lbl = F.pad(lbl, (0, pad_w, 0, pad_h), mode="constant", value=int(label_pad_value))
                inst = F.pad(inst, (0, pad_w, 0, pad_h), mode="constant", value=0)

            images.append(img.float())
            labels.append(lbl.long())
            instances.append(inst.long())
            has_instances.append(item["has_instances"])
            is_real.append(item["is_real"])
            image_paths.append(item.get("image_path", ""))
            label_paths.append(item.get("label_path", ""))
            meta_paths.append(item.get("meta_path", ""))

        return {
            "image": torch.stack(images, dim=0),
            "label": torch.stack(labels, dim=0),
            "instance": torch.stack(instances, dim=0),
            "has_instances": torch.stack(has_instances, dim=0),
            "is_real": torch.stack(is_real, dim=0),
            "image_path": image_paths,
            "label_path": label_paths,
            "meta_path": meta_paths,
        }

    return _collate


def discriminative_embedding_loss(
    embeddings: torch.Tensor,
    instances: torch.Tensor,
    seg_targets: torch.Tensor,
    *,
    ignore_label: int = 255,
    moving_class: int = 1,
    delta_var: float = 0.5,
    delta_dist: float = 1.5,
    alpha: float = 1.0,
    beta: float = 1.0,
    gamma: float = 0.001,
) -> tuple[torch.Tensor, dict[str, float]]:
    """Instance discriminative loss over moving-class pixels.

    embeddings: [B, D, H, W]
    instances: [B, H, W], 0 is background/unknown
    seg_targets: [B, H, W], moving class is `moving_class`
    """

    bsz = int(embeddings.shape[0])
    total_var = embeddings.new_tensor(0.0)
    total_dist = embeddings.new_tensor(0.0)
    total_reg = embeddings.new_tensor(0.0)
    n_valid_samples = 0
    n_instance_pixels = 0

    emb_hw = embeddings.permute(0, 2, 3, 1).contiguous()  # [B, H, W, D]

    for b in range(bsz):
        inst = instances[b]
        seg = seg_targets[b]
        valid = seg != int(ignore_label)
        mask = torch.logical_and(valid, seg == int(moving_class))
        mask = torch.logical_and(mask, inst > 0)

        pix_count = int(torch.count_nonzero(mask).item())
        if pix_count <= 1:
            continue
        n_instance_pixels += pix_count

        ids = torch.unique(inst[mask])
        ids = ids[ids > 0]
        if ids.numel() == 0:
            continue

        sample_var = embeddings.new_tensor(0.0)
        centers: list[torch.Tensor] = []
        used_ids = 0

        for tid in ids:
            sel = torch.logical_and(mask, inst == tid)
            pix = emb_hw[b][sel]  # [N, D]
            if pix.ndim != 2 or pix.shape[0] == 0:
                continue
            center = torch.mean(pix, dim=0)
            centers.append(center)
            dist = torch.norm(pix - center.unsqueeze(0), dim=1)
            sample_var = sample_var + torch.mean(torch.clamp(dist - float(delta_var), min=0.0) ** 2)
            used_ids += 1

        if used_ids == 0:
            continue

        sample_var = sample_var / float(used_ids)
        centers_t = torch.stack(centers, dim=0)  # [K, D]
        sample_reg = torch.mean(torch.norm(centers_t, dim=1))

        sample_dist = embeddings.new_tensor(0.0)
        if centers_t.shape[0] > 1:
            c1 = centers_t.unsqueeze(0)
            c2 = centers_t.unsqueeze(1)
            dmat = torch.norm(c1 - c2, dim=2)
            eye = torch.eye(centers_t.shape[0], device=dmat.device, dtype=torch.bool)
            pair_d = dmat[~eye]
            if pair_d.numel() > 0:
                sample_dist = torch.mean(torch.clamp(2.0 * float(delta_dist) - pair_d, min=0.0) ** 2)

        total_var = total_var + sample_var
        total_dist = total_dist + sample_dist
        total_reg = total_reg + sample_reg
        n_valid_samples += 1

    if n_valid_samples == 0:
        zero = embeddings.sum() * 0.0
        return zero, {
            "embed_var": 0.0,
            "embed_dist": 0.0,
            "embed_reg": 0.0,
            "embed_loss": 0.0,
            "instance_pixels": float(n_instance_pixels),
            "instance_samples": 0.0,
        }

    var = total_var / float(n_valid_samples)
    dist = total_dist / float(n_valid_samples)
    reg = total_reg / float(n_valid_samples)
    loss = float(alpha) * var + float(beta) * dist + float(gamma) * reg
    return loss, {
        "embed_var": float(var.detach().item()),
        "embed_dist": float(dist.detach().item()),
        "embed_reg": float(reg.detach().item()),
        "embed_loss": float(loss.detach().item()),
        "instance_pixels": float(n_instance_pixels),
        "instance_samples": float(n_valid_samples),
    }


def run_epoch(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer | None,
    criterion_synth: nn.Module,
    device: torch.device,
    *,
    ignore_label: int,
    real_ignore_label: int,
    lambda_embed: float,
    delta_var: float,
    delta_dist: float,
    embed_alpha: float,
    embed_beta: float,
    embed_gamma: float,
    synth_seg_weight: float,
    real_seg_weight: float,
    synth_embed_weight: float,
    real_embed_weight: float,
    real_loss_mode: str,
    real_bootstrap_ignore_prob: float,
    real_bootstrap_max_frac: float,
    enable_real_bootstrap: bool,
) -> dict[str, float]:
    training = optimizer is not None
    model.train(training)

    total_loss = 0.0
    total_seg = 0.0
    total_embed = 0.0
    total_items = 0
    total_instance_pixels = 0.0
    total_instance_samples = 0.0
    total_seg_synth = 0.0
    total_seg_real = 0.0
    total_embed_synth = 0.0
    total_embed_real = 0.0
    seen_real = 0
    seen_synth = 0
    tp = fp = fn = 0.0

    for batch in loader:
        x = batch["image"].to(device)
        y = batch["label"].to(device)
        inst = batch["instance"].to(device)
        is_real = batch["is_real"].to(device).reshape(-1).bool()
        synth_mask = ~is_real

        if training:
            optimizer.zero_grad(set_to_none=True)

        seg_logits, embeddings = model(x)
        loss_zero = seg_logits.sum() * 0.0

        seg_parts: list[torch.Tensor] = []
        seg_weight_sum = 0.0
        seg_loss_synth = loss_zero
        seg_loss_real = loss_zero
        if bool(torch.any(synth_mask).item()):
            seg_loss_synth = criterion_synth(seg_logits[synth_mask], y[synth_mask])
            seg_parts.append(float(max(0.0, synth_seg_weight)) * seg_loss_synth)
            seg_weight_sum += float(max(0.0, synth_seg_weight))
            seen_synth += int(torch.count_nonzero(synth_mask).item())
        if bool(torch.any(is_real).item()):
            if str(real_loss_mode).lower() == "three_class":
                seg_loss_real = criterion_synth(seg_logits[is_real], y[is_real])
            else:
                bootstrap_prob = float(real_bootstrap_ignore_prob) if enable_real_bootstrap else 0.0
                seg_loss_real = real_loss_moving_vs_rest(
                    seg_logits[is_real],
                    y[is_real],
                    ignore_label=int(real_ignore_label),
                    bootstrap_ignore_prob=bootstrap_prob,
                    bootstrap_max_frac=float(real_bootstrap_max_frac),
                )
            seg_parts.append(float(max(0.0, real_seg_weight)) * seg_loss_real)
            seg_weight_sum += float(max(0.0, real_seg_weight))
            seen_real += int(torch.count_nonzero(is_real).item())
        if seg_parts and seg_weight_sum > 0.0:
            seg_loss = sum(seg_parts) / float(seg_weight_sum)
        else:
            seg_loss = loss_zero

        embed_parts: list[torch.Tensor] = []
        embed_weight_sum = 0.0
        embed_stats = {
            "embed_loss": 0.0,
            "instance_pixels": 0.0,
            "instance_samples": 0.0,
        }
        embed_loss_synth = loss_zero
        embed_loss_real = loss_zero
        if bool(torch.any(synth_mask).item()):
            embed_loss_synth, stats_synth = discriminative_embedding_loss(
                embeddings[synth_mask],
                inst[synth_mask],
                y[synth_mask],
                ignore_label=int(ignore_label),
                moving_class=1,
                delta_var=float(delta_var),
                delta_dist=float(delta_dist),
                alpha=float(embed_alpha),
                beta=float(embed_beta),
                gamma=float(embed_gamma),
            )
            embed_parts.append(float(max(0.0, synth_embed_weight)) * embed_loss_synth)
            embed_weight_sum += float(max(0.0, synth_embed_weight))
            embed_stats["instance_pixels"] += float(stats_synth["instance_pixels"])
            embed_stats["instance_samples"] += float(stats_synth["instance_samples"])
            total_embed_synth += float(stats_synth["embed_loss"]) * int(torch.count_nonzero(synth_mask).item())
        if bool(torch.any(is_real).item()):
            embed_loss_real, stats_real = discriminative_embedding_loss(
                embeddings[is_real],
                inst[is_real],
                y[is_real],
                ignore_label=int(real_ignore_label),
                moving_class=1,
                delta_var=float(delta_var),
                delta_dist=float(delta_dist),
                alpha=float(embed_alpha),
                beta=float(embed_beta),
                gamma=float(embed_gamma),
            )
            embed_parts.append(float(max(0.0, real_embed_weight)) * embed_loss_real)
            embed_weight_sum += float(max(0.0, real_embed_weight))
            embed_stats["instance_pixels"] += float(stats_real["instance_pixels"])
            embed_stats["instance_samples"] += float(stats_real["instance_samples"])
            total_embed_real += float(stats_real["embed_loss"]) * int(torch.count_nonzero(is_real).item())
        if embed_parts and embed_weight_sum > 0.0:
            embed_loss = sum(embed_parts) / float(embed_weight_sum)
            embed_stats["embed_loss"] = float(embed_loss.detach().item())
        else:
            embed_loss = loss_zero

        loss = seg_loss + float(lambda_embed) * embed_loss

        if training:
            loss.backward()
            optimizer.step()

        bsz = x.shape[0]
        total_loss += float(loss.item()) * bsz
        total_seg += float(seg_loss.item()) * bsz
        total_embed += float(embed_stats["embed_loss"]) * bsz
        total_items += bsz
        total_instance_pixels += float(embed_stats["instance_pixels"])
        total_instance_samples += float(embed_stats["instance_samples"])
        total_seg_synth += float(seg_loss_synth.item()) * int(torch.count_nonzero(synth_mask).item())
        total_seg_real += float(seg_loss_real.item()) * int(torch.count_nonzero(is_real).item())

        if bool(torch.any(synth_mask).item()):
            btp, bfp, bfn = moving_stats_from_batch(
                seg_logits[synth_mask].detach(),
                y[synth_mask],
                ignore_label=int(ignore_label),
            )
            tp += btp
            fp += bfp
            fn += bfn
        if bool(torch.any(is_real).item()):
            btp, bfp, bfn = moving_stats_from_batch(
                seg_logits[is_real].detach(),
                y[is_real],
                ignore_label=int(real_ignore_label),
            )
            tp += btp
            fp += bfp
            fn += bfn

    out = {
        "loss": total_loss / max(1, total_items),
        "seg_loss": total_seg / max(1, total_items),
        "embed_loss": total_embed / max(1, total_items),
        "instance_pixels": total_instance_pixels,
        "instance_samples": total_instance_samples,
        "seg_loss_synth": total_seg_synth / max(1, seen_synth),
        "seg_loss_real": total_seg_real / max(1, seen_real),
        "embed_loss_synth": total_embed_synth / max(1, seen_synth),
        "embed_loss_real": total_embed_real / max(1, seen_real),
        "seen_real": float(seen_real),
        "seen_synth": float(seen_synth),
        "real_fraction_seen": float(seen_real / max(1, seen_real + seen_synth)),
    }
    out.update(moving_metrics(tp, fp, fn))
    return out


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Train instance-aware kymograph model (segmentation + embeddings)")
    ap.add_argument("--train_manifest", required=True)
    ap.add_argument("--val_manifest", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--pretrained", default="", help="Optional seg/instance checkpoint to warm-start from")
    ap.add_argument("--epochs", type=int, default=30)
    ap.add_argument("--batch", type=int, default=8)
    ap.add_argument(
        "--val_batch",
        type=int,
        default=0,
        help="Validation batch size (0 uses --batch).",
    )
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--num_workers", type=int, default=0)
    ap.add_argument("--seed", type=int, default=123)
    ap.add_argument("--device", default=default_device())
    ap.add_argument("--base_channels", type=int, default=32)
    ap.add_argument("--embedding_dim", type=int, default=8)
    ap.add_argument("--lambda_embed", type=float, default=0.2)
    ap.add_argument("--delta_var", type=float, default=0.5)
    ap.add_argument("--delta_dist", type=float, default=1.5)
    ap.add_argument("--embed_alpha", type=float, default=1.0)
    ap.add_argument("--embed_beta", type=float, default=1.0)
    ap.add_argument("--embed_gamma", type=float, default=0.001)
    ap.add_argument("--label_pad_value", type=int, default=255)
    ap.add_argument("--include_static_instances", action="store_true")
    ap.add_argument("--instance_thickness", type=int, default=3)
    ap.add_argument(
        "--train_manifest_is_real",
        default="auto",
        help="Default is_real for train manifest rows lacking is_real (auto|true|false)",
    )
    ap.add_argument(
        "--val_manifest_is_real",
        default="auto",
        help="Default is_real for val manifest rows lacking is_real (auto|true|false)",
    )
    ap.add_argument(
        "--class_weights_from",
        choices=["all", "synth_only"],
        default="synth_only",
        help="Compute segmentation class weights from all train samples or synth-only samples.",
    )
    ap.add_argument("--synth_seg_weight", type=float, default=1.0)
    ap.add_argument("--real_seg_weight", type=float, default=1.0)
    ap.add_argument("--synth_embed_weight", type=float, default=1.0)
    ap.add_argument("--real_embed_weight", type=float, default=0.35)
    ap.add_argument(
        "--real_loss_mode",
        choices=["moving_vs_rest", "three_class"],
        default="moving_vs_rest",
        help="Real-data segmentation supervision mode.",
    )
    ap.add_argument("--real_ignore_label", type=int, default=255)
    ap.add_argument("--real_bootstrap_ignore_prob", type=float, default=0.90)
    ap.add_argument("--real_bootstrap_warmup_epochs", type=int, default=2)
    ap.add_argument("--real_bootstrap_max_frac", type=float, default=0.2)
    ap.add_argument("--freeze_backbone_epochs", type=int, default=0)
    ap.add_argument(
        "--synth_style_aug_prob",
        type=float,
        default=0.85,
        help="Style-augmentation probability for synthetic rows in training.",
    )
    ap.add_argument(
        "--real_style_aug_prob",
        type=float,
        default=0.25,
        help="Style-augmentation probability for real rows in training.",
    )
    ap.add_argument(
        "--style_aug_strength",
        type=float,
        default=1.0,
        help="Global strength multiplier for style augmentation.",
    )
    ap.add_argument(
        "--max_T",
        type=int,
        default=0,
        help="Optional max time-axis size; larger samples are cropped (0 disables).",
    )
    ap.add_argument(
        "--max_X",
        type=int,
        default=0,
        help="Optional max space-axis size; larger samples are cropped (0 disables).",
    )
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    for key in ("synth_style_aug_prob", "real_style_aug_prob"):
        v = float(getattr(args, key))
        if not (0.0 <= v <= 1.0):
            raise ValueError(f"{key} must be in [0, 1]")
    if float(args.style_aug_strength) < 0.0:
        raise ValueError("style_aug_strength must be >= 0")
    if int(args.max_T) < 0 or int(args.max_X) < 0:
        raise ValueError("max_T/max_X must be >= 0")
    if int(args.val_batch) < 0:
        raise ValueError("val_batch must be >= 0")
    seed_everything(args.seed)

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    tb_dir = out_dir / "tb"
    tb_dir.mkdir(parents=True, exist_ok=True)

    moving_only_instances = not bool(args.include_static_instances)
    train_default_is_real = _resolve_manifest_default_is_real(args.train_manifest_is_real, args.train_manifest)
    val_default_is_real = _resolve_manifest_default_is_real(args.val_manifest_is_real, args.val_manifest)
    print(
        f"[info] manifest defaults: train_is_real={train_default_is_real} "
        f"val_is_real={val_default_is_real}",
        flush=True,
    )
    train_ds = KymoInstanceDataset(
        args.train_manifest,
        is_real=train_default_is_real,
        augment=True,
        seed=args.seed,
        moving_only_instances=moving_only_instances,
        default_instance_thickness=int(args.instance_thickness),
        synth_style_aug_prob=float(args.synth_style_aug_prob),
        real_style_aug_prob=float(args.real_style_aug_prob),
        style_aug_strength=float(args.style_aug_strength),
    )
    val_ds = KymoInstanceDataset(
        args.val_manifest,
        is_real=val_default_is_real,
        augment=False,
        seed=args.seed,
        moving_only_instances=moving_only_instances,
        default_instance_thickness=int(args.instance_thickness),
        synth_style_aug_prob=float(args.synth_style_aug_prob),
        real_style_aug_prob=float(args.real_style_aug_prob),
        style_aug_strength=float(args.style_aug_strength),
    )
    print(f"[info] dataset sizes: train={len(train_ds)} val={len(val_ds)}", flush=True)

    device = torch.device(args.device)
    use_pin_memory = device.type == "cuda"
    loader_kwargs = {
        "num_workers": args.num_workers,
        "pin_memory": use_pin_memory,
        "persistent_workers": args.num_workers > 0,
    }
    if args.num_workers > 0:
        try:
            torch.multiprocessing.set_sharing_strategy("file_system")
        except Exception:
            pass
        loader_kwargs["prefetch_factor"] = 2

    train_collate = make_pad_collate(
        label_pad_value=int(args.label_pad_value),
        max_h=int(args.max_T),
        max_w=int(args.max_X),
        random_crop=True,
    )
    val_collate = make_pad_collate(
        label_pad_value=int(args.label_pad_value),
        max_h=int(args.max_T),
        max_w=int(args.max_X),
        random_crop=False,
    )
    val_batch = int(args.val_batch) if int(args.val_batch) > 0 else int(args.batch)
    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch,
        shuffle=True,
        collate_fn=train_collate,
        **loader_kwargs,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=val_batch,
        shuffle=False,
        collate_fn=val_collate,
        **loader_kwargs,
    )

    model = UNetInstance(
        in_channels=1,
        num_classes=3,
        base_channels=args.base_channels,
        embedding_dim=args.embedding_dim,
        normalize_embeddings=True,
    ).to(device)

    if args.pretrained:
        loaded, total = load_pretrained_into_instance(model, args.pretrained)
        print(f"[info] loaded pretrained tensors: {loaded}/{total} from {args.pretrained}")
        if loaded == 0:
            print("[warn] no compatible tensors were loaded; check base_channels/model family.")

    class_weights = make_class_weights(
        train_ds,
        log_progress=True,
        synth_only=(str(args.class_weights_from) == "synth_only"),
    ).to(device)
    criterion = nn.CrossEntropyLoss(weight=class_weights, ignore_index=int(args.label_pad_value))
    optimizer = AdamW(model.parameters(), lr=args.lr)
    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, args.epochs))

    freeze_epochs = int(max(0, args.freeze_backbone_epochs))
    if freeze_epochs > 0:
        set_backbone_trainable(model, False)
        print(f"[info] freezing backbone for first {freeze_epochs} epoch(s)", flush=True)

    writer = None
    try:
        from torch.utils.tensorboard import SummaryWriter

        writer = SummaryWriter(log_dir=str(tb_dir))
    except Exception:
        writer = None

    history_path = out_dir / "history.jsonl"
    best_f1 = -1.0

    for epoch in range(1, args.epochs + 1):
        if freeze_epochs > 0 and epoch == (freeze_epochs + 1):
            set_backbone_trainable(model, True)
            print("[info] unfroze backbone", flush=True)

        use_real_bootstrap = bool(epoch > int(args.real_bootstrap_warmup_epochs))
        train_metrics = run_epoch(
            model,
            train_loader,
            optimizer=optimizer,
            criterion_synth=criterion,
            device=device,
            ignore_label=int(args.label_pad_value),
            real_ignore_label=int(args.real_ignore_label),
            lambda_embed=float(args.lambda_embed),
            delta_var=float(args.delta_var),
            delta_dist=float(args.delta_dist),
            embed_alpha=float(args.embed_alpha),
            embed_beta=float(args.embed_beta),
            embed_gamma=float(args.embed_gamma),
            synth_seg_weight=float(args.synth_seg_weight),
            real_seg_weight=float(args.real_seg_weight),
            synth_embed_weight=float(args.synth_embed_weight),
            real_embed_weight=float(args.real_embed_weight),
            real_loss_mode=str(args.real_loss_mode),
            real_bootstrap_ignore_prob=float(args.real_bootstrap_ignore_prob),
            real_bootstrap_max_frac=float(args.real_bootstrap_max_frac),
            enable_real_bootstrap=use_real_bootstrap,
        )
        val_metrics = run_epoch(
            model,
            val_loader,
            optimizer=None,
            criterion_synth=criterion,
            device=device,
            ignore_label=int(args.label_pad_value),
            real_ignore_label=int(args.real_ignore_label),
            lambda_embed=float(args.lambda_embed),
            delta_var=float(args.delta_var),
            delta_dist=float(args.delta_dist),
            embed_alpha=float(args.embed_alpha),
            embed_beta=float(args.embed_beta),
            embed_gamma=float(args.embed_gamma),
            synth_seg_weight=float(args.synth_seg_weight),
            real_seg_weight=float(args.real_seg_weight),
            synth_embed_weight=float(args.synth_embed_weight),
            real_embed_weight=float(args.real_embed_weight),
            real_loss_mode=str(args.real_loss_mode),
            real_bootstrap_ignore_prob=float(args.real_bootstrap_ignore_prob),
            real_bootstrap_max_frac=float(args.real_bootstrap_max_frac),
            enable_real_bootstrap=False,
        )
        scheduler.step()

        row = {
            "epoch": epoch,
            "lr": optimizer.param_groups[0]["lr"],
            "use_real_bootstrap": bool(use_real_bootstrap),
            "train": train_metrics,
            "val": val_metrics,
        }
        with history_path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(row) + "\n")

        if writer is not None:
            writer.add_scalar("train/loss", train_metrics["loss"], epoch)
            writer.add_scalar("train/seg_loss", train_metrics["seg_loss"], epoch)
            writer.add_scalar("train/embed_loss", train_metrics["embed_loss"], epoch)
            writer.add_scalar("train/seg_loss_real", train_metrics["seg_loss_real"], epoch)
            writer.add_scalar("train/seg_loss_synth", train_metrics["seg_loss_synth"], epoch)
            writer.add_scalar("train/embed_loss_real", train_metrics["embed_loss_real"], epoch)
            writer.add_scalar("train/embed_loss_synth", train_metrics["embed_loss_synth"], epoch)
            writer.add_scalar("train/real_fraction_seen", train_metrics["real_fraction_seen"], epoch)
            writer.add_scalar("train/moving_f1", train_metrics["moving_f1"], epoch)
            writer.add_scalar("val/loss", val_metrics["loss"], epoch)
            writer.add_scalar("val/seg_loss", val_metrics["seg_loss"], epoch)
            writer.add_scalar("val/embed_loss", val_metrics["embed_loss"], epoch)
            writer.add_scalar("val/seg_loss_real", val_metrics["seg_loss_real"], epoch)
            writer.add_scalar("val/seg_loss_synth", val_metrics["seg_loss_synth"], epoch)
            writer.add_scalar("val/embed_loss_real", val_metrics["embed_loss_real"], epoch)
            writer.add_scalar("val/embed_loss_synth", val_metrics["embed_loss_synth"], epoch)
            writer.add_scalar("val/real_fraction_seen", val_metrics["real_fraction_seen"], epoch)
            writer.add_scalar("val/moving_f1", val_metrics["moving_f1"], epoch)

        ckpt = {
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
            "epoch": epoch,
            "args": vars(args),
            "val_metrics": val_metrics,
        }
        torch.save(ckpt, out_dir / "last.pt")

        if val_metrics["moving_f1"] >= best_f1:
            best_f1 = val_metrics["moving_f1"]
            torch.save(ckpt, out_dir / "best.pt")

        print(
            "epoch {ep:03d} | train loss={tl:.4f} seg={ts:.4f} emb={te:.4f} f1={tf:.4f}"
            " | val loss={vl:.4f} seg={vs:.4f} emb={ve:.4f} f1={vf:.4f} real_frac={vr:.3f}".format(
                ep=epoch,
                tl=float(train_metrics["loss"]),
                ts=float(train_metrics["seg_loss"]),
                te=float(train_metrics["embed_loss"]),
                tf=float(train_metrics["moving_f1"]),
                vl=float(val_metrics["loss"]),
                vs=float(val_metrics["seg_loss"]),
                ve=float(val_metrics["embed_loss"]),
                vf=float(val_metrics["moving_f1"]),
                vr=float(val_metrics["real_fraction_seen"]),
            )
        )

    if writer is not None:
        writer.close()

    summary = {
        "best_moving_f1": float(best_f1),
        "class_weights": class_weights.detach().cpu().numpy().tolist(),
        "embedding_dim": int(args.embedding_dim),
        "lambda_embed": float(args.lambda_embed),
        "delta_var": float(args.delta_var),
        "delta_dist": float(args.delta_dist),
        "class_weights_from": str(args.class_weights_from),
        "real_loss_mode": str(args.real_loss_mode),
        "real_bootstrap_ignore_prob": float(args.real_bootstrap_ignore_prob),
        "real_bootstrap_warmup_epochs": int(args.real_bootstrap_warmup_epochs),
        "real_bootstrap_max_frac": float(args.real_bootstrap_max_frac),
        "freeze_backbone_epochs": int(args.freeze_backbone_epochs),
        "synth_seg_weight": float(args.synth_seg_weight),
        "real_seg_weight": float(args.real_seg_weight),
        "synth_embed_weight": float(args.synth_embed_weight),
        "real_embed_weight": float(args.real_embed_weight),
        "synth_style_aug_prob": float(args.synth_style_aug_prob),
        "real_style_aug_prob": float(args.real_style_aug_prob),
        "style_aug_strength": float(args.style_aug_strength),
        "max_T": int(args.max_T),
        "max_X": int(args.max_X),
        "val_batch": int(val_batch),
    }
    (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()
