#!/usr/bin/env python3
"""Train MapVGGT FULL -- VGGT-Omega backbone (frozen) + per-pixel static Gaussians
+ MAGT map-anchored Gaussian tokens (contrib ①) + PointForward scene-graph dynamics,
with MapGS map-aware losses (map-depth ②, sub-surface free-space ②) and the s(t)
tempering curriculum. Train on Waymo; hold out ONE (pinned) clip for val."""
import argparse, time, os, random, copy
import numpy as np
import torch
import torch.nn.functional as F
from gsplat import rasterization

from mapgs.config import load_config
from mapgs.data import UnifiedClipDataset
from mapgs.hdmap.rasterize_map import rasterize_map_depth
from mapgs.losses import mapdepth_loss, Tempering
from mapgs.eval.metrics import psnr, ssim
from mapvggt import MapVGGT
from mapvggt.heads import cat_gaussians

DEV = "cuda"


def prep(s, n_in, device):
    I = min(s.box_centers.shape[0], 8)
    dyn = None
    if I > 0:
        dyn = dict(c=s.box_centers[:I].to(device), R=s.box_rots[:I].to(device),
                   valid=s.box_valid[:I].to(device), canon=s.canon_idx[:I].to(device),
                   radius=(0.5 * s.box_size[:I].max(-1).values).to(device).clamp(0.1, 5.0))
    return dict(
        in_img=s.ctx_images[:n_in].to(device), in_K=s.ctx_K[:n_in].to(device),
        in_c2w=s.ctx_c2w[:n_in].to(device),
        sup_img=s.sup_images.to(device), sup_K=s.sup_K.to(device), sup_c2w=s.sup_c2w.to(device),
        sup_frame=s.sup_frame.to(device), ground=s.ground.to(device),
        ap=s.anchor_pos[None].to(device), at=s.anchor_type[None].to(device),
        an=s.anchor_normal[None].to(device), dyn=dyn)


def render_scene(model, gsm, d, H, W, gain):
    """gsm = static+map gaussians (frame-independent). Place dynamics per sup-view frame,
    union, render each view. Returns rgb [S,3,H,W], depth [S,H,W]."""
    canon = model.dyn_canonical() if (model.with_dyn and d["dyn"] is not None) else None
    rgbs, deps = [], []
    for i in range(d["sup_c2w"].shape[0]):
        full = gsm
        if canon is not None:
            dn = d["dyn"]
            dg = model.dyn_head.place(canon, int(d["sup_frame"][i]), dn["c"], dn["R"],
                                      dn["canon"], dn["valid"], dn["radius"], gain=gain)
            if dg["means"].shape[0] > 0:
                full = cat_gaussians(gsm, dg)
        out, _, _ = rasterization(means=full["means"], quats=full["quats"], scales=full["scales"],
                                  opacities=full["opacities"], colors=full["colors"],
                                  viewmats=torch.inverse(d["sup_c2w"][i:i + 1]), Ks=d["sup_K"][i:i + 1],
                                  width=W, height=H, near_plane=0.01, far_plane=500.0,
                                  render_mode="RGB+ED")
        rgbs.append(out[0, ..., :3].clamp(0, 1).permute(2, 0, 1)); deps.append(out[0, ..., 3])
    return torch.stack(rgbs), torch.stack(deps)


