"""Training loop + the per-step loss assembly (§3).

``compute_step_losses`` runs one feed-forward pass and assembles every term of
the objective (§2.6); the ``Trainer`` wraps it with the AdamW/cosine schedule,
the tempering curriculum (§2.7), the L_extrap ramp-in, grad clipping, logging
and checkpointing.
"""

from __future__ import annotations

import os
import time
from typing import Optional

import torch
from torch.utils.data import DataLoader

from mapgs.config import MapGSConfig
from mapgs.data import collate_samples, build_dataset
from mapgs.hdmap.rasterize_map import rasterize_map_depth, project_polylines, render_lane_mask
from mapgs.losses import (
    MapGSCriterion, mapdepth_loss, free_space_loss, ground_coupling_loss,
    silog_loss, non_ground_static_mask, visibility_loss,
    perturb_pose, compute_warp_targets, warp_loss, lane_loss,
)
from mapgs.model import MapGS
from mapgs.model.dynamic import place_dynamic_gaussians
from mapgs.render.gaussians import GROUP_FREE
from mapgs.render.rasterizer import GaussianRasterizer
from mapgs.train.optim import build_optimizer, cosine_warmup
from mapgs.train.render import batched_plucker, render_scene_views, extract_scene_dynamic


def _move_batch(batch: dict, device):
    for k, v in batch.items():
        if torch.is_tensor(v):
            batch[k] = v.to(device)
    if batch.get("dynamic") is not None:
        batch["dynamic"] = {k: v.to(device) for k, v in batch["dynamic"].items()}
    batch["grounds"] = [g.to(device) for g in batch["grounds"]]
    batch["lanes"] = [[l.to(device) for l in ls] for ls in batch["lanes"]]
    return batch


def compute_step_losses(
    model: MapGS,
    ras: GaussianRasterizer,
    batch: dict,
    criterion: MapGSCriterion,
    it: int,
    cfg: MapGSConfig,
    device,
    use_extrap: Optional[bool] = None,
):
    H, W = cfg.data.height, cfg.data.width
    temper = criterion.tempering
    s_t, eps = temper.s(it), temper.eps(it)
    if use_extrap is None:
        use_extrap = it >= cfg.train.extrap_ramp_iter

    plucker = batched_plucker(batch["ctx_K"], batch["ctx_c2w"], H, W)
    # autocast covers only the (matmul-heavy) encoder/decoder/head; rendering and
    # losses run in fp32 so the gsplat / grid_sample boundaries stay dtype-consistent.
    amp_ctx = torch.autocast("cuda", dtype=(torch.bfloat16 if cfg.train.amp_dtype == "bf16" else torch.float16)) \
        if cfg.train.amp else _nullctx()
    with amp_ctx:
        decoded = model(
            batch["ctx_images"], plucker, batch["ctx_tids"],
            batch["anchor_pos"], batch["anchor_type"], batch["anchor_normal"],
            s_t=s_t, dynamic=batch.get("dynamic"),
        )
    decoded = decoded.to_float()
    B = batch["ctx_images"].shape[0]
    uses_feat = model.uses_features

    acc = {k: 0.0 for k in ["rgb", "mapdepth", "free_space", "ground_coupling",
                            "vert", "vis", "extrap_warp", "extrap_lane"]}
    n_ext = 0
    for b in range(B):
        g_canon = decoded.scene(b)
        dyn_b = extract_scene_dynamic(batch.get("dynamic"), b)
        ground_b = batch["grounds"][b]
        skb, scb, sfb = batch["sup_K"][b], batch["sup_c2w"][b], batch["sup_frame"][b]

        rend = render_scene_views(model, ras, g_canon, dyn_b, skb, scb, sfb, H, W, uses_feat)
        map_depth, gmask = rasterize_map_depth(ground_b, skb, scb, H, W)

        l_rgb, _ = criterion.rgb_loss(rend["rgb"], batch["sup_images"][b])
        acc["rgb"] += l_rgb
        acc["mapdepth"] += mapdepth_loss(rend["depth"], map_depth, gmask, eps, cfg.loss.huber_delta)
        acc["free_space"] += free_space_loss(g_canon, ground_b, cfg.map.ground_z_below)
        acc["ground_coupling"] += ground_coupling_loss(g_canon, ground_b, eps, cfg.loss.huber_delta)

        vmask = non_ground_static_mask(gmask, rend["alpha"])
        acc["vert"] += silog_loss(rend["depth"], batch["sup_mono"][b], vmask)

        free_sel = (decoded.group[b] == GROUP_FREE) & decoded.valid[b]
        acc["vis"] += visibility_loss(decoded.means[b][free_sel].float(), skb, scb, H, W)

        if use_extrap:
            n_ext += 1
            Kref, cref = batch["ctx_K"][b, 0], batch["ctx_c2w"][b, 0]
            md_ref, gm_ref = rasterize_map_depth(ground_b, Kref[None], cref[None], H, W)
            lat = float(torch.empty(1).uniform_(*cfg.loss.lateral_shift_range))
            yaw = float(torch.empty(1).uniform_(-cfg.loss.yaw_jitter_deg, cfg.loss.yaw_jitter_deg))
            pitch = float(torch.empty(1).uniform_(-cfg.loss.pitch_jitter_deg, cfg.loss.pitch_jitter_deg))
            Pdev = perturb_pose(cref, lat, yaw, pitch)
            uv, z, ctgt = compute_warp_targets(
                batch["ctx_images"][b, 0], md_ref[0], gm_ref[0], Kref, cref, Kref, Pdev
            )
            g_ref = g_canon if dyn_b is None else place_dynamic_gaussians(
                g_canon, dyn_b["box_centers"], dyn_b["box_rots"], dyn_b["canon_idx"], 0)
            dev_out = ras.render(g_ref, Kref[None], Pdev[None], H, W)
            dev_rgb = (model.feature_to_rgb(dev_out.color) if uses_feat else dev_out.color[:, :3].clamp(0, 1))[0].float()
            acc["extrap_warp"] += warp_loss(dev_rgb, dev_out.depth[0], uv, z, ctgt, eps, cfg.loss.warp_tau)
            if dev_out.aux is not None and len(batch["lanes"][b]) > 0:
                map_uv = project_polylines(batch["lanes"][b], Kref[None], Pdev[None], H, W)[0]
                tgt_lane = render_lane_mask(map_uv, H, W)
                acc["extrap_lane"] += lane_loss(dev_out.aux[0, 0], tgt_lane)

    parts = {}
    for k, v in acc.items():
        if k in ("extrap_warp", "extrap_lane"):
            parts[k] = (v / n_ext) if n_ext > 0 else None
        else:
            parts[k] = v / B
    return criterion.total(parts, it)


