#!/usr/bin/env python3
"""Zero-shot (NO finetune) TokenGS dl3dv_base PSNR on a few TRAIN clips:
sample 2 clips from av2 train + 2 from waymo train, eval novel-view PSNR/SSIM.
Same 448x768 / n_in=8 / s=s_max protocol as the reported warm-start number."""
import argparse, random, copy
import torch
from safetensors.torch import load_file

from tokengs import options
from mapgs.config import load_config
from mapgs.data import UnifiedClipDataset
from maptokengs import MapTokenGS

import scripts.finetune_maptokengs as F  # reuse prep / render_dyn / held_out_psnr

AV2 = "/mnt/william/data/unified/av2"
WAYMO = "/mnt/william/data/unified/waymo"
CKPT = "/mnt/william/tokengs_ckpts/checkpoints/dl3dv_base.safetensors"


def pick(ds, k, seed):
    rng = random.Random(seed)
    idx = sorted(rng.sample(range(len(ds)), k))
    sub = copy.copy(ds)
    sub.clips = [ds.clips[i] for i in idx]
    return sub, [ds.clips[i].split("/")[-1] for i in idx]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--height", type=int, default=448)
    ap.add_argument("--width", type=int, default=768)
    ap.add_argument("--n-in", type=int, default=8)
    ap.add_argument("--k", type=int, default=2)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    dev = "cuda"
    H, W = args.height, args.width

    opt = options.config_defaults["train_dl3dv_base"]; opt = opt() if callable(opt) else opt
    opt.img_size = (H, W)
    model = MapTokenGS(opt, s0=0.5, s_max=2.0, max_instances=8, dyn_per_instance=32).to(dev)
    miss, unexp = model.load_state_dict(load_file(CKPT), strict=False)
    model.eval()
    print(f"warm-started dl3dv_base (ZERO-SHOT, no finetune) | fresh {len(miss)} unexpected {len(unexp)} "
          f"| res {H}x{W} n_in {args.n_in}", flush=True)

    cfg = load_config(overrides=["data.name=unified", f"data.root={AV2},{WAYMO}",
                                  f"data.height={H}", f"data.width={W}",
                                  "model.tokens.n_map=2048", "data.max_instances=8"])

    ds_av2 = UnifiedClipDataset(cfg, roots=[AV2], split="train", n_sup_views=4)
    ds_way = UnifiedClipDataset(cfg, roots=[WAYMO], split="train", n_sup_views=4)
    sub_av2, names_av2 = pick(ds_av2, args.k, args.seed)
    sub_way, names_way = pick(ds_way, args.k, args.seed)
    print("av2 clips:  ", names_av2, flush=True)
    print("waymo clips:", names_way, flush=True)

    with torch.no_grad():
        pa, sa, da, na = F.held_out_psnr(model, sub_av2, args.k, args.n_in, 1.0, dev)
        pw, sw, dw, nw = F.held_out_psnr(model, sub_way, args.k, args.n_in, 1.0, dev)

        # combined over all clips
        ps, ss = [], []
        for sub in (sub_av2, sub_way):
            for i in range(len(sub.clips)):
                d = F.prep(sub[i], args.n_in, 1.0, dev)
                g = model.forward_reconstruction_mapped(d["mi"], d["anchors"], d["atype"], d["anormal"])
                gg = model.gaussian_group()[0]; ig = model.gaussian_instance()[0]
                model.cur_s = model.s_max
                rgb, _, _ = F.render_dyn(model, g, gg, ig, d["mi"].decoder.cam_view,
                                         d["mi"].decoder.intrinsics, d["sup_frame"], d["dyn"])
                rgb = rgb.clamp(0, 1)
                p, s = float(F.psnr(rgb, d["gt"])), float(F.ssim(rgb, d["gt"]))
                if p == p and abs(p) != float("inf"):
                    ps.append(p); ss.append(s)

    print(f"\nAV2   (n={na}): PSNR {pa:.2f}+-{da:.2f}  SSIM {sa:.3f}")
    print(f"Waymo (n={nw}): PSNR {pw:.2f}+-{dw:.2f}  SSIM {sw:.3f}")
    mp = sum(ps)/len(ps); sd = (sum((x-mp)**2 for x in ps)/len(ps))**0.5
    print(f"ALL   (n={len(ps)}): PSNR {mp:.2f}+-{sd:.2f}  SSIM {sum(ss)/len(ss):.3f}")
    print("per-clip PSNR:", [round(x,2) for x in ps])


if __name__ == "__main__":
    main()