#!/usr/bin/env python3
"""部署加速引擎（ref/gc25d-live-demo 的 FastStream：静态 KV 环形缓冲 + torch.compile + CUDA graph）在 480p / 704p 下的
每格耗时。耗时与权重数值无关，所以直接用包里的 v8 DMD 权重（ids 版 GcDiT）：只换 lat_hw / 滑窗长度 / 去噪步数。
稳态 = 缓存已满（n = cap）：每格 = 历史 K 旋转 1 次 + N 趟去噪 + 1 趟提交，另测 TAEHV 解 1 格（4 帧）。
不含：语义条件嵌入（每格 1 次，很小）、文本 K/V（按 prompt 缓存，换按键组合那格多 6–7 ms）、代码世界、网络。
只 import 包里的代码（PYTHONPATH 只放包，见 usage），不碰 maliang 的 gc25d。
usage:
  R=ref/gc25d-live-demo; CUDA_VISIBLE_DEVICES=7 PYTHONPATH=$R/code/gc25d:$R/code/arope:$R/code/diffsynth_src AROPE_ROOT=$R/code/arope \
  TAEHV_DIR=$R/code/taehv DIFFSYNTH_SKIP_DOWNLOAD=True TMPDIR=.tmp TORCHINDUCTOR_CACHE_DIR=.tmp/inductor TRITON_CACHE_DIR=.tmp/triton \
  .venv/bin/python scripts/bench_fast704.py --out outputs/bench_fast/run.json"""
import argparse
import json
import os
import statistics
import subprocess
import time

import torch

from codeworld.fast_live import FastStream
from codeworld.live import TinyStreamVAE, empty_wan_dit
from gcdistill.gc_dit import GcDiT, load_gc_ckpt

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
R = f"{M}/ref/gc25d-live-demo"
RES = {"480p": (30, 52), "704p": (44, 80)}           # 832×480 / 1280×704 的 latent 网格（VAE ×16）
T_GRID = (1000.0, 750.0, 500.0, 250.0)


def gpu_apps():
    """本卡上的计算进程（判断测的时候有没有别的进程在抢卡）。"""
    try:
        idx = os.environ.get("CUDA_VISIBLE_DEVICES", "0").split(",")[0]
        q = subprocess.run(["nvidia-smi", "-i", idx, "--query-compute-apps=pid,used_memory", "--format=csv,noheader"],
                           capture_output=True, text=True, timeout=30).stdout.strip()
        return [ln for ln in q.splitlines() if ln and str(os.getpid()) not in ln]
    except Exception as e:  # noqa: BLE001
        return [f"nvidia-smi failed: {e}"]


def cell_ms(fs, x, steps, n_cells):
    out = []
    for _ in range(n_cells):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        fs._rotate()
        for t in T_GRID[:steps]:
            fs.forward(x, t, False)
        fs.forward(x, 0.0, True)
        torch.cuda.synchronize()
        out.append((time.perf_counter() - t0) * 1e3)
    return out


def summ(v):
    v = sorted(v)
    return {"median": round(statistics.median(v), 1), "p90": round(v[int(0.9 * (len(v) - 1))], 1), "min": round(v[0], 1),
            "n": len(v)}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default=f"{R}/weights/gc25d_v8_dmd_ema500.safetensors")
    ap.add_argument("--tae", default=f"{R}/weights/taew2_2_super.pth")
    ap.add_argument("--configs", default="480p:17:4,704p:17:4,704p:17:3,704p:11:4,704p:11:3",
                    help="res:window:steps，逗号分隔；sink 固定 3")
    ap.add_argument("--cells", type=int, default=30)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    dev = "cuda"
    rep = {"gpu": torch.cuda.get_device_name(0), "others_on_gpu_before": gpu_apps(), "configs": {}, "taehv": {}}
    print(rep, flush=True)
    model = GcDiT(dit=empty_wan_dit(), causal=True).to(dev).eval().requires_grad_(False)
    model.move_rope_table(dev)
    load_gc_ckpt(model, a.ckpt)
    model.to(torch.bfloat16)
    g = torch.Generator(device=dev).manual_seed(0)
    fs_cache = {}
    for spec in a.configs.split(","):
        res, win, steps = spec.split(":")
        win, steps = int(win), int(steps)
        key = (res, win)
        if key not in fs_cache:
            fs_cache.clear()
            torch.cuda.empty_cache()
            t0 = time.time()
            fs = FastStream(model, kv_window=win, sink_size=3, rope_cap=3 + win, lat_hw=RES[res], graphs=True, compile=True)
            fs.ktxt.normal_(generator=g).mul_(0.5)
            fs.vtxt.normal_(generator=g).mul_(0.5)
            fs.order, fs.n = list(range(fs.cap)), fs.cap                     # 稳态：缓存已满（之后每格提交驱逐最老的非 sink 格）
            fs.pos.copy_(torch.arange(fs.cap, device=dev))
            fs.kraw.normal_(generator=g).mul_(0.5)
            fs.vbuf.normal_(generator=g).mul_(0.5)
            x = torch.randn(1, 48, 1, *RES[res], generator=g, device=dev, dtype=torch.bfloat16)
            cell_ms(fs, x, 4, 3)                                              # 编译 + 录 graph（(cap, 不提交) / (cap, 提交)）
            fs_cache[key] = (fs, x)
            print(f"[{res} w{win}] build+compile+capture {time.time() - t0:.1f}s  bufs {fs.nbytes() / 2**30:.2f} GiB", flush=True)
        fs, x = fs_cache[key]
        cell_ms(fs, x, steps, 3)
        v = cell_ms(fs, x, steps, a.cells)
        rep["configs"][spec] = {**summ(v), "tokens_per_cell": fs.hw, "kv_cells": fs.cap,
                                "mem_reserved_GiB": round(torch.cuda.memory_reserved() / 2**30, 1)}
        print(spec, rep["configs"][spec], flush=True)
    fs_cache.clear()
    torch.cuda.empty_cache()
    vae = TinyStreamVAE(a.tae, dev)
    for res, hw in RES.items():
        vae.reset()
        z = torch.randn(1, 48, 1, *hw, generator=g, device=dev, dtype=torch.bfloat16)
        for _ in range(4):
            vae.decode_one(z)
        v = []
        for _ in range(a.cells):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            fr = vae.decode_one(z)
            torch.cuda.synchronize()
            v.append((time.perf_counter() - t0) * 1e3)
        rep["taehv"][res] = {**summ(v), "frames": list(fr.shape)}
        print("taehv", res, rep["taehv"][res], flush=True)
    rep["others_on_gpu_after"] = gpu_apps()
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    json.dump(rep, open(a.out, "w"), indent=1)
    print(json.dumps(rep, indent=1))


if __name__ == "__main__":
    main()