import csv
import math
import os
import signal
import time
from pathlib import Path
from types import FrameType
from typing import Callable

import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm

from giant.model.schedule import (
    CosineSchedule,
    flow_matching_loss,
    flow_matching_loss_secondary,
)
from giant.validate import validate_marginals

_METRICS_FIELDS = [
    "epoch",
    "train_loss",
    "train_loss_s1",
    "train_loss_nsec",
    "train_loss_s2",
    "val_loss",
    "val_loss_s1",
    "val_loss_nsec",
    "val_loss_s2",
    "lr",
    "epoch_time_s",
]

_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)


class _GracefulShutdown:
    """Turns SIGINT/SIGTERM into a flag check instead of an immediate crash.

    A second signal while already shutting down restores the default
    handler and re-sends the signal, so an unresponsive run can still be
    force-killed.
    """

    def __init__(self) -> None:
        self.requested = False
        self._previous: dict[
            int,
            Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
        ] = {}

    def __enter__(self) -> "_GracefulShutdown":
        for sig in _CATCHABLE_SIGNALS:
            self._previous[sig] = signal.getsignal(sig)
            signal.signal(sig, self._handle)
        return self

    def __exit__(self, *exc_info) -> None:
        for sig, handler in self._previous.items():
            signal.signal(sig, handler)

    def _handle(self, signum: int, frame) -> None:
        if self.requested:
            signal.signal(signum, self._previous[signum])
            os.kill(os.getpid(), signum)
            return
        self.requested = True
        print(
            f"\nreceived {signal.Signals(signum).name} — finishing the current "
            "batch, then saving a checkpoint and exiting (send again to force-quit)"
        )


def _build_sec_x1(
    sec_cont: torch.Tensor,
    sec_pdg_idx: torch.Tensor,
    pdg_emb_weight: torch.Tensor,
) -> torch.Tensor:
    """Assemble the Stage-2 flow target by appending type embeddings.

    sec_cont: (B, K_MAX, 4) — [stick_logit, dir_local]
    sec_pdg_idx: (B, K_MAX) — integer PDG model-indices
    pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights

    Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)).

    Detaches the looked-up rows: this tensor becomes x1 in the flow-matching
    loss (u_t = x1 - x0), so without detaching, the Stage-2 loss could pull
    the embedding table itself toward whatever the decoder already predicts
    (a moving, self-referential regression target) instead of only pulling
    the decoder toward the table. The table is still trained normally via
    its Stage-1 conditioning role and `predict_n_sec`.
    """
    type_emb = pdg_emb_weight[sec_pdg_idx].detach()  # (B, K_MAX, emb_dim)
    x1_s2 = torch.cat([sec_cont, type_emb], dim=-1)  # (B, K_MAX, 4+emb_dim)
    return x1_s2.flatten(1)  # (B, SEC_DIM)