@torch.no_grad()
def evaluate(model, vds, n_in, device):
    model.eval(); model.cur_s = model.s_max
    ps, ss = [], []
    for i in range(len(vds.clips)):
        d = prep(vds[i], n_in, device)
        gsm = model(d["in_img"], d["in_K"], d["in_c2w"], d["ap"], d["at"], d["an"])
        rgb, _ = render_scene(model, gsm, d, *d["sup_img"].shape[-2:], gain=1.0)
        p = float(psnr(rgb, d["sup_img"]))
        if p == p and abs(p) != float("inf"):
            ps.append(p); ss.append(float(ssim(rgb, d["sup_img"])))
    model.train()
    if not model.finetune_backbone:
        model.vggt.eval()
    mp = sum(ps) / max(1, len(ps)); sd = (sum((x - mp) ** 2 for x in ps) / max(1, len(ps))) ** 0.5
    return mp, sum(ss) / max(1, len(ss)), sd


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--roots", default="/mnt/william/data/unified/waymo")
    ap.add_argument("--iters", type=int, default=3000)
    ap.add_argument("--n-in", type=int, default=8)
    ap.add_argument("--height", type=int, default=256)
    ap.add_argument("--width", type=int, default=448)
    ap.add_argument("--lr", type=float, default=3e-4)        # static per-pixel head
    ap.add_argument("--lr-ext", type=float, default=1e-4)    # map/dyn heads (gentler)
    ap.add_argument("--lam-md", type=float, default=0.5)
    ap.add_argument("--lam-fs", type=float, default=0.1)
    ap.add_argument("--lam-sparse", type=float, default=0.02)  # opacity sparsity: keep extras OFF unless useful
    ap.add_argument("--dyn-ramp", type=int, default=800)
    ap.add_argument("--val-segs", type=int, default=40)      # held-out SEGMENTS (scene-disjoint)
    ap.add_argument("--no-map", action="store_true")         # ablation: drop MAGT map tokens
    ap.add_argument("--no-dyn", action="store_true")         # ablation: drop scene-graph dynamics
    ap.add_argument("--finetune-backbone", action="store_true")
    ap.add_argument("--lr-vggt", type=float, default=1e-5)
    ap.add_argument("--eval-every", type=int, default=250)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default="/mnt/william/runs/mapvggt_full.safetensors")
    args = ap.parse_args()
    H, W = args.height, args.width
    random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed)

    model = MapVGGT(with_map=not args.no_map, with_dyn=not args.no_dyn,
                    finetune_backbone=args.finetune_backbone).to(DEV)
    cfg = load_config(overrides=["data.name=unified", f"data.root={args.roots}",
                                 f"data.height={H}", f"data.width={W}", "model.tokens.n_map=2048"])
    full = UnifiedClipDataset(cfg, roots=args.roots.split(","), split="train", n_sup_views=6)
    # SCENE-disjoint split: hold out whole SEGMENTS (a Waymo segment -> ~9 clips that share
    # the same physical scene). Splitting by clip index leaks scenes; split by segment id.
    def segid(p):
        return "_".join(os.path.basename(p.rstrip("/")).split("_")[:2])     # train_<segid>
    segs = sorted(set(segid(c) for c in full.clips))
    val_segs = set(segs[:args.val_segs])
    ds = copy.copy(full);  ds.clips = [c for c in full.clips if segid(c) not in val_segs]
    # val = ONE clip per held-out segment (scene-diverse, fast eval), all scene-disjoint from train
    seen = set(); vclips = []
    for c in full.clips:
        sgi = segid(c)
        if sgi in val_segs and sgi not in seen:
            seen.add(sgi); vclips.append(c)
    vds = copy.copy(full); vds.clips = vclips
    leak = sum(segid(c) in val_segs for c in ds.clips)
    print(f"MapVGGT (VGGT-Omega) map={not args.no_map} dyn={not args.no_dyn} | train {len(ds.clips)} clips "
          f"({len(segs)-len(val_segs)} segs) | val {len(vds.clips)} scene-disjoint segs | "
          f"train/val segment overlap {leak} | {H}x{W} n_in {args.n_in} ft={args.finetune_backbone}", flush=True)

    temper = Tempering(cfg.loss, cfg.model.tokens, args.iters)
    ext_params = []
    if model.with_map: ext_params += list(model.map_head.parameters())
    if model.with_dyn: ext_params += list(model.dyn_head.parameters())
    ext_ids = {id(p) for p in ext_params}
    vggt_ids = {id(p) for p in model.vggt.parameters()}
    groups = [
        {"params": [p for p in model.parameters() if p.requires_grad and id(p) not in ext_ids | vggt_ids], "lr": args.lr},
        {"params": [p for p in model.parameters() if p.requires_grad and id(p) in ext_ids], "lr": args.lr_ext},
    ]
    if args.finetune_backbone:
        groups.append({"params": [p for p in model.vggt.parameters() if p.requires_grad], "lr": args.lr_vggt})
    opt = torch.optim.AdamW(groups, betas=(0.9, 0.95))
    from safetensors.torch import save_file
    def tsd():
        if args.finetune_backbone:
            return model.state_dict()
        return {k: v for k, v in model.state_dict().items() if not k.startswith("vggt.")}

    b_ps, b_ss, b_sd = evaluate(model, vds, args.n_in, DEV)
    print(f"BEFORE (heads untrained): val PSNR {b_ps:.2f}±{b_sd:.2f} SSIM {b_ss:.3f} (n={len(vds.clips)})", flush=True)
    best_path = args.out.replace(".safetensors", "_best.safetensors"); best = b_ps
    t = time.time()
    for step in range(args.iters):
        model.cur_s = temper.s(step)
        gain = min(1.0, step / max(1, args.dyn_ramp))
        d = prep(ds[step % len(ds.clips)], args.n_in, DEV)
        gsm = model(d["in_img"], d["in_K"], d["in_c2w"], d["ap"], d["at"], d["an"])
        rgb, depth = render_scene(model, gsm, d, H, W, gain)
        l_rgb = F.l1_loss(rgb, d["sup_img"]) + 0.1 * (1 - ssim(rgb, d["sup_img"]))
        with torch.no_grad():
            md, mask = rasterize_map_depth(d["ground"], d["sup_K"], d["sup_c2w"], H, W)
        eps = temper.eps(step)
        l_md = mapdepth_loss(depth, md, mask, eps, cfg.loss.huber_delta) if mask.any() else depth.sum() * 0
        l_fs = F.relu(depth - md)[mask].mean() if mask.any() else depth.sum() * 0
        l_sp = model._map_opacity.mean() if model._map_opacity is not None else rgb.sum() * 0
        loss = l_rgb + args.lam_md * l_md + args.lam_fs * l_fs + args.lam_sparse * l_sp
        opt.zero_grad(set_to_none=True)
        if torch.isfinite(loss):
            loss.backward()
            gn = torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
            if torch.isfinite(gn):
                opt.step()
        if step % 50 == 0 or step < 4:
            print(f"it {step:5d} | loss {float(loss):.4f} rgb {float(l_rgb):.4f} md {float(l_md):.4f} "
                  f"fs {float(l_fs):.4f} | s {model.cur_s:.2f} gain {gain:.2f} G {gsm['means'].shape[0]//1000}k "
                  f"| {time.time()-t:.0f}s", flush=True)
        if step > 0 and step % args.eval_every == 0:
            e_ps, e_ss, e_sd = evaluate(model, vds, args.n_in, DEV)
            tag = ""
            if e_ps > best:
                best = e_ps; save_file(tsd(), best_path); tag = " *best"
            save_file(tsd(), args.out)
            print(f"  [eval @ {step}] val PSNR {e_ps:.2f}±{e_sd:.2f} SSIM {e_ss:.3f}{tag} | {time.time()-t:.0f}s", flush=True)

    a_ps, a_ss, a_sd = evaluate(model, vds, args.n_in, DEV)
    if a_ps > best:
        best = a_ps; save_file(tsd(), best_path)
    save_file(tsd(), args.out)
    print(f"\nAFTER ({args.iters} it): val PSNR {a_ps:.2f} SSIM {a_ss:.3f} | BEFORE {b_ps:.2f} | BEST {best:.2f}", flush=True)


if __name__ == "__main__":
    main()