#!/usr/bin/env python3
"""把手部网格（MANO、SOMA 或任意三角网格）投影到 RGB 上，生成半透明叠加图。

只依赖 numpy 和 Pillow。可单独运行，也被 server_sync.py 调用：

    python tools/render_overlay.py --image frame.png --mesh hand.npz --out frame_gt.png

mesh.npz 需要：
    verts  (N, 3)  相机坐标系下的顶点，单位米，z 朝前
    faces  (F, 3)  三角面索引
    K      (3, 3)  内参（也可用 --K 传入 JSON）
可选：
    verts 可以是 (H, N, 3)，表示多只手；faces 共用
    joints (J, 3) 或 (H, J, 3)，会画成小圆点
"""
import argparse
import json

import numpy as np
from PIL import Image, ImageDraw

HAND_COLORS = [(70, 140, 230), (235, 140, 50)]  # 右手蓝、左手橙（按传入顺序）


def project(p, K):
    z = np.clip(p[:, 2], 1e-6, None)
    return np.stack([K[0, 0] * p[:, 0] / z + K[0, 2], K[1, 1] * p[:, 1] / z + K[1, 2]], 1)


def overlay(img, verts, faces, K, joints=None, alpha=0.55, colors=HAND_COLORS):
    """img: PIL.Image；verts: (N,3) 或 (H,N,3)；返回新的 PIL.Image。"""
    base = img.convert("RGBA")
    layer = Image.new("RGBA", base.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(layer)
    V = np.asarray(verts, float)
    if V.ndim == 2:
        V = V[None]
    F = np.asarray(faces, int)
    K = np.asarray(K, float)
    light = np.array([0.3, -0.4, -0.85])
    light /= np.linalg.norm(light)

    tris = []  # (depth, polygon, rgba)
    for h, v in enumerate(V):
        col = np.array(colors[h % len(colors)], float)
        uv = project(v, K)
        a, b, c = v[F[:, 0]], v[F[:, 1]], v[F[:, 2]]
        n = np.cross(b - a, c - a)
        n /= np.linalg.norm(n, axis=1, keepdims=True) + 1e-12
        shade = 0.45 + 0.55 * np.abs(n @ light)
        depth = (a[:, 2] + b[:, 2] + c[:, 2]) / 3
        ok = (a[:, 2] > 0) & (b[:, 2] > 0) & (c[:, 2] > 0)
        for i in np.nonzero(ok)[0]:
            poly = [tuple(uv[j]) for j in F[i]]
            rgb = tuple(int(x) for x in np.clip(col * shade[i], 0, 255))
            tris.append((depth[i], poly, rgb + (int(255 * alpha),)))
    for _, poly, rgba in sorted(tris, key=lambda t: -t[0]):  # 远的先画
        draw.polygon(poly, fill=rgba)
    out = Image.alpha_composite(base, layer)

    if joints is not None:
        J = np.asarray(joints, float)
        if J.ndim == 2:
            J = J[None]
        d = ImageDraw.Draw(out)
        r = max(2, round(min(out.size) / 220))
        for j in J:
            for x, y in project(j, K):
                d.ellipse([x - r, y - r, x + r, y + r], fill=(255, 255, 255, 255), outline=(20, 20, 20, 255))
    return out.convert("RGB")


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--image", required=True)
    ap.add_argument("--mesh", required=True, help="npz：verts, faces, 可选 K / joints")
    ap.add_argument("--K", help="内参 JSON，例如 [[fx,0,cx],[0,fy,cy],[0,0,1]]")
    ap.add_argument("--alpha", type=float, default=0.55)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()
    m = np.load(a.mesh)
    K = np.array(json.loads(a.K)) if a.K else m["K"]
    img = overlay(Image.open(a.image), m["verts"], m["faces"], K,
                  joints=m["joints"] if "joints" in m.files else None, alpha=a.alpha)
    img.save(a.out)
    print("写出", a.out)


if __name__ == "__main__":
    main()