#!/usr/bin/env python3
"""Real-data train->eval on the converted Argoverse 2 unified clips.

Trains on data/unified/av2/train (snapshot at launch) and evaluates on
data/unified/av2/val. Interpolation + lane consistency only (real logs have no
deviated GT; the synthetic-style extrapolation sweep needs a known world).
GPU-only training coexists with the CPU-bound conversion still running."""

import argparse, time
import torch
from mapgs.config import load_config
from mapgs.data import UnifiedClipDataset, collate_samples
from mapgs.eval import Evaluator
from mapgs.train import Trainer


def fmt(d): return {k: (round(v, 3) if isinstance(v, float) else v) for k, v in d.items()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--iters", type=int, default=1200)
    ap.add_argument("--root", default="/mnt/william/data/unified/av2")
    args = ap.parse_args()
    cfg = load_config("configs/base.yaml", [
        "data.name=unified", f"data.root={args.root}", "data.num_frames=20",
        "data.height=256", "data.width=384",
        "model.embed_dim=512", "model.enc_depth=3", "model.dec_depth=6", "model.n_heads=8",
        "model.tokens.gaussians_per_token=8", "model.feature_dim=32",
        "train.amp=true", "train.batch_size=1", "train.num_workers=3", "train.grad_checkpoint=true",
        "train.lr=1.0e-4", "train.warmup=100", f"train.iters={args.iters}",
        "train.extrap_ramp_iter=100000",  # off: deviated render less useful w/o lane-change GT here
        "train.log_every=50", "train.ckpt_every=0", "train.out_dir=runs/mapgs_av2",
    ])
    train_ds = UnifiedClipDataset(cfg, roots=args.root, split="train", n_sup_views=4)
    val_ds = UnifiedClipDataset(cfg, roots=args.root, split="val", n_sup_views=6)
    print(f"AV2 train clips: {len(train_ds)} | val clips: {len(val_ds)}")

    trainer = Trainer(cfg)
    ev = Evaluator(trainer.model, cfg, device="cuda")
    print("BEFORE:", fmt(ev.interpolation(val_ds, max_scenes=40)))
    t = time.time()
    trainer.fit(train_ds, max_iters=args.iters)
    print(f"trained {args.iters} iters in {time.time()-t:.0f}s")
    trainer.model.eval()
    after = ev.interpolation(val_ds, max_scenes=40)
    lane = ev.lane_consistency(val_ds, max_scenes=30, shift=3.0, frame=cfg.data.num_frames // 2)
    print("AFTER :", fmt(after))
    print("LANE  :", fmt(lane))
    trainer.save("runs/mapgs_av2/ckpt_av2.pt")


if __name__ == "__main__":
    main()