#!/usr/bin/env python3
"""角色清晰度诊断·生成端：一个变体（ckpt + cfg）在 N 条 val clip 上 50 步采样，**无损**存解码帧（不经 mp4）。
--recon 另存 VAE 重建与 GT（832×480 + 源 1280×704）。输出 <out>/<variant>/<clip>.npz（uint8 [81,480,832,3]）。"""
import argparse
import json
import os
import sys
import types

import numpy as np
import torch

import infer as INF
import sample as S

M = "/opt/dlami/nvme/zhiyangdeng/maliang"


def pick_clips(n: int) -> list[str]:
    rows = [json.loads(l) for l in open(f"{M}/data/gc25d_v10/index/clips.jsonl")]
    out = []
    for sub in ("v10_real", "v9_stylized"):
        c = [r for r in rows if r["split"] == "val" and not r["drop"] and r["subset"] == sub and r.get("dead_frac", 1) == 0
             and r.get("player_hidden_frac", 1) < 0.05 and r.get("player_px_mean", 0) > 3000]
        c.sort(key=lambda r: (r["run"], r["start"]))
        runs, pick = {}, []
        for r in c:                       # 每个 run 先取一条，凑不够再补
            if r["run"] not in runs:
                runs[r["run"]] = 1; pick.append(r)
        out += [r["clip_id"] for r in (pick + [r for r in c if r not in pick])[: n // 2]]
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", required=True)
    ap.add_argument("--variant", required=True)
    ap.add_argument("--cfg", type=float, default=1.0)
    ap.add_argument("--sem_cfg", type=float, default=1.0)
    ap.add_argument("--n", type=int, default=8)
    ap.add_argument("--recon", action="store_true")
    ap.add_argument("--out", default=f"{M}/outputs/char_diag")
    a = ap.parse_args()
    clips = pick_clips(a.n)
    os.makedirs(f"{a.out}/{a.variant}", exist_ok=True)
    json.dump(clips, open(f"{a.out}/clips.json", "w"), indent=1)
    index = S.clip_index(S.DATA_ROOT)
    args = types.SimpleNamespace(device="cuda", steps=50, shift=5.0, cfg=a.cfg, sem_cfg=a.sem_cfg, seed=0)
    dit, info = S.load_dit(a.ckpt, False, "parts", S.MODEL_DIR, "cuda")
    tt = S.TextTable(os.path.join(S.DATA_ROOT, "text_table_gc25d.pt"))
    vae = INF.load_vae("cuda", S.MODEL_DIR)
    for cid in clips:
        dst = f"{a.out}/{a.variant}/{cid}.npz"
        if os.path.exists(dst):
            continue
        c = S.load_clip(S.DATA_ROOT, cid)
        z0 = c["latent"][:, 0:1].unsqueeze(0).to("cuda")
        lat, sec = S.generate(args, dit, tt, z0, S.model_ids(c["ids"], info["vocab"], "cond"), c["prompts"])
        np.savez(dst, frames=INF.decode_video(vae, lat))
        print(a.variant, cid, f"{sec:.1f}s", flush=True)
        if a.recon:
            row = index[cid]
            os.makedirs(f"{a.out}/_ref", exist_ok=True)
            raw = INF.read_raw_frames(row["scene"], row["run"], int(row["start"]), 81, "jpg",
                                      raw_root=os.path.join(S.RAW_ROOT, row["subset"]))
            np.savez(f"{a.out}/_ref/{cid}.npz", recon=INF.decode_video(vae, c["latent"].unsqueeze(0).to("cuda")),
                     gt=S.gt_frames(S.RAW_ROOT, row), gt_src=np.stack(raw), ids=c["ids"])
    print("DONE", a.variant, flush=True)


if __name__ == "__main__":
    sys.exit(main())