"""Forward-project an IR depth image into RGB pixel coordinates.

Vectorised equivalent of librealsense `rs.align(stream.color)`. Each IR pixel
with valid depth is back-projected to 3D in the IR camera frame, transformed to
the RGB camera frame via the rigid extrinsics, and projected with RGB
intrinsics. On collisions the closest point wins (z-buffer).

Output is a uint16 array in millimeters, shape (rgb_height, rgb_width), with 0
on pixels that received no IR projection.
"""
from __future__ import annotations

import numpy as np

from .depth_to_rgb import CameraRegistration

_SENTINEL = np.iinfo(np.uint32).max


def align_depth_to_rgb(
    depth_ir_mm: np.ndarray,
    reg: CameraRegistration,
) -> np.ndarray:
    """Forward-project IR depth → RGB pixel grid.

    Parameters
    ----------
    depth_ir_mm : (H_ir, W_ir) uint16
        Depth in millimeters, IR camera frame. Zero means "no depth".
    reg : CameraRegistration
        IR and RGB intrinsics + IR→RGB rigid transform (translation in meters).

    Returns
    -------
    (reg.rgb_height, reg.rgb_width) uint16
        Depth in millimeters in RGB pixel coordinates; 0 where no IR pixel
        projected.
    """
    if depth_ir_mm.ndim != 2:
        raise ValueError(f"depth_ir_mm must be 2D, got shape {depth_ir_mm.shape}")
    if depth_ir_mm.shape != (reg.ir_height, reg.ir_width):
        raise ValueError(
            f"depth_ir_mm shape {depth_ir_mm.shape} ≠ IR resolution "
            f"({reg.ir_height}, {reg.ir_width}) from registration"
        )

    H_rgb, W_rgb = reg.rgb_height, reg.rgb_width

    valid = depth_ir_mm > 0
    if not valid.any():
        return np.zeros((H_rgb, W_rgb), dtype=np.uint16)

    v_idx, u_idx = np.where(valid)
    z_ir_m = depth_ir_mm[v_idx, u_idx].astype(np.float64) / 1000.0

    # Back-project IR pixel → 3D point in IR camera frame.
    K_ir_inv = np.linalg.inv(reg.K_ir)
    pixels_h = np.vstack([u_idx.astype(np.float64),
                          v_idx.astype(np.float64),
                          np.ones_like(u_idx, dtype=np.float64)])
    rays_ir = K_ir_inv @ pixels_h
    P_ir = rays_ir * z_ir_m

    # IR camera frame → RGB camera frame.
    P_rgb = reg.R_ir_to_rgb @ P_ir + reg.t_ir_to_rgb.reshape(3, 1)

    z_rgb = P_rgb[2]
    in_front = z_rgb > 1e-6
    P_rgb = P_rgb[:, in_front]
    if P_rgb.shape[1] == 0:
        return np.zeros((H_rgb, W_rgb), dtype=np.uint16)

    z_rgb = P_rgb[2]
    u_rgb = reg.K_rgb[0, 0] * P_rgb[0] / z_rgb + reg.K_rgb[0, 2]
    v_rgb = reg.K_rgb[1, 1] * P_rgb[1] / z_rgb + reg.K_rgb[1, 2]

    u_int = np.round(u_rgb).astype(np.int64)
    v_int = np.round(v_rgb).astype(np.int64)
    z_mm = np.round(z_rgb * 1000.0).astype(np.int64)

    in_bounds = (
        (u_int >= 0) & (u_int < W_rgb)
        & (v_int >= 0) & (v_int < H_rgb)
        & (z_mm > 0) & (z_mm < 65535)
    )
    u_int = u_int[in_bounds]
    v_int = v_int[in_bounds]
    z_mm = z_mm[in_bounds].astype(np.uint32)
    if z_mm.size == 0:
        return np.zeros((H_rgb, W_rgb), dtype=np.uint16)

    # Z-buffer: initialise to sentinel, reduce-min, then sentinel→0.
    zbuf = np.full(H_rgb * W_rgb, _SENTINEL, dtype=np.uint32)
    flat_idx = v_int * W_rgb + u_int
    np.minimum.at(zbuf, flat_idx, z_mm)
    zbuf[zbuf == _SENTINEL] = 0
    return zbuf.reshape(H_rgb, W_rgb).astype(np.uint16)