#!/usr/bin/env python3
"""分辨率决策的两个快测：(1) 角色区域的 VAE 重建上限：832×480 / 960×544 / 1280×704（都在源分辨率上和 GT 比）；
(2) DiT 前向耗时：21 cell 双向前向在三种分辨率下的相对开销（实时可行性的粗估）。"""
import json
import os
import sys
import time

import cv2
import numpy as np
import torch

import geometry as GEO
import infer as INF
import model as GM
import sample as S

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
D = f"{M}/outputs/char_diag"
RES = {"832x480": (832, 480), "960x544": (960, 544), "1024x576": (1024, 576), "1152x640": (1152, 640), "1280x704": (1280, 704)}
NF = 41


def enc(vae, frames):
    """infer.encode_video 的任意分辨率版（fp32 编码，非 tiled）。"""
    x = torch.from_numpy(frames).permute(3, 0, 1, 2).float().div_(127.5).sub_(1.0)
    with torch.no_grad(), INF._vae_fp32(vae, True):
        z = vae.vae.encode([x], device=vae.device, tiled=False)
    return z.to(dtype=torch.bfloat16, device=vae.device)


def dec(vae, z):
    with torch.no_grad():
        v = vae.vae.decode(z.to(vae.torch_dtype), device=vae.device, tiled=False)
    return v[0].float().add_(1.0).mul_(127.5).round_().clamp_(0, 255).to(torch.uint8).permute(1, 2, 3, 0).cpu().numpy()


def to_res(src, wh):
    """源 1280×704 → 目标分辨率：832/960 先按 gc25d 约定裁成目标宽高比再 INTER_AREA；1280×704 原样。"""
    W, H = wh
    if (W, H) == (1280, 704):
        return src
    cw = int(round(704 * W / H))
    x0 = (1280 - cw) // 2
    return np.stack([cv2.resize(f[:, x0:x0 + cw], (W, H), interpolation=cv2.INTER_AREA) for f in src]), x0, cw


def main():
    clips = json.load(open(f"{D}/clips.json"))
    vae = INF.load_vae("cuda", S.MODEL_DIR)
    out = {k: [] for k in RES}
    for cid in clips:
        ref = np.load(f"{D}/_ref/{cid}.npz")
        src, ids = ref["gt_src"][:NF], ref["ids"]
        for name, wh in RES.items():
            if name == "1280x704":
                x, x0, cw = src, 0, 1280
            else:
                x, x0, cw = to_res(src, wh)
            rec = dec(vae, enc(vae, np.ascontiguousarray(x)))
            for t in range(8, NF, 8):
                m = np.isin(ids[t], GM.PLAYER_PART_IDS)
                if m.sum() < 600:
                    continue
                ys, xs = np.nonzero(m)
                p = GEO.dst_to_src(np.array([[xs.min() - 10, ys.min() - 10], [xs.max() + 10, ys.max() + 10]], float))
                sx0, sy0 = int(max(0, p[0, 0])), int(max(0, p[0, 1]))
                sx1, sy1 = int(min(1280, p[1, 0])), int(min(704, p[1, 1]))
                gt_c = src[t, sy0:sy1, sx0:sx1]
                # 重建帧映回源分辨率坐标再裁同一块（832/960 的重建上采样回源尺寸，比较同一份 GT）
                if name == "1280x704":
                    rc = rec[t, sy0:sy1, sx0:sx1]
                else:
                    up = cv2.resize(rec[t], (cw, 704), interpolation=cv2.INTER_CUBIC)
                    full = np.zeros_like(src[t]); full[:, x0:x0 + cw] = up
                    rc = full[sy0:sy1, sx0:sx1]
                mse = np.mean((rc.astype(np.float32) - gt_c.astype(np.float32)) ** 2)
                g = lambda a: float(np.abs(cv2.Laplacian(cv2.cvtColor(a, cv2.COLOR_RGB2GRAY).astype(np.float32), cv2.CV_32F)).mean())
                out[name].append((10 * np.log10(255 ** 2 / max(mse, 1e-6)), g(rc) / max(g(gt_c), 1e-6)))
        print(cid, {k: len(v) for k, v in out.items()}, flush=True)
    res = {k: {"n": len(v), "player_crop_psnr_vs_native_gt": round(float(np.mean([a for a, _ in v])), 2),
               "detail_ratio_vs_native_gt": round(float(np.mean([b for _, b in v])), 3)} for k, v in out.items()}
    print(json.dumps(res, indent=1), flush=True)
    del vae
    torch.cuda.empty_cache()

    # (2) DiT 前向耗时（bf16，no_grad，21 cell，双向；只看相对比例）
    dit = GM.load_dit(S.MODEL_DIR, "cuda")
    GM.install_sem_cond(dit, "parts")
    dit = dit.cuda().eval().requires_grad_(False)
    dit.freqs = tuple(t.cuda() for t in dit.freqs)
    tt = S.TextTable(os.path.join(S.DATA_ROOT, "text_table_gc25d.pt"))
    ctx, cids = tt.context([sorted(tt.table)[0]] * 21, "cuda")
    timing = {}
    for name, (W, H) in RES.items():
        lat = torch.randn(1, 48, 21, H // 16, W // 16, dtype=torch.bfloat16, device="cuda")
        sem = torch.ones(1, 81, H, W, dtype=torch.uint8, device="cuda")
        ts = []
        for i in range(4):
            torch.cuda.synchronize(); t0 = time.time()
            with torch.no_grad():
                GM.gc_forward(dit, lat, torch.tensor([500.0], dtype=torch.bfloat16, device="cuda"), ctx, sem,
                              first_frame_cond=True, context_ids=cids)
            torch.cuda.synchronize(); ts.append(time.time() - t0)
        timing[name] = {"tokens_per_cell": (H // 32) * (W // 32), "sec_per_forward": round(float(np.median(ts[1:])), 3)}
        print(name, timing[name], flush=True)
    base = timing["832x480"]["sec_per_forward"]
    for k in timing:
        timing[k]["ratio_vs_480"] = round(timing[k]["sec_per_forward"] / base, 2)
    json.dump({"vae_ceiling": res, "dit_timing": timing}, open(f"{D}/res_test.json", "w"), indent=1)
    print(json.dumps(timing, indent=1))


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