def _compute_losses(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    batch: tuple,
    mode: str,
    ddpm_schedule,
    device: torch.device,
    lambda_nsec: float,
    lambda_s2: float,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Compute (total_loss, L_s1, L_nsec, L_s2) for one batch."""
    cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx = batch
    cond_cont = cond_cont.to(device)
    cond_cat = cond_cat.to(device)
    x1_s1 = x1_s1.to(device)
    n_sec = n_sec.to(device)
    sec_cont = sec_cont.to(device)
    sec_pdg_idx = sec_pdg_idx.to(device)

    # Stage-1 flow loss
    if mode == "flow":
        l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
    else:
        assert ddpm_schedule is not None
        l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat)

    # n_sec classification loss
    n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat)
    l_nsec = F.cross_entropy(n_sec_logits, n_sec)

    # Stage-2 secondary flow loss
    # Use a noiseless Stage-1 target as context (detach to avoid back-prop
    # coupling between the two flow paths through the same embedding table).
    # The type-embedding lookup itself is also detached inside _build_sec_x1,
    # so the shared PDG table is shaped only by its Stage-1 conditioning role
    # and predict_n_sec, not by chasing the Stage-2 decoder's predictions.
    from giant.constants import K_MAX

    pdg_emb_weight = stage1_model.pdg_embedding_weight()
    x1_s2 = _build_sec_x1(sec_cont, sec_pdg_idx, pdg_emb_weight)

    sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
    l_s2 = flow_matching_loss_secondary(
        sec_decoder,
        x1_s2,
        cond_cont,
        cond_cat,
        x1_s1.detach(),
        sec_mask,
    )

    total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
    return total, l_s1, l_nsec, l_s2


def train(
    stage1_model: torch.nn.Module,
    sec_decoder: torch.nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    mode: str,
    epochs: int,
    lr: float,
    warmup_epochs: int,
    device: torch.device,
    out_dir: str | Path,
    lambda_nsec: float = 0.1,
    lambda_s2: float = 1.0,
    normalizer_dict: dict | None = None,
    pdg_map: dict | None = None,
    mat_map: dict | None = None,
    model_config: dict | None = None,
    resume_path: str | Path | None = None,
    validate_every: int = 0,
    validate_steps: int = 10,
    total_train_batches: int = 0,
) -> None:
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    stage1_model = stage1_model.to(device)
    sec_decoder = sec_decoder.to(device)

    all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters())
    optimizer = optim.AdamW(all_params, lr=lr)

    def _lr_lambda(epoch: int) -> float:
        if warmup_epochs > 0 and epoch < warmup_epochs:
            return (epoch + 1) / warmup_epochs
        t = epoch - warmup_epochs
        T = max(epochs - warmup_epochs, 1)
        return 0.5 * (1.0 + math.cos(math.pi * t / T))

    lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda)

    ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None

    start_epoch = 1
    best_val_loss = float("inf")
    if resume_path is not None:
        ckpt = torch.load(resume_path, map_location=device, weights_only=False)
        stage1_model.load_state_dict(ckpt["model"])
        sec_decoder.load_state_dict(ckpt["sec_decoder"])
        optimizer.load_state_dict(ckpt["optimizer"])
        lr_sched.load_state_dict(ckpt["lr_sched"])
        start_epoch = ckpt.get("epoch", 0) + 1
        best_val_loss = ckpt.get("best_val_loss", float("inf"))

        # optimizer/lr_sched.load_state_dict() above restore the checkpoint's
        # own base LR, which would otherwise silently override an explicit
        # `lr` argument. Make `lr` authoritative again, applied at whatever
        # point the cosine/warmup schedule has already reached.
        lr_sched.base_lrs = [lr for _ in lr_sched.base_lrs]
        resumed_lr = lr * _lr_lambda(lr_sched.last_epoch)
        for group in optimizer.param_groups:
            group["lr"] = resumed_lr

        if start_epoch > epochs:
            print(
                f"checkpoint already completed epoch {start_epoch - 1} "
                f"(>= --epochs {epochs}) — nothing to train"
            )
            return

    metrics_path = out_dir / "metrics.csv"
    resuming_existing_metrics = resume_path is not None and metrics_path.exists()
    write_header = not resuming_existing_metrics
    metrics_file = open(
        metrics_path, "a" if resuming_existing_metrics else "w", newline=""
    )
    metrics_writer = csv.DictWriter(metrics_file, fieldnames=_METRICS_FIELDS)
    if write_header:
        metrics_writer.writeheader()

    epoch_w = len(str(epochs))
    last_completed_epoch = start_epoch - 1
    with _GracefulShutdown() as shutdown:
        for epoch in range(start_epoch, epochs + 1):
            epoch_start = time.monotonic()
            current_lr = optimizer.param_groups[0]["lr"]
            stage1_model.train()
            sec_decoder.train()
            train_loss_sum = 0.0
            train_s1_sum = 0.0
            train_nsec_sum = 0.0
            train_s2_sum = 0.0
            train_n = 0
            ema_loss = 0.0
            bar = tqdm(
                train_loader,
                desc=f"  epoch {epoch:{epoch_w}d}/{epochs}",
                total=total_train_batches or None,
                leave=False,
                unit="batch",
                dynamic_ncols=True,
            )
            for batch in bar:
                loss, l_s1, l_nsec, l_s2 = _compute_losses(
                    stage1_model,
                    sec_decoder,
                    batch,
                    mode,
                    ddpm_schedule,
                    device,
                    lambda_nsec,
                    lambda_s2,
                )
                optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(all_params, 1.0)
                optimizer.step()

                B = batch[0].size(0)
                batch_loss = loss.item()
                train_loss_sum += batch_loss * B
                train_s1_sum += l_s1.item() * B
                train_nsec_sum += l_nsec.item() * B
                train_s2_sum += l_s2.item() * B
                train_n += B
                ema_loss = (
                    batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss
                )
                bar.set_postfix_str(f"loss={ema_loss:.4f}", refresh=False)

                if shutdown.requested:
                    break
            bar.close()

            if shutdown.requested:
                break

            train_loss = train_loss_sum / max(train_n, 1)
            lr_sched.step()

            stage1_model.eval()
            sec_decoder.eval()
            val_loss_sum = 0.0
            val_s1_sum = 0.0
            val_nsec_sum = 0.0
            val_s2_sum = 0.0
            val_n = 0
            with torch.no_grad():
                for batch in val_loader:
                    loss, l_s1, l_nsec, l_s2 = _compute_losses(
                        stage1_model,
                        sec_decoder,
                        batch,
                        mode,
                        ddpm_schedule,
                        device,
                        lambda_nsec,
                        lambda_s2,
                    )
                    B = batch[0].size(0)
                    val_loss_sum += loss.item() * B
                    val_s1_sum += l_s1.item() * B
                    val_nsec_sum += l_nsec.item() * B
                    val_s2_sum += l_s2.item() * B
                    val_n += B
            val_loss = val_loss_sum / max(val_n, 1)
            epoch_time = time.monotonic() - epoch_start

            is_best = val_loss < best_val_loss
            marker = "  [best]" if is_best else ""
            print(
                f"epoch {epoch:{epoch_w}d}/{epochs}"
                f"  train {train_loss:.4f}"
                f"  (s1={train_s1_sum / max(train_n, 1):.3f}"
                f"  nsec={train_nsec_sum / max(train_n, 1):.3f}"
                f"  s2={train_s2_sum / max(train_n, 1):.3f})"
                f"  val {val_loss:.4f}"
                f"  lr {current_lr:.2e}  {epoch_time:.1f}s{marker}"
            )
            metrics_writer.writerow(
                {
                    "epoch": epoch,
                    "train_loss": train_loss,
                    "train_loss_s1": train_s1_sum / max(train_n, 1),
                    "train_loss_nsec": train_nsec_sum / max(train_n, 1),
                    "train_loss_s2": train_s2_sum / max(train_n, 1),
                    "val_loss": val_loss,
                    "val_loss_s1": val_s1_sum / max(val_n, 1),
                    "val_loss_nsec": val_nsec_sum / max(val_n, 1),
                    "val_loss_s2": val_s2_sum / max(val_n, 1),
                    "lr": current_lr,
                    "epoch_time_s": epoch_time,
                }
            )
            metrics_file.flush()

            if validate_every > 0 and epoch % validate_every == 0:
                print(f"[epoch {epoch}] marginal validation:")
                validate_marginals(
                    stage1_model,
                    val_loader,
                    mode=mode,
                    schedule=ddpm_schedule,
                    device=device,
                    steps=validate_steps,
                    sec_decoder=sec_decoder,
                )

            ckpt: dict = {
                "model": stage1_model.state_dict(),
                "sec_decoder": sec_decoder.state_dict(),
                "optimizer": optimizer.state_dict(),
                "lr_sched": lr_sched.state_dict(),
                "epoch": epoch,
                "best_val_loss": best_val_loss,
            }
            if normalizer_dict is not None:
                ckpt["normalizer"] = normalizer_dict
            if pdg_map is not None:
                ckpt["pdg_map"] = pdg_map
            if mat_map is not None:
                ckpt["mat_map"] = mat_map
            if model_config is not None:
                ckpt["model_config"] = model_config

            if val_loss < best_val_loss:
                best_val_loss = val_loss
                ckpt["best_val_loss"] = best_val_loss
                torch.save(ckpt, out_dir / "best.pt")

            torch.save(ckpt, out_dir / "last.pt")
            last_completed_epoch = epoch

            if shutdown.requested:
                break

    metrics_file.close()

    if shutdown.requested:
        print(
            f"stopped after epoch {last_completed_epoch} due to shutdown signal — "
            f"resume with --resume {out_dir / 'last.pt'}"
        )
