from __future__ import annotations

import argparse
import json
import random
import time
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 KymoSegmentationDataset, compute_class_histogram
from .models import UNetSmall
from .train_seg import moving_metrics, seed_everything


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 _next_batch(it, loader):
    try:
        return next(it), it
    except StopIteration:
        it = iter(loader)
        return next(it), it


def _maybe_crop_hw(
    img: torch.Tensor,
    lbl: torch.Tensor,
    *,
    max_h: int,
    max_w: int,
    random_crop: bool,
) -> tuple[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, :]
        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]
    return img, lbl


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 size in that batch."""

    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"]
            h = int(img.shape[-2])
            w = int(img.shape[-1])
            if cap_h > 0:
                h = min(h, cap_h)
            if cap_w > 0:
                w = min(w, cap_w)
            batch_h = max(batch_h, h)
            batch_w = max(batch_w, w)

        images = []
        labels = []
        is_real = []
        image_paths = []
        label_paths = []

        for item in batch:
            img = item["image"]
            lbl = item["label"]
            img, lbl = _maybe_crop_hw(
                img,
                lbl,
                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))

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

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

    return _collate


def _masked_moving_stats(logits: torch.Tensor, targets: torch.Tensor, ignore_label: int) -> tuple[float, float, float]:
    pred = torch.argmax(logits, dim=1)
    valid = targets != int(ignore_label)
    pred_m = torch.logical_and(pred == 1, valid)
    gt_m = torch.logical_and(targets == 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 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 non-moving with optional ignore masking.

    - labels==ignore_label are excluded from the loss.
    - optional bootstrap ignore masks high-confidence moving predictions on unlabeled pixels
      to reduce false-negative pressure from incomplete manual annotations.
    """
    valid_mask = targets != int(ignore_label)
    targets = 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 bootstrap_ignore_prob > 0:
        with torch.no_grad():
            probs = torch.softmax(two_logits, dim=1)[:, 1]
            candidates = torch.logical_and(valid_mask, targets == 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, targets, 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 run_val_real(model: nn.Module, loader: DataLoader, device: torch.device, ignore_label: int) -> dict[str, float]:
    model.eval()
    total_loss = 0.0
    total_items = 0
    tp = fp = fn = 0.0

    with torch.no_grad():
        for batch in loader:
            x = batch["image"].to(device)
            y = batch["label"].to(device)
            logits = model(x)
            loss = real_loss_moving_vs_rest(logits, y, ignore_label=ignore_label)

            bsz = x.shape[0]
            total_loss += float(loss.item()) * bsz
            total_items += bsz

            btp, bfp, bfn = _masked_moving_stats(logits, y, ignore_label=ignore_label)
            tp += btp
            fp += bfp
            fn += bfn

    metrics = {"loss": total_loss / max(1, total_items)}
    metrics.update(moving_metrics(tp, fp, fn))
    return metrics


def make_synth_class_weights(dataset, max_samples: int | None = None) -> torch.Tensor:
    hist = compute_class_histogram(dataset, num_classes=3, max_samples=max_samples)
    hist = np.asarray(hist, dtype=np.float64)
    present = hist > 0
    if not np.any(present):
        w = np.ones((3,), dtype=np.float64)
    else:
        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:
            w = np.ones((3,), dtype=np.float64)
        else:
            w = np.zeros_like(hist, dtype=np.float64)
            w[present] = (inv[present] / z) * float(np.count_nonzero(present))
    return torch.tensor(w, dtype=torch.float32)


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Fine-tune segmentation model on real labeled kymographs")
    ap.add_argument("--pretrained", required=True)
    ap.add_argument("--real_train_manifest", required=True)
    ap.add_argument("--real_val_manifest", required=True)
    ap.add_argument("--synth_train_manifest", default=None)
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=int, default=20)
    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("--mix_real", type=float, default=0.3)
    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(
        "--synth_hist_max_samples",
        type=int,
        default=512,
        help="Max synth samples to use for class-weight histogram (0 means all samples).",
    )
    ap.add_argument(
        "--real_ignore_label",
        type=int,
        default=255,
        help="Label value in real masks to ignore in loss/metrics. Use 255 for unlabeled/unknown.",
    )
    ap.add_argument(
        "--real_bootstrap_ignore_prob",
        type=float,
        default=0.90,
        help="If >0, ignore unlabeled real pixels with predicted moving prob above this threshold.",
    )
    ap.add_argument(
        "--real_bootstrap_warmup_epochs",
        type=int,
        default=2,
        help="Enable bootstrap ignore after this many epochs.",
    )
    ap.add_argument(
        "--real_bootstrap_max_frac",
        type=float,
        default=0.2,
        help="Max fraction of valid real pixels to drop via bootstrap ignore per batch.",
    )
    ap.add_argument(
        "--log_every_steps",
        type=int,
        default=25,
        help="Print intra-epoch training progress every N steps (<=0 disables).",
    )
    ap.add_argument(
        "--synth_style_aug_prob",
        type=float,
        default=0.85,
        help="Style-augmentation probability for synthetic training samples.",
    )
    ap.add_argument(
        "--real_style_aug_prob",
        type=float,
        default=0.25,
        help="Style-augmentation probability for real training samples.",
    )
    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 _resolve_device(requested: str) -> str:
    req = str(requested).lower().strip()
    mps_ok = bool(getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available())
    cuda_ok = bool(torch.cuda.is_available())

    if req == "mps" and not mps_ok:
        print(
            "[warn] Requested --device mps but MPS is not available in this runtime. "
            "Falling back to cpu."
        )
        return "cpu"
    if req == "cuda" and not cuda_ok:
        print(
            "[warn] Requested --device cuda but CUDA is not available in this runtime. "
            "Falling back to cpu."
        )
        return "cpu"
    return req


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)
    resolved_device = _resolve_device(args.device)
    mps_built = bool(getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_built())
    mps_ok = bool(getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available())
    cuda_ok = bool(torch.cuda.is_available())
    print(
        f"[info] device request={args.device} resolved={resolved_device} "
        f"(mps_built={mps_built}, mps_available={mps_ok}, cuda_available={cuda_ok})"
    )

    device = torch.device(resolved_device)
    print("[info] loading manifests and datasets...")
    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

    real_train_ds = KymoSegmentationDataset(
        args.real_train_manifest,
        is_real=True,
        augment=True,
        seed=args.seed,
        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),
    )
    real_val_ds = KymoSegmentationDataset(
        args.real_val_manifest,
        is_real=True,
        augment=False,
        seed=args.seed,
        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),
    )

    synth_train_ds = None
    if args.synth_train_manifest:
        synth_train_ds = KymoSegmentationDataset(
            args.synth_train_manifest,
            is_real=False,
            augment=True,
            seed=args.seed,
            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: real_train={len(real_train_ds)} real_val={len(real_val_ds)} "
        f"synth_train={(len(synth_train_ds) if synth_train_ds is not None else 0)}"
    )
    print("[info] building dataloaders...")

    train_collate_real = make_pad_collate(
        label_pad_value=int(args.real_ignore_label),
        max_h=int(args.max_T),
        max_w=int(args.max_X),
        random_crop=True,
    )
    train_collate_synth = make_pad_collate(
        label_pad_value=0,
        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.real_ignore_label),
        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)

    real_loader = DataLoader(
        real_train_ds,
        batch_size=args.batch,
        shuffle=True,
        collate_fn=train_collate_real,
        **loader_kwargs,
    )
    real_val_loader = DataLoader(
        real_val_ds,
        batch_size=val_batch,
        shuffle=False,
        collate_fn=val_collate,
        **loader_kwargs,
    )
    synth_loader = (
        DataLoader(
            synth_train_ds,
            batch_size=args.batch,
            shuffle=True,
            collate_fn=train_collate_synth,
            **loader_kwargs,
        )
        if synth_train_ds is not None
        else None
    )

    model = UNetSmall(in_channels=1, num_classes=3, base_channels=args.base_channels).to(device)

    print(f"[info] loading checkpoint: {args.pretrained}")
    ckpt = torch.load(args.pretrained, map_location="cpu")
    state_dict = ckpt.get("model_state") if isinstance(ckpt, dict) and "model_state" in ckpt else ckpt
    model.load_state_dict(state_dict, strict=True)

    optimizer = AdamW(model.parameters(), lr=args.lr)
    scheduler = CosineAnnealingLR(optimizer, T_max=max(1, args.epochs))

    synth_class_weights = None
    if synth_train_ds is not None:
        max_samples = None if int(args.synth_hist_max_samples) <= 0 else int(args.synth_hist_max_samples)
        if max_samples is None:
            print("[info] computing synth class weights on all synth samples...")
        else:
            print(f"[info] computing synth class weights on up to {max_samples} synth samples...")
        synth_class_weights = make_synth_class_weights(synth_train_ds, max_samples=max_samples).to(device)
        print(f"[info] synth class weights: {synth_class_weights.detach().cpu().tolist()}")
    synth_criterion = nn.CrossEntropyLoss(weight=synth_class_weights)

    history_path = out_dir / "history.jsonl"
    best_f1 = -1.0
    print("[info] starting training loop...")

    for epoch in range(1, args.epochs + 1):
        model.train()
        total_loss = 0.0
        steps = 0
        real_steps = 0
        synth_steps = 0

        real_it = iter(real_loader)
        synth_it = iter(synth_loader) if synth_loader is not None else None

        n_real_steps = max(1, len(real_loader))
        n_synth_steps = max(1, len(synth_loader)) if synth_loader is not None else 0
        steps_per_epoch = max(n_real_steps, n_synth_steps) if synth_loader is not None else n_real_steps

        epoch_start = time.perf_counter()
        for step_idx in range(1, steps_per_epoch + 1):
            use_real = synth_loader is None or (random.random() < float(args.mix_real))
            if use_real:
                batch, real_it = _next_batch(real_it, real_loader)
                real_steps += 1
            else:
                batch, synth_it = _next_batch(synth_it, synth_loader)
                synth_steps += 1

            x = batch["image"].to(device)
            y = batch["label"].to(device)

            optimizer.zero_grad(set_to_none=True)
            logits = model(x)

            if use_real:
                bootstrap_prob = None
                if int(epoch) > int(args.real_bootstrap_warmup_epochs):
                    if float(args.real_bootstrap_ignore_prob) > 0:
                        bootstrap_prob = float(args.real_bootstrap_ignore_prob)
                loss = real_loss_moving_vs_rest(
                    logits,
                    y,
                    ignore_label=int(args.real_ignore_label),
                    bootstrap_ignore_prob=bootstrap_prob,
                    bootstrap_max_frac=float(args.real_bootstrap_max_frac),
                )
            else:
                loss = synth_criterion(logits, y)

            loss.backward()
            optimizer.step()

            total_loss += float(loss.item())
            steps += 1

            log_every = int(args.log_every_steps)
            should_log = (
                log_every > 0
                and (
                    step_idx == 1
                    or step_idx == steps_per_epoch
                    or (step_idx % log_every == 0)
                )
            )
            if should_log:
                elapsed = time.perf_counter() - epoch_start
                it_s = float(step_idx / max(1e-6, elapsed))
                avg_loss = float(total_loss / max(1, steps))
                print(
                    f"[train] epoch {epoch:03d} step {step_idx:04d}/{steps_per_epoch:04d} "
                    f"loss={float(loss.item()):.4f} avg={avg_loss:.4f} "
                    f"real_steps={real_steps} synth_steps={synth_steps} "
                    f"speed={it_s:.2f} it/s",
                    flush=True,
                )

        scheduler.step()

        train_loss = total_loss / max(1, steps)
        print(f"[info] epoch {epoch:03d} running validation...", flush=True)
        val_metrics = run_val_real(model, real_val_loader, device=device, ignore_label=int(args.real_ignore_label))

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

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

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

        print(
            f"epoch {epoch:03d} | train loss={train_loss:.4f} | "
            f"val loss={val_metrics['loss']:.4f} f1={val_metrics['moving_f1']:.4f}"
        )

    summary = {
        "best_moving_f1": best_f1,
        "mix_real": float(args.mix_real),
        "used_synth": synth_train_ds is not None,
        "real_ignore_label": int(args.real_ignore_label),
        "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),
        "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()
