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 KymoSegmentationDataset, compute_class_histogram
from .models import UNetSmall


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 run_epoch(
    model: nn.Module,
    loader: DataLoader,
    optimizer: torch.optim.Optimizer | None,
    criterion: nn.Module,
    device: torch.device,
    *,
    ignore_label: int = 255,
) -> dict[str, float]:
    training = optimizer is not None
    model.train(training)

    total_loss = 0.0
    total_items = 0
    tp = fp = fn = 0.0

    for batch in loader:
        x = batch["image"].to(device)
        y = batch["label"].to(device)

        if training:
            optimizer.zero_grad(set_to_none=True)

        logits = model(x)
        loss = criterion(logits, y)

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

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

        btp, bfp, bfn = moving_stats_from_batch(logits.detach(), y, ignore_label=int(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 _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) -> torch.Tensor:
    n = len(dataset)
    progress_every = max(1, n // 20) if (log_progress and n > 0) else 0
    if log_progress:
        print(f"[info] computing class histogram over {n} training sample(s)...", flush=True)
    hist = compute_class_histogram(
        dataset,
        num_classes=num_classes,
        progress_every=progress_every,
        progress_prefix="class histogram",
    )
    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 _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"]
            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 = []
        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 parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Train kymograph segmentation on synthetic data")
    ap.add_argument("--train_manifest", required=True)
    ap.add_argument("--val_manifest", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=int, default=30)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument(
        "--val_batch",
        type=int,
        default=0,
        help="Validation batch size (0 uses --batch).",
    )
    ap.add_argument("--lr", type=float, default=1e-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_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 samples (used only if manifest rows are tagged real).",
    )
    ap.add_argument(
        "--style_aug_strength",
        type=float,
        default=1.0,
        help="Global strength multiplier for style augmentation.",
    )
    ap.add_argument(
        "--label_pad_value",
        type=int,
        default=255,
        help="Padding label value for variable-size batches; ignored in loss/metrics.",
    )
    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.label_pad_value) < 0:
        raise ValueError("label_pad_value 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)

    train_ds = KymoSegmentationDataset(
        args.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),
    )
    val_ds = KymoSegmentationDataset(
        args.val_manifest,
        is_real=False,
        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),
    )
    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 = UNetSmall(in_channels=1, num_classes=3, base_channels=args.base_channels).to(device)

    class_weights = make_class_weights(train_ds, log_progress=True).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))

    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):
        train_metrics = run_epoch(
            model,
            train_loader,
            optimizer=optimizer,
            criterion=criterion,
            device=device,
            ignore_label=int(args.label_pad_value),
        )
        val_metrics = run_epoch(
            model,
            val_loader,
            optimizer=None,
            criterion=criterion,
            device=device,
            ignore_label=int(args.label_pad_value),
        )
        scheduler.step()

        row = {
            "epoch": epoch,
            "lr": optimizer.param_groups[0]["lr"],
            "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/moving_f1", train_metrics["moving_f1"], epoch)
            writer.add_scalar("val/loss", val_metrics["loss"], 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(
            f"epoch {epoch:03d} | train loss={train_metrics['loss']:.4f} f1={train_metrics['moving_f1']:.4f}"
            f" | val loss={val_metrics['loss']:.4f} f1={val_metrics['moving_f1']:.4f}"
        )

    if writer is not None:
        writer.close()

    summary = {
        "best_moving_f1": best_f1,
        "class_weights": class_weights.detach().cpu().numpy().tolist(),
        "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),
        "label_pad_value": int(args.label_pad_value),
        "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()
