#!/usr/bin/env python3
"""res_test.py 的补充（分辨率决策的证据）：detail_ratio = |Laplacian| 能量比只看高频「多不多」，不看「在不在对的位置」——
没有重采样低通的 1280×704 会被奖励，哪怕多出来的高频是 VAE 伪影；crop PSNR 又偏爱模糊。这里在同一批 player crop、
同一条 VAE 往返（res_test.enc / dec / to_res 原样复用）上加三个看「对齐」的指标：
  ssim         灰度 SSIM（高斯窗 σ 1.5，标准常数）
  lap_corr     Laplacian(重建) 与 Laplacian(GT) 的 Pearson 相关：高频要在对的位置才高
  edge_f1      Canny(50,150) 边缘，1 px 容差的 F1（precision = 重建边缘落在 GT 边缘旁的比例；recall 反之）
psnr / detail_ratio 也重算一遍，对拍 outputs/char_diag/res_test.json（应相同）。写 outputs/char_diag/res_test_edge.json，
不动 res_test.json。默认 GC_RES（_ref 的 ids 是 832×480）。

    CUDA_VISIBLE_DEVICES=0 $PY scripts/res_test_edge.py
"""
import json
import os
import sys

import cv2
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import res_test as RT  # noqa: E402  （enc / dec / to_res / RES / NF / D 原样复用）

import geometry as GEO  # noqa: E402
import infer as INF  # noqa: E402
import model as GM  # noqa: E402
import sample as S  # noqa: E402

K3 = np.ones((3, 3), np.uint8)


def gray(a):
    return cv2.cvtColor(np.ascontiguousarray(a), cv2.COLOR_RGB2GRAY)


def ssim(a, b):
    a, b = gray(a).astype(np.float64), gray(b).astype(np.float64)
    c1, c2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
    blur = lambda x: cv2.GaussianBlur(x, (11, 11), 1.5)  # noqa: E731
    ma, mb = blur(a), blur(b)
    va, vb, cov = blur(a * a) - ma * ma, blur(b * b) - mb * mb, blur(a * b) - ma * mb
    return float((((2 * ma * mb + c1) * (2 * cov + c2)) / ((ma * ma + mb * mb + c1) * (va + vb + c2))).mean())


def lap_corr(a, b):
    la = cv2.Laplacian(gray(a).astype(np.float32), cv2.CV_32F).ravel()
    lb = cv2.Laplacian(gray(b).astype(np.float32), cv2.CV_32F).ravel()
    return float(np.corrcoef(la, lb)[0, 1])


def edge_f1(rc, gt):
    er, eg = cv2.Canny(gray(rc), 50, 150) > 0, cv2.Canny(gray(gt), 50, 150) > 0
    if not er.any() or not eg.any():
        return 0.0, 0.0, 0.0
    dr, dg = cv2.dilate(er.astype(np.uint8), K3) > 0, cv2.dilate(eg.astype(np.uint8), K3) > 0
    p, r = float((er & dg).sum() / er.sum()), float((eg & dr).sum() / eg.sum())
    return p, r, (2 * p * r / (p + r) if p + r > 0 else 0.0)


def main():
    assert GEO.RES_TAG == GEO.LEGACY_RES, "_ref 的 ids 是 832×480：用默认 GC_RES 跑"
    clips = json.load(open(f"{RT.D}/clips.json"))
    vae = INF.load_vae("cuda", S.MODEL_DIR)
    out = {k: [] for k in RT.RES}
    for cid in clips:
        ref = np.load(f"{RT.D}/_ref/{cid}.npz")
        src, ids = ref["gt_src"][:RT.NF], ref["ids"]
        for name, wh in RT.RES.items():
            x, x0, cw = (src, 0, 1280) if name == "1280x704" else RT.to_res(src, wh)
            rec = RT.dec(vae, RT.enc(vae, np.ascontiguousarray(x)))
            for t in range(8, RT.NF, 8):                       # 与 res_test.py 同一套 crop
                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]
                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(gray(a).astype(np.float32), cv2.CV_32F)).mean())  # noqa: E731
                ep, er, ef = edge_f1(rc, gt_c)
                out[name].append(dict(psnr=10 * np.log10(255 ** 2 / max(mse, 1e-6)), detail=g(rc) / max(g(gt_c), 1e-6),
                                      ssim=ssim(rc, gt_c), lap_corr=lap_corr(rc, gt_c), edge_p=ep, edge_r=er, edge_f1=ef))
        print(cid, {k: len(v) for k, v in out.items()}, flush=True)
    res = {k: {"n": len(v), **{m: round(float(np.mean([d[m] for d in v])), 4) for m in v[0]}} for k, v in out.items()}
    base = json.load(open(f"{RT.D}/res_test.json"))["vae_ceiling"]
    for k in res:                                              # 对拍 res_test.json（同一条路径 ⇒ 同数）
        res[k]["matches_res_test_json"] = (abs(res[k]["psnr"] - base[k]["player_crop_psnr_vs_native_gt"]) < 0.01
                                           and abs(res[k]["detail"] - base[k]["detail_ratio_vs_native_gt"]) < 0.001)
    print(json.dumps(res, indent=1), flush=True)
    json.dump({"vae_ceiling_alignment": res, "note": __doc__}, open(f"{RT.D}/res_test_edge.json", "w"), indent=1,
              ensure_ascii=False)


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