"""Training loop: AdamW, linear warm-up + cosine decay, gradient clipping, early stopping on the validation
probabilistic score, fixed seeds, crash-safe checkpoints with resume, JSONL tracking, throughput/memory measurement."""
from __future__ import annotations

import math
import os
import random
import time
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

from ..models.model import MRPT, MRPTConfig
from .losses import forecast_loss
from .sampling import OriginDataset, collate
from .tracking import Tracker
from . import checkpoint as ckpt


def seed_all(seed: int) -> None:
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)


def pick_device(name: str | None = None) -> torch.device:
    if name:
        return torch.device(name)
    if torch.backends.mps.is_available():
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def to_device(batch: dict[str, torch.Tensor], dev: torch.device) -> dict[str, torch.Tensor]:
    return {k: v.to(dev, non_blocking=True) for k, v in batch.items()}


def lr_lambda(warmup: int, total: int):
    def f(step):
        if step < warmup:
            return (step + 1) / max(1, warmup)
        p = (step - warmup) / max(1, total - warmup)
        return 0.5 * (1 + math.cos(math.pi * min(1.0, p))) * 0.9 + 0.1
    return f


def mem_mb(dev: torch.device) -> float:
    if dev.type == "mps":
        return torch.mps.current_allocated_memory() / 1e6
    if dev.type == "cuda":
        return torch.cuda.max_memory_allocated() / 1e6
    return 0.0


@torch.no_grad()
def evaluate_nll(model: MRPT, loader: DataLoader, dev: torch.device, head: str, levels, max_batches: int | None = None):
    model.eval()
    tot, n = 0.0, 0
    for i, b in enumerate(loader):
        if max_batches is not None and i >= max_batches:
            break
        b = to_device(b, dev)
        out = model(b)
        loss, _ = forecast_loss(out, b, head, levels)
        k = int(b["y_mask"].sum())
        tot += float(loss) * k; n += k
    model.train()
    return tot / max(1, n)


def train(model: MRPT, train_ds: OriginDataset, val_ds: OriginDataset, out_dir: Path, *, batch_size: int = 64,
          max_steps: int = 1000, lr: float = 3e-4, warmup_steps: int = 100, weight_decay: float = 0.01,
          grad_clip: float = 1.0, eval_every: int = 200, patience: int = 5, seed: int = 1, num_workers: int = 4,
          device: str | None = None, aux_vol_weight: float = 0.0, resume: bool = True, val_max_batches: int = 50,
          copula_only: bool = False) -> dict:
    seed_all(seed)
    dev = pick_device(device)
    model.to(dev)
    head = model.cfg.head
    params = [p for n, p in model.named_parameters() if (not copula_only) or n.startswith("copula.")]
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay, betas=(0.9, 0.98))
    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_lambda(warmup_steps, max_steps))
    tracker = Tracker(out_dir / "train_log.jsonl")
    ck_path = out_dir / "last.pt"; best_path = out_dir / "best.pt"
    step, best, bad = 0, float("inf"), 0
    if resume and ck_path.exists():
        ck = ckpt.load(ck_path, model, opt, sched, map_location=dev)
        step, best = ck["step"], ck["best"]
        tracker.log(event="resume", step=step, best=best)
    g = torch.Generator(); g.manual_seed(seed + step)
    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate,
                              drop_last=True, generator=g, persistent_workers=num_workers > 0, prefetch_factor=4 if num_workers > 0 else None)
    val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate)
    model.train()
    t0 = time.time(); n_ex = 0; t_data = 0.0; t_last = time.time()
    done = step >= max_steps
    while not done:
        for b in train_loader:
            t_data += time.time() - t_last
            b = to_device(b, dev)
            out = model(b)
            if copula_only:
                from ..models.joint_path import pit_student_t, z_from_pit
                with torch.no_grad():
                    u = pit_student_t(b["y"].cpu().numpy(), out["mu"].detach().cpu().numpy(),
                                      out["scale"].detach().cpu().numpy(), out["nu"].detach().cpu().numpy())
                    z = torch.from_numpy(z_from_pit(u).astype(np.float32)).to(dev)
                loss = model.copula.nll(out["hidden"], z, b["y_mask"].bool())
                logs = {"copula_nll": float(loss.detach())}
            else:
                loss, logs = forecast_loss(out, b, head, model.cfg.quantiles, aux_vol_weight)
            if not torch.isfinite(loss):
                tracker.log(event="nonfinite_loss", step=step)
                opt.zero_grad(set_to_none=True); t_last = time.time()
                continue
            opt.zero_grad(set_to_none=True)
            loss.backward()
            gn = torch.nn.utils.clip_grad_norm_(params, grad_clip)
            opt.step(); sched.step(); step += 1; n_ex += batch_size
            if step % 20 == 0:
                el = time.time() - t0
                tracker.log(event="train", step=step, loss=float(loss.detach()), grad_norm=float(gn), lr=sched.get_last_lr()[0],
                            ex_per_s=n_ex / el, data_frac=t_data / el, mem_mb=mem_mb(dev), **logs)
            if step % eval_every == 0 or step >= max_steps:
                if copula_only:
                    vl = float(loss)
                else:
                    vl = evaluate_nll(model, val_loader, dev, head, model.cfg.quantiles, val_max_batches)
                improved = vl < best - 1e-5
                if improved:
                    best, bad = vl, 0
                    ckpt.save(best_path, model, None, None, step, best, {"val": vl})
                else:
                    bad += 1
                ckpt.save(ck_path, model, opt, sched, step, best, {"val": vl})
                tracker.log(event="val", step=step, val_nll=vl, best=best, bad=bad, mem_mb=mem_mb(dev),
                            elapsed_s=time.time() - t0)
                if bad >= patience:
                    tracker.log(event="early_stop", step=step)
                    done = True
            if step >= max_steps:
                done = True
            t_last = time.time()
            if done:
                break
    el = time.time() - t0
    summary = {"steps": step, "best_val": best, "wall_s": el, "ex_per_s": n_ex / max(el, 1e-9),
               "peak_mem_mb": mem_mb(dev), "n_params": model.n_params(), "device": str(dev),
               "data_time_frac": t_data / max(el, 1e-9)}
    tracker.log(event="done", **summary)
    return summary