#!/usr/bin/env python3
"""角色清晰度诊断·分析端：按 ids 里的角色掩码裁框，比较 GT(源 1280×704) / GT(832×480) / VAE 重建 / 各生成变体。
指标（框内，相对 GT832）：sharp = Sobel 梯度均值之比，lap = Laplacian 方差之比，psnr。另出无损裁剪对比图。"""
import json
import os
import sys

import cv2
import numpy as np

import geometry as GEO
import model as GM

M = "/opt/dlami/nvme/zhiyangdeng/maliang"
D = f"{M}/outputs/char_diag"
VARIANTS = ["s4000_cfg1", "s8000_cfg1", "s8000_semcfg2", "s8000_semcfg3"]
ROLES = {"player": GM.PLAYER_PART_IDS, "boss": GM.BOSS_PART_IDS}
PAD, MIN_PX, UP = 10, 600, 3


def bbox(mask):
    ys, xs = np.nonzero(mask)
    return max(0, ys.min() - PAD), min(mask.shape[0], ys.max() + PAD + 1), max(0, xs.min() - PAD), min(mask.shape[1], xs.max() + PAD + 1)


def grad(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.sqrt(gx ** 2 + gy ** 2).mean()), float(cv2.Laplacian(g, cv2.CV_32F).var())


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


def src_crop(src, box):
    """GT 源分辨率里与 832×480 框对应的区域，缩放到框的尺寸 ×UP（看 1280×704 原生能有多清楚）。"""
    y0, y1, x0, x1 = box
    p = GEO.dst_to_src(np.array([[x0 - 0.5, y0 - 0.5], [x1 - 0.5, y1 - 0.5]]))
    sx0, sy0 = int(max(0, np.floor(p[0, 0]))), int(max(0, np.floor(p[0, 1])))
    sx1, sy1 = int(min(src.shape[1], np.ceil(p[1, 0]))), int(min(src.shape[0], np.ceil(p[1, 1])))
    return cv2.resize(src[sy0:sy1, sx0:sx1], ((x1 - x0) * UP, (y1 - y0) * UP), interpolation=cv2.INTER_CUBIC)


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():
    clips = json.load(open(f"{D}/clips.json"))
    variants = [v for v in VARIANTS if all(os.path.exists(f"{D}/{v}/{c}.npz") for c in clips)]
    stats = {r: {k: [] for k in ["recon"] + variants} for r in ROLES}
    sheet_rows = []
    for ci, cid in enumerate(clips):
        ref = np.load(f"{D}/_ref/{cid}.npz")
        gt, rec, src, ids = ref["gt"], ref["recon"], ref["gt_src"], ref["ids"]
        gens = {v: np.load(f"{D}/{v}/{cid}.npz")["frames"] for v in variants}
        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 = box = bbox(m)
                g_ref, l_ref = grad(gt[t, y0:y1, x0:x1])
                for k, fr in [("recon", rec)] + list(gens.items()):
                    c = fr[t, y0:y1, x0:x1]
                    g, l = grad(c)
                    stats[role][k].append((g / max(g_ref, 1e-6), l / max(l_ref, 1e-6), psnr(c, gt[t, y0:y1, x0:x1])))
                if role == "player" and t in (40, 72) and len(sheet_rows) < 12:
                    up = lambda a: cv2.resize(a, ((x1 - x0) * UP, (y1 - y0) * UP), interpolation=cv2.INTER_NEAREST)
                    tiles = [label(src_crop(src[t], box), "GT 1280x704 (native)"), label(up(gt[t, y0:y1, x0:x1]), "GT 832x480"),
                             label(up(rec[t, y0:y1, x0:x1]), "VAE recon")]
                    tiles += [label(up(gens[v][t, y0:y1, x0:x1]), v) for v in variants]
                    h = max(x.shape[0] for x in tiles)
                    tiles = [cv2.copyMakeBorder(x, 0, h - x.shape[0], 0, 4, cv2.BORDER_CONSTANT) for x in tiles]
                    sheet_rows.append(np.concatenate(tiles, 1))
    summary = {}
    for role, d in stats.items():
        summary[role] = {k: {"n": len(v), "sharp_ratio": round(float(np.mean([x[0] for x in v])), 3),
                             "lapvar_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}
    json.dump(summary, open(f"{D}/summary.json", "w"), indent=1)
    print(json.dumps(summary, indent=1))
    if sheet_rows:
        w = max(r.shape[1] for r in sheet_rows)
        rows = [cv2.copyMakeBorder(r, 0, 6, 0, w - r.shape[1], cv2.BORDER_CONSTANT) for r in sheet_rows]
        sheet = np.concatenate(rows, 0)
        cv2.imwrite(f"{D}/player_crops.png", sheet[:, :, ::-1])
        print("sheet", sheet.shape)


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