#!/usr/bin/env python3
"""角色清晰度：480p（s8000）vs 704p（--variant）在原生 1280×704 坐标上比（不依赖 GC_RES：480p 的裁剪/缩放规则这里写死）。
每个 clip、每 8 帧，按 480p ids 的玩家 / Boss 掩码取框 → 映回源坐标；480p 的帧（生成 / VAE 重建）先 INTER_CUBIC 放回 1220×704
再贴到 x0=30；704p 的帧原样。指标（框内，相对原生 GT）：detail = |Laplacian| 均值之比（res_test 同款），sharp = Sobel 梯度
均值之比，psnr。另出 player_crops_480_vs_704.jpg（同尺度裁剪，×2 最近邻放大）与 compare_480_vs_704.mp4（整帧三联）。
usage: char_cmp704.py --variant s5000_704 [--v480 s8000_cfg1]"""
import argparse
import json
import os
import sys

import cv2
import imageio.v3 as iio
import numpy as np

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
D = f"{M}/outputs/char_diag"
ROLES = {"player": (10, 12, 13, 14, 15, 16), "boss": (11, 17, 18, 19, 20, 21)}
X0, CW, SW, SH, W4, H4 = 30, 1220, 1280, 704, 832, 480
PAD, MIN_PX, UP = 10, 600, 2


def up480(fr):
    """832×480 → 源 1280×704 画布（两侧 30 px 黑边）。"""
    out = np.zeros((SH, SW, 3), np.uint8)
    out[:, X0:X0 + CW] = cv2.resize(fr, (CW, SH), interpolation=cv2.INTER_CUBIC)
    return out


def src_box(mask):
    ys, xs = np.nonzero(mask)
    x0, x1 = xs.min() - PAD, xs.max() + PAD + 1
    y0, y1 = ys.min() - PAD, ys.max() + PAD + 1
    sx0, sx1 = X0 + x0 * CW / W4, X0 + x1 * CW / W4
    sy0, sy1 = y0 * SH / H4, y1 * SH / H4
    return int(max(0, sy0)), int(min(SH, np.ceil(sy1))), int(max(X0, sx0)), int(min(X0 + CW, np.ceil(sx1)))


def stats(img):
    g = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).astype(np.float32)
    gx, gy = cv2.Sobel(g, cv2.CV_32F, 1, 0), cv2.Sobel(g, cv2.CV_32F, 0, 1)
    return float(np.abs(cv2.Laplacian(g, cv2.CV_32F)).mean()), float(np.sqrt(gx ** 2 + gy ** 2).mean())


def psnr(a, b):
    return 10 * np.log10(255.0 ** 2 / max(float(np.mean((a.astype(np.float32) - b.astype(np.float32)) ** 2)), 1e-6))


def label(img, text):
    out = img.copy()
    cv2.rectangle(out, (0, 0), (out.shape[1], 18), (0, 0, 0), -1)
    cv2.putText(out, text, (3, 13), cv2.FONT_HERSHEY_SIMPLEX, 0.42, (255, 230, 120), 1, cv2.LINE_AA)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--variant", required=True)
    ap.add_argument("--v480", default="s8000_cfg1")
    ap.add_argument("--video_clips", type=int, default=3)
    ap.add_argument("--n", type=int, default=0, help="只用前 n 个 clip（测试用）")
    a = ap.parse_args()
    clips = json.load(open(f"{D}/clips.json"))
    clips = clips[:a.n] if a.n else clips
    keys = ["recon480", "gen480", "recon704", "gen704"]
    st = {r: {k: [] for k in keys} for r in ROLES}
    rows, vid = [], []
    for ci, cid in enumerate(clips):
        ref = np.load(f"{D}/_ref/{cid}.npz")
        gt, ids = ref["gt_src"], ref["ids"]
        fr = {"recon480": ref["recon"], "gen480": np.load(f"{D}/{a.v480}/{cid}.npz")["frames"],
              "recon704": np.load(f"{D}/_ref704/{cid}.npz")["recon"], "gen704": np.load(f"{D}/{a.variant}/{cid}.npz")["frames"]}
        assert fr["gen704"].shape == gt.shape, (fr["gen704"].shape, gt.shape)
        for role, pids in ROLES.items():
            for t in range(8, 81, 8):
                m = np.isin(ids[t], pids)
                if m.sum() < MIN_PX:
                    continue
                y0, y1, x0, x1 = src_box(m)
                g = gt[t, y0:y1, x0:x1]
                lg, sg = stats(g)
                crops = {}
                for k in keys:
                    full = up480(fr[k][t]) if k.endswith("480") else fr[k][t]
                    c = full[y0:y1, x0:x1]
                    crops[k] = c
                    l, s = stats(c)
                    st[role][k].append((l / max(lg, 1e-6), s / max(sg, 1e-6), psnr(c, g)))
                if role == "player" and t in (40, 72) and len(rows) < 14:
                    big = lambda x: cv2.resize(x, (x.shape[1] * UP, x.shape[0] * UP), interpolation=cv2.INTER_NEAREST)
                    names = {"recon480": "VAE recon 480p", "gen480": "gen 480p s8000", "recon704": "VAE recon 704p",
                             "gen704": f"gen 704p {a.variant}"}
                    tiles = [label(big(g), "GT native 1280x704")] + [label(big(crops[k]), names[k]) for k in keys]
                    rows.append(np.concatenate([cv2.copyMakeBorder(x, 0, 0, 0, 4, cv2.BORDER_CONSTANT) for x in tiles], 1))
        if ci < a.video_clips:
            for t in range(81):
                p = [label(gt[t], "GT native"), label(up480(fr["gen480"][t]), "gen 480p (upscaled)"),
                     label(fr["gen704"][t], f"gen 704p {a.variant}")]
                vid.append(np.concatenate(p, 1))
    summary = {r: {k: {"n": len(v), "detail_ratio": round(float(np.mean([x[0] for x in v])), 3),
                       "sharp_ratio": round(float(np.mean([x[1] for x in v])), 3),
                       "psnr": round(float(np.mean([x[2] for x in v])), 2)} for k, v in d.items() if v} for r, d in st.items()}
    json.dump(summary, open(f"{D}/summary_480_vs_704.json", "w"), indent=1)
    print(json.dumps(summary, indent=1))
    if rows:
        w = max(r.shape[1] for r in rows)
        sheet = np.concatenate([cv2.copyMakeBorder(r, 0, 6, 0, w - r.shape[1], cv2.BORDER_CONSTANT) for r in rows], 0)
        cv2.imwrite(f"{D}/player_crops_480_vs_704.jpg", sheet[:, :, ::-1], [cv2.IMWRITE_JPEG_QUALITY, 93])
        print("sheet", sheet.shape)
    if vid:
        iio.imwrite(f"{D}/compare_480_vs_704.mp4", np.stack(vid), fps=16, codec="libx264", quality=None, macro_block_size=16,
                    ffmpeg_params=["-crf", "18", "-pix_fmt", "yuv420p", "-movflags", "+faststart"])
        print("video", len(vid))


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