class Trainer:
    def __init__(self, cfg: MapGSConfig, model: Optional[MapGS] = None, device: Optional[str] = None):
        self.cfg = cfg
        self.device = device or cfg.device
        self.model = (model or MapGS(cfg)).to(self.device)
        if cfg.train.grad_checkpoint:
            self.model.set_grad_checkpoint(True)
        self.ras = GaussianRasterizer()
        self.criterion = MapGSCriterion(cfg, cfg.train.iters).to(self.device)
        self.opt = build_optimizer(self.model, cfg.train.lr, cfg.train.weight_decay)
        self.sched = cosine_warmup(self.opt, cfg.train.warmup, cfg.train.iters, cfg.train.min_lr_ratio)
        self.it = 0
        self.grad_accum = max(1, cfg.train.grad_accum)
        self._micro = 0
        self.amp = cfg.train.amp
        self.amp_dtype = torch.bfloat16 if cfg.train.amp_dtype == "bf16" else torch.float16
        os.makedirs(cfg.train.out_dir, exist_ok=True)

    def train_step(self, batch: dict) -> dict:
        """One micro-batch. With grad_accum>1, the optimizer steps every
        ``grad_accum`` calls (literal batch fills memory; effective batch =
        batch_size * grad_accum)."""
        self.model.train()
        batch = _move_batch(batch, self.device)
        if self._micro == 0:
            self.opt.zero_grad(set_to_none=True)
        total, log = compute_step_losses(self.model, self.ras, batch, self.criterion,
                                         self.it, self.cfg, self.device)
        finite = bool(torch.isfinite(total))
        if finite:
            (total / self.grad_accum).backward()  # skip a bad micro-batch rather than poison weights
        else:
            log["skipped"] = 1.0
        self._micro += 1
        log["lr"] = self.sched.get_last_lr()[0]
        log["grad_norm"] = 0.0
        if self._micro >= self.grad_accum:
            gn = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.train.grad_clip)
            # finite LOSS can still yield non-finite GRADIENTS (e.g. a degenerate gaussian /
            # sqrt(0) in a backward); stepping then would poison the weights permanently.
            if torch.isfinite(gn):
                self.opt.step()
            else:
                self.opt.zero_grad(set_to_none=True)
                log["skipped"] = 1.0
            self.sched.step()
            self._micro = 0
            self.it += 1
            log["grad_norm"] = float(gn) if torch.isfinite(gn) else -1.0
        return log

    def fit(self, dataset=None, max_iters: Optional[int] = None):
        if dataset is None:
            dataset = build_dataset(self.cfg, "train")
        loader = DataLoader(
            dataset, batch_size=self.cfg.train.batch_size, shuffle=True,
            num_workers=self.cfg.train.num_workers, collate_fn=collate_samples,
            drop_last=True, persistent_workers=self.cfg.train.num_workers > 0,
        )
        target = max_iters or self.cfg.train.iters
        t0 = time.time()
        while self.it < target:
            for batch in loader:
                if self.it >= target:
                    break
                log = self.train_step(batch)
                if self.it % self.cfg.train.log_every == 0:
                    dt = time.time() - t0
                    print(f"it {self.it:>7} | loss {log['total']:.4f} | rgb {log.get('rgb',0):.4f} "
                          f"| md {log.get('mapdepth',0):.4f} | lr {log['lr']:.2e} | {dt:.1f}s")
                if self.cfg.train.ckpt_every and self.it % self.cfg.train.ckpt_every == 0:
                    self.save(os.path.join(self.cfg.train.out_dir, f"ckpt_{self.it}.pt"))
        return self

    def save(self, path: str):
        torch.save({"model": self.model.state_dict(), "opt": self.opt.state_dict(),
                    "it": self.it, "cfg": self.cfg.to_dict()}, path)

    def load(self, path: str):
        ckpt = torch.load(path, map_location=self.device, weights_only=False)
        self.model.load_state_dict(ckpt["model"])
        if "opt" in ckpt:
            self.opt.load_state_dict(ckpt["opt"])
        self.it = ckpt.get("it", 0)
        return self


class _nullctx:
    def __enter__(self):
        return self

    def __exit__(self, *a):
        return False