"""Training loop: exactly ``training.max_steps`` optimizer updates, with EMA weights, gradient clipping and
validation events.

A validation event fires after global update ``step`` when ``step % val_every_updates == 0``, at ``max_steps`` and
at every step listed in ``val_extra_steps`` (0 = the untrained model). An event evaluates the training objective on
the validation pairs (``val_loss``, and ``val_loss_ema`` for the EMA weights), writes one row of ``train_log.csv``
(``event`` = 1), saves ``checkpoints/checkpoint_step<NNNNNNNN>.pt`` and the full-state ``restart_latest.pt``. With
``log_loss_every_epochs: k`` a loss-only row (``event`` = 0, no checkpoint) is written at the end of every k-th epoch.
``train_loss`` of a row is the mean over the updates since the previous row. The validation passes run under a
forked RNG with a fixed seed and restore the loader's generator state, so they never perturb the training
trajectory. ``checkpoints/checkpoint.pt`` holds the final weights; every checkpoint stores the raw and the EMA
weights and the run configuration.

EMA and LR ramp-up follow the convention of NVlabs/edm (``training/training_loop.py``): ``ema_halflife_kimg``,
``ema_rampup_ratio`` and ``lr_rampup_kimg`` come from the training section (overrides) or the method section.
"""
from __future__ import annotations

import copy
import os
import random
import time
from pathlib import Path
from typing import Any, Dict

import numpy as np
import torch

RESTART_SCHEMA = "restart_v1"
RESTART_NAME = "restart_latest.pt"
LOG_COLUMNS = "step,epoch,batch_index,event,train_loss,val_loss,val_loss_ema,train_eval_loss,train_eval_loss_ema,lr,elapsed_s"


def _ema_beta(cur_nimg: int, batch_size: int, ema_halflife_kimg: float, ema_rampup_ratio: float | None) -> float:
    ema_halflife_nimg = ema_halflife_kimg * 1000
    if ema_rampup_ratio is not None:
        ema_halflife_nimg = min(ema_halflife_nimg, cur_nimg * ema_rampup_ratio)
    return 0.5 ** (batch_size / max(ema_halflife_nimg, 1e-8))


def _set_rampup_lr(optimizer: torch.optim.Optimizer, base_lr: float, cur_nimg: int, lr_rampup_kimg: float) -> None:
    factor = min(cur_nimg / max(lr_rampup_kimg * 1000, 1e-8), 1.0)
    for g in optimizer.param_groups:
        g["lr"] = base_lr * factor


def _eval_loss(model, loader, method, device, max_batches=None) -> float:
    losses = []
    with torch.no_grad():
        for i, batch in enumerate(loader):
            if max_batches is not None and i >= max_batches:
                break
            batch = tuple(t.to(device) for t in batch)
            losses.append(method.loss(model, batch).item())
    return float(np.mean(losses)) if losses else float("nan")


def _save_checkpoint(path: Path, model, ema_model, cfg: Dict[str, Any], epoch: int, step: int) -> None:
    payload = {"model_state_dict": model.state_dict(), "epoch": epoch, "config": cfg, "step": int(step)}
    if ema_model is not None:
        payload["ema_state_dict"] = ema_model.state_dict()
    torch.save(payload, path)


def _save_restart(path: Path, payload: Dict[str, Any]) -> None:
    tmp = path.with_name(path.name + ".tmp")            # atomic: a job killed mid-save leaves no truncated file
    torch.save(payload, tmp)
    os.replace(tmp, path)


def _truncate_log(log_path: Path, keep_step: int) -> int:
    """Drop the rows after ``keep_step`` (and a partial final line)."""
    lines = log_path.read_text().splitlines()
    header, rows = lines[0], lines[1:]
    kept = []
    for line in rows:
        try:
            if int(line.split(",", 1)[0]) <= keep_step:
                kept.append(line)
        except ValueError:
            continue
    log_path.write_text("\n".join([header] + kept) + "\n")
    return len(rows) - len(kept)


class Trainer:
    def __init__(self, model, method, train_loader, val_loader, cfg: Dict[str, Any], run_dir: Path, device: str):
        self.model = model
        self.method = method
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.cfg = cfg
        self.run_dir = Path(run_dir)
        self.device = device

        tr, m = cfg["training"], cfg["method"]
        allow_tf32 = tr.get("allow_tf32")
        if allow_tf32 is not None:                    # one arithmetic class on every GPU (False = strict FP32)
            torch.backends.cuda.matmul.allow_tf32 = bool(allow_tf32)
            torch.backends.cudnn.allow_tf32 = bool(allow_tf32)
            print(f"precision: allow_tf32={bool(allow_tf32)} (matmul + cudnn)")
        if bool(tr.get("deterministic", False)):      # a resumed run is bit-identical to an uninterrupted one
            os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
            torch.backends.cudnn.deterministic = True
            torch.backends.cudnn.benchmark = False
            torch.use_deterministic_algorithms(True)
            print("determinism: use_deterministic_algorithms(True), cudnn.deterministic, benchmark off, "
                  "CUBLAS_WORKSPACE_CONFIG=" + os.environ["CUBLAS_WORKSPACE_CONFIG"])

        self.batch_size = int(tr["batch_size"])
        if "lr" in tr:
            self.base_lr = float(tr["lr"])
        elif "lr" in m:
            self.base_lr = float(m["lr"])
        else:
            raise KeyError("'lr' must be set in the training or the method section")
        self.val_seed = int(tr.get("val_seed", 777))
        self.restart_checkpointing = bool(tr.get("restart_checkpointing", True))

        if tr.get("max_steps") is None or tr.get("val_every_updates") is None:
            raise KeyError("training.max_steps and training.val_every_updates are required")
        self.max_steps = int(tr["max_steps"])
        self.val_every_updates = int(tr["val_every_updates"])
        if self.max_steps <= 0 or self.val_every_updates <= 0:
            raise ValueError("max_steps and val_every_updates must be positive")
        self.val_extra_steps = tr.get("val_extra_steps")
        if self.val_extra_steps is not None:
            if not isinstance(self.val_extra_steps, (list, tuple)) or \
                    any(isinstance(s, bool) or int(s) != s for s in self.val_extra_steps):
                raise ValueError("training.val_extra_steps must be a list of integers")
            xs = sorted({int(s) for s in self.val_extra_steps})
            bad = [s for s in xs if s < 0 or s >= self.max_steps]
            if bad:
                raise ValueError(f"training.val_extra_steps must lie in [0, max_steps): {bad}")
            self.val_extra_steps = xs
        self.log_loss_every_epochs = tr.get("log_loss_every_epochs")
        if self.log_loss_every_epochs is not None:
            k = self.log_loss_every_epochs
            if isinstance(k, bool) or int(k) != k or int(k) <= 0:
                raise ValueError("training.log_loss_every_epochs must be a positive integer")
            self.log_loss_every_epochs = int(k)

        def _opt(key, default=None):                  # training section overrides the method section
            return tr[key] if key in tr else m.get(key, default)

        self.ema_halflife_kimg = _opt("ema_halflife_kimg")           # None: no EMA
        self.ema_rampup_ratio = _opt("ema_rampup_ratio", 0.05)
        self.lr_rampup_kimg = _opt("lr_rampup_kimg")                 # None: no LR ramp
        self.max_grad_norm = _opt("max_grad_norm")                   # None: no clipping
        self.weight_decay = _opt("weight_decay")
        self.use_ema = self.ema_halflife_kimg is not None and float(self.ema_halflife_kimg) > 0
        self.use_lr_rampup = self.lr_rampup_kimg is not None and float(self.lr_rampup_kimg) > 0

    def fit(self, resume: Dict[str, Any] | None = None) -> None:
        self.run_dir.mkdir(parents=True, exist_ok=True)
        ckpt_dir = self.run_dir / "checkpoints"
        ckpt_dir.mkdir(parents=True, exist_ok=True)

        model = self.model.to(self.device)
        if self.use_ema:
            ema_model = copy.deepcopy(model).eval().to(self.device)
            for p in ema_model.parameters():
                p.requires_grad_(False)
        else:
            ema_model = None
        optimizer = self.method.configure_optimizer(model, {"lr": self.base_lr, "weight_decay": self.weight_decay})

        start_epoch = 1
        if resume is not None:
            if resume.get("schema") != RESTART_SCHEMA:
                raise AssertionError(f"restart checkpoint schema {resume.get('schema')!r} != {RESTART_SCHEMA}")
            model.load_state_dict(resume["model_state_dict"])
            if ema_model is not None:
                if resume.get("ema_state_dict") is None:
                    raise AssertionError("EMA is on but the restart checkpoint has no EMA state")
                ema_model.load_state_dict(resume["ema_state_dict"])
            optimizer.load_state_dict(resume["optimizer_state_dict"])
            start_epoch = int(resume["epoch"])        # the interrupted epoch is replayed

        n_train_pairs = len(self.train_loader.dataset)
        steps_per_epoch = len(self.train_loader)
        n_epochs = max(1, -(-self.max_steps // steps_per_epoch))
        print(f"schedule: max_steps = {self.max_steps}, validation + checkpoint every {self.val_every_updates} updates "
              f"(+ at max_steps); {n_train_pairs} training pairs, batch {self.batch_size}, {steps_per_epoch} steps/epoch -> "
              f"{n_epochs} epochs (last one partial: {self.max_steps - (n_epochs - 1) * steps_per_epoch} steps)")
        if self.use_ema:
            seen_kimg = (self.max_steps - 1) * self.batch_size / 1000.0
            eff_kimg = min(float(self.ema_halflife_kimg), float(self.ema_rampup_ratio) * seen_kimg)
            print(f"EMA: half-life {self.ema_halflife_kimg} kimg, ramp-up ratio {self.ema_rampup_ratio} "
                  f"(effective half-life at the last update ~= {eff_kimg:.1f} kimg)")
        else:
            print("EMA: disabled")
        print(f"LR ramp-up: {self.lr_rampup_kimg if self.use_lr_rampup else 'disabled'}; "
              f"gradient clipping: {self.max_grad_norm if self.max_grad_norm is not None else 'disabled'}")

        log_path = self.run_dir / "train_log.csv"
        if resume is not None:
            if int(resume["step"]) >= self.max_steps:
                raise AssertionError(f"resume step {int(resume['step'])} >= max_steps {self.max_steps} (nothing to continue)")
            if not log_path.exists():
                raise FileNotFoundError(f"--resume requires the existing {log_path}")
            dropped = _truncate_log(log_path, int(resume["step"]))
            print(f"resume: continuing at step {int(resume['step'])} (epoch {start_epoch}, batch "
                  f"{int(resume['batch_index'])}; {dropped} log row(s) dropped)")
        else:
            log_path.write_text(LOG_COLUMNS + "\n")

        cur_nimg = 0
        if resume is not None:                        # RNG streams restored last, after every setup step
            cur_nimg = int(resume["cur_nimg"])
            torch.set_rng_state(resume["torch_rng_state"])
            if str(self.device).startswith("cuda"):
                torch.cuda.set_rng_state_all(resume["cuda_rng_state_all"])
            np.random.set_state(resume["numpy_rng_state"])
            random.setstate(resume["python_rng_state"])
            gen = self.train_loader.generator
            if resume.get("loader_generator_state_epoch_start") is not None:
                if gen is None:
                    raise AssertionError("restart checkpoint has a loader generator state but the train loader has no generator")
                gen.set_state(resume["loader_generator_state_epoch_start"])
            print(f"resume: restored optimizer, EMA, counters (cur_nimg={cur_nimg}), RNG streams and loader generator", flush=True)
        self._fit_updates(model=model, ema_model=ema_model, optimizer=optimizer, ckpt_dir=ckpt_dir, log_path=log_path,
                          resume=resume, start_epoch=start_epoch, cur_nimg=cur_nimg, n_val_batches=len(self.val_loader))

    def _fit_updates(self, *, model, ema_model, optimizer, ckpt_dir, log_path, resume, start_epoch, cur_nimg, n_val_batches):
        max_steps, every = self.max_steps, self.val_every_updates
        extra_steps = set(self.val_extra_steps or [])
        loss_every = int(self.log_loss_every_epochs) if self.log_loss_every_epochs else 0
        devices_all = [self.device] if str(self.device).startswith("cuda") else []
        losses_since_epoch_start: list[float] = []
        last_event_step = -1
        t_row = time.time()
        if resume is not None:
            losses_since_epoch_start = [float(x) for x in resume.get("losses_since_epoch_start", [])]
            last_event_step = int(resume["step"])
        step, start_batch = 0, 0
        if resume is not None:
            step, start_batch = int(resume["step"]), int(resume["batch_index"])
        epoch = start_epoch
        losses_since_event: list[float] = []
        t_event = time.time()
        done = False
        gen = self.train_loader.generator

        def _event(epoch_i: int, batch_i: int, gen_state_epoch_start):
            nonlocal losses_since_event, t_event, last_event_step, t_row
            train_loss = float(np.mean(losses_since_event)) if losses_since_event else float("nan")
            cur_lr = optimizer.param_groups[0]["lr"]
            model.eval()
            g_state = gen.get_state() if gen is not None else None
            with torch.random.fork_rng(devices=devices_all):
                torch.manual_seed(self.val_seed)
                val_loss = _eval_loss(model, self.val_loader, self.method, self.device)
                train_eval = _eval_loss(model, self.train_loader, self.method, self.device, max_batches=n_val_batches)
                if self.use_ema:
                    val_loss_ema = _eval_loss(ema_model, self.val_loader, self.method, self.device)
                    train_eval_ema = _eval_loss(ema_model, self.train_loader, self.method, self.device, max_batches=n_val_batches)
                else:
                    val_loss_ema = train_eval_ema = float("nan")
            if g_state is not None:
                gen.set_state(g_state)
            elapsed = time.time() - t_event
            t_event = t_row = time.time()
            last_event_step = step
            print(f"step {step} (epoch {epoch_i}, batch {batch_i}): train_loss={train_loss:.6f}  val={val_loss:.6f} "
                  f"val_ema={val_loss_ema:.6f}  lr={cur_lr:.2e}  elapsed={elapsed:.2f}s", flush=True)
            with open(log_path, "a") as f:
                f.write(f"{step},{epoch_i},{batch_i},1,{train_loss},{val_loss},{val_loss_ema},"
                        f"{train_eval},{train_eval_ema},{cur_lr},{elapsed}\n")
            losses_since_event = []
            ckpt = ckpt_dir / f"checkpoint_step{step:08d}.pt"
            _save_checkpoint(ckpt, model, ema_model, self.cfg, epoch_i, step=step)
            print(f"  saved: {ckpt.name}", flush=True)
            if self.restart_checkpointing:
                _save_restart(ckpt_dir / RESTART_NAME, {
                    "schema": RESTART_SCHEMA, "epoch": epoch_i, "step": step, "batch_index": batch_i, "cur_nimg": cur_nimg,
                    "model_state_dict": model.state_dict(),
                    "ema_state_dict": (ema_model.state_dict() if ema_model is not None else None),
                    "optimizer_state_dict": optimizer.state_dict(),
                    "torch_rng_state": torch.get_rng_state(),
                    "cuda_rng_state_all": (torch.cuda.get_rng_state_all() if str(self.device).startswith("cuda") else []),
                    "numpy_rng_state": np.random.get_state(), "python_rng_state": random.getstate(),
                    "loader_generator_state_epoch_start": gen_state_epoch_start,
                    "losses_since_epoch_start": list(losses_since_epoch_start),
                    "config": self.cfg})
            model.train()

        def _loss_row(epoch_i: int, batch_i: int):
            nonlocal losses_since_epoch_start, t_row
            if not losses_since_epoch_start or step == last_event_step:
                losses_since_epoch_start = []
                return                                # nothing to log, or an event just logged this step
            train_loss = float(np.mean(losses_since_epoch_start))
            cur_lr = optimizer.param_groups[0]["lr"]
            model.eval()
            g_state = gen.get_state() if gen is not None else None
            with torch.random.fork_rng(devices=devices_all):
                torch.manual_seed(self.val_seed)
                val_loss = _eval_loss(model, self.val_loader, self.method, self.device)
                val_loss_ema = (_eval_loss(ema_model, self.val_loader, self.method, self.device) if self.use_ema else float("nan"))
            if g_state is not None:
                gen.set_state(g_state)
            model.train()
            nan = float("nan")
            elapsed = time.time() - t_row
            t_row = time.time()
            with open(log_path, "a") as f:
                f.write(f"{step},{epoch_i},{batch_i},0,{train_loss},{val_loss},{val_loss_ema},{nan},{nan},{cur_lr},{elapsed}\n")
            losses_since_epoch_start = []

        if 0 in extra_steps and resume is None and step == 0:
            # the untrained model: its train_loss is a forward-only loss on the first batch; RNG and loader untouched
            g_state = gen.get_state() if gen is not None else None
            model.eval()
            with torch.random.fork_rng(devices=devices_all):
                torch.manual_seed(self.val_seed)
                with torch.no_grad():
                    b0 = tuple(t.to(self.device) for t in next(iter(self.train_loader)))
                    loss0 = float(self.method.loss(model, b0).item())
            model.train()
            if g_state is not None:
                gen.set_state(g_state)
            losses_since_event = [loss0]
            _event(epoch, 0, gen.get_state() if gen is not None else None)

        while not done:
            gen_state_epoch_start = gen.get_state() if gen is not None else None
            replay = resume is not None and epoch == start_epoch and start_batch > 0
            if replay:                                # the interrupted epoch: same permutation, consumed batches skipped
                gen_state_epoch_start = resume["loader_generator_state_epoch_start"]
            model.train()
            if not replay:
                losses_since_epoch_start = []
            it = iter(self.train_loader)
            skip = start_batch if (resume is not None and epoch == start_epoch) else 0
            for _ in range(skip):
                next(it)
            batch_i = skip
            for batch in it:
                batch = tuple(t.to(self.device) for t in batch)
                if self.use_lr_rampup:
                    _set_rampup_lr(optimizer, self.base_lr, cur_nimg, float(self.lr_rampup_kimg))
                optimizer.zero_grad()
                loss = self.method.loss(model, batch)
                loss.backward()
                if self.max_grad_norm is not None:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), float(self.max_grad_norm))
                optimizer.step()
                if self.use_ema:
                    beta = _ema_beta(cur_nimg, self.batch_size, float(self.ema_halflife_kimg), self.ema_rampup_ratio)
                    for p_ema, p_net in zip(ema_model.parameters(), model.parameters()):
                        p_ema.copy_(p_net.detach().lerp(p_ema, beta))
                    for b_ema, b_net in zip(ema_model.buffers(), model.buffers()):
                        b_ema.copy_(b_net)
                cur_nimg += batch[0].shape[0]
                _li = loss.item()
                losses_since_event.append(_li)
                losses_since_epoch_start.append(_li)
                step += 1
                batch_i += 1
                if step % every == 0 or step == max_steps or step in extra_steps:
                    _event(epoch, batch_i, gen_state_epoch_start)
                if step >= max_steps:
                    done = True
                    break
            if not done and loss_every and epoch % loss_every == 0:
                _loss_row(epoch, batch_i)
            epoch += 1

        print(f"trained exactly {max_steps} updates", flush=True)
        _save_checkpoint(ckpt_dir / "checkpoint.pt", model, ema_model, self.cfg, epoch - 1, step=max_steps)