# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Render the predicted MANO mesh on the *original* image.

Projection convention derived from ``hamer/utils/renderer.py`` upstream.
The upstream applies a 180° X-rotation to vertices (trimesh) *and* then
renders with pyrender/OpenGL which does a y-flip (NDC y-up → image y-down).
Those two y-negations cancel.  Our equivalent — without pyrender — is to
project the MANO vertices **directly** without any rotation:

    proj_x = focal * (vx + tx) / (vz + tz) + W/2
    proj_y = focal * (vy + ty) / (vz + tz) + H/2

where (tx, ty, tz) = cam_crop_to_full(pred_cam, ...).  Applying only the
rotation (without the matching y-flip) inverts every vertex's y relative
to the hand center, making a downward-facing palm appear upward.

Triangles are rasterized as alpha-blended polygons with simple diffuse
shading from face-normal · view-vector — no ``pyrender`` / ``trimesh``
dependency.  The shading is enough to make finger curl readable instead of
the flat silhouette we used to ship.
"""
from __future__ import annotations

from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import torch

# matplotlib is imported lazily inside ``overlay_mesh`` so the projection
# helpers (``cam_crop_to_full`` / ``project_mesh``) stay importable in a
# runtime that ships no matplotlib (the serving image).

FOCAL_LENGTH = 5000.0
MODEL_IMAGE_SIZE = 256


def cam_crop_to_full(
    cam_bbox: np.ndarray,
    box_center: np.ndarray,
    box_size: float,
    img_w: int,
    img_h: int,
    focal: float,
) -> np.ndarray:
    """Replicates ``hamer.utils.renderer.cam_crop_to_full`` on numpy.

    Returns the 3-D camera translation ``(tx, ty, tz)`` in original-image
    pixel-scaled world coordinates.
    """
    s, tx_b, ty_b = float(cam_bbox[0]), float(cam_bbox[1]), float(cam_bbox[2])
    cx, cy = float(box_center[0]), float(box_center[1])
    bs = box_size * s + 1e-9
    tz = 2.0 * focal / bs
    tx = 2.0 * (cx - img_w / 2.0) / bs + tx_b
    ty = 2.0 * (cy - img_h / 2.0) / bs + ty_b
    return np.array([tx, ty, tz], dtype=np.float64)


def project_mesh(
    verts: torch.Tensor,
    cam_bbox: torch.Tensor,
    box_center: Tuple[float, float],
    box_size: float,
    img_w: int,
    img_h: int,
    is_right: bool = True,
) -> Tuple[np.ndarray, np.ndarray]:
    """Project MANO vertices to ORIGINAL-image pixel coords.

    For a LEFT hand the upstream pipeline runs the (right-hand) model on a
    horizontally-mirrored input crop, then flips ``cam.x`` and ``verts.x``
    back before rendering — see ``demo.py`` lines 136-138 / 178-179.
    Reproduce that here.

    Returns ``(verts_2d, verts_camera_z)`` — the second array is per-vertex
    camera-space z (positive = depth in front of the camera) used for
    back-to-front face sorting in rendering.
    """
    focal = FOCAL_LENGTH / MODEL_IMAGE_SIZE * max(img_w, img_h)
    cam_bbox_np = cam_bbox.detach().cpu().numpy().copy()
    if not is_right:
        cam_bbox_np[1] *= -1.0
    cam_t = cam_crop_to_full(cam_bbox_np, np.array(box_center), box_size, img_w, img_h, focal)

    v = verts.detach().cpu().numpy().astype(np.float64).copy()
    if not is_right:
        v[:, 0] *= -1.0
    # Direct projection — no rotation applied (upstream's 180° X rotation and
    # pyrender's NDC y-flip cancel; applying only the rotation inverts y).
    v_cam = v + cam_t
    z = v_cam[:, 2]
    proj = focal * v_cam[:, :2] / z[:, None]
    proj[:, 0] += img_w / 2.0
    proj[:, 1] += img_h / 2.0
    return proj, z


def overlay_mesh(
    ax,
    image_uint8: np.ndarray,
    verts_2d: np.ndarray,
    verts_z: np.ndarray,
    faces: np.ndarray,
    color=(0.20, 0.85, 1.00),
    alpha: float = 0.65,
    is_right: bool = True,
) -> None:
    """Rasterize the projected mesh on ``image_uint8`` with face-normal
    diffuse shading + back-to-front depth ordering."""
    from matplotlib.collections import PolyCollection  # noqa: WPS433 (visualisation only)

    h, w = image_uint8.shape[:2]
    ax.imshow(image_uint8)
    ax.set_xlim(0, w)
    ax.set_ylim(h, 0)
    ax.set_aspect("equal")
    ax.axis("off")

    tri_z = verts_z[faces].mean(axis=1)               # (F,) camera z per face
    tris2d = verts_2d[faces]                          # (F, 3, 2)

    # Cull triangles whose mean z is behind the camera or way off-screen.
    inside = (
        (tris2d[..., 0] > -50).all(-1)
        & (tris2d[..., 0] < w + 50).all(-1)
        & (tris2d[..., 1] > -50).all(-1)
        & (tris2d[..., 1] < h + 50).all(-1)
        & (tri_z > 0.05)
    )
    tris2d, tri_z = tris2d[inside], tri_z[inside]

    # Winding: without the 180° X rotation, projecting v directly means
    # vy > 0 maps to larger proj_y (lower in image).  A right-hand
    # front-facing triangle that is CCW in 3D world projects to CW in
    # y-down image (cross_z > 0).  The x-flip for left hands reverses
    # winding, so front-facing left-hand triangles have cross_z < 0.
    e1 = tris2d[:, 1] - tris2d[:, 0]
    e2 = tris2d[:, 2] - tris2d[:, 0]
    cross_z = e1[:, 0] * e2[:, 1] - e1[:, 1] * e2[:, 0]
    winding_sign = 1.0 if is_right else -1.0
    facing = np.clip(winding_sign * cross_z / (np.abs(cross_z).max() + 1e-9), -1, 1)
    diffuse = 0.55 + 0.45 * facing

    # Painter's order: render farthest faces first.
    order = np.argsort(-tri_z)
    tris2d = tris2d[order]
    diffuse = diffuse[order]

    facecolors = np.empty((len(tris2d), 4), dtype=np.float32)
    facecolors[:, 0] = color[0] * diffuse
    facecolors[:, 1] = color[1] * diffuse
    facecolors[:, 2] = color[2] * diffuse
    facecolors[:, 3] = alpha

    coll = PolyCollection(
        tris2d,
        facecolors=facecolors,
        edgecolors=(0.05, 0.10, 0.15, 0.06),
        linewidths=0.25,
    )
    ax.add_collection(coll)