"""3D → 2D projection for trajectory overlay.

One function. Used by:
- ``vis_dataset.py`` — project recorded EE poses (FK on saved joints)
  to the scene cam image for the dataset start-frame visualization.
- Deploy (``deploy.py``) — project predicted scene-cam-frame trajectory
  into the wrist image.
- Training viz — same as deploy, on a batch.

Generic — name it ``project_world_to_image``, not ``into_wrist``,
because either camera works.
"""
from __future__ import annotations

import numpy as np


def project_world_to_image(
    points_world: np.ndarray,        # (T, 3) world-frame 3D points
    K: np.ndarray,                   # (3, 3) camera intrinsics
    T_world_to_cam: np.ndarray,      # (4, 4) world → cam (= inv(T_cam_in_world))
    image_wh: tuple[int, int],       # (W, H) for in-frustum test
) -> tuple[np.ndarray, np.ndarray]:
    """Return ``(pix, in_frustum)``.

    ``pix``: (T, 2) float pixel coords (xy). Unclipped — caller decides.
    ``in_frustum``: (T,) bool — True if in front of cam AND inside image.
    """
    P = np.asarray(points_world, dtype=np.float64)
    R = T_world_to_cam[:3, :3]; t = T_world_to_cam[:3, 3]
    cam_pts = (R @ P.T).T + t                          # (T, 3)
    z = cam_pts[:, 2]
    in_front = z > 1e-6
    uv = (K @ cam_pts.T).T                             # (T, 3)
    pix = np.stack([uv[:, 0] / np.where(in_front, z, 1.0),
                    uv[:, 1] / np.where(in_front, z, 1.0)], axis=1)
    W, H = image_wh
    in_bounds = (pix[:, 0] >= 0) & (pix[:, 0] < W) & (pix[:, 1] >= 0) & (pix[:, 1] < H)
    return pix, (in_front & in_bounds)