"""Unit tests for align_depth_to_rgb."""
from __future__ import annotations

import numpy as np

from presa_mano.perception.depth_align import align_depth_to_rgb
from presa_mano.perception.depth_to_rgb import CameraRegistration


def _identity_reg(W: int = 64, H: int = 48) -> CameraRegistration:
    """Same K for RGB and IR, no extrinsic offset."""
    K = np.array([[60.0, 0.0, W / 2.0],
                  [0.0,  60.0, H / 2.0],
                  [0.0,  0.0,  1.0]], dtype=np.float64)
    return CameraRegistration(
        K_rgb=K.copy(), K_ir=K.copy(),
        R_ir_to_rgb=np.eye(3, dtype=np.float64),
        t_ir_to_rgb=np.zeros(3, dtype=np.float64),
        rgb_width=W, rgb_height=H,
        ir_width=W, ir_height=H,
    )


def test_identity_alignment_preserves_depth():
    """K_ir == K_rgb and no extrinsic → every IR pixel maps to itself."""
    reg = _identity_reg()
    depth_ir = np.full((reg.ir_height, reg.ir_width), 1000, dtype=np.uint16)
    depth_rgb = align_depth_to_rgb(depth_ir, reg)
    assert depth_rgb.shape == (reg.rgb_height, reg.rgb_width)
    assert depth_rgb.dtype == np.uint16
    nonzero = depth_rgb[depth_rgb > 0]
    # Coverage close to full (rounding can drop a few border pixels).
    assert nonzero.size >= 0.95 * depth_rgb.size, f"coverage too low: {nonzero.size}/{depth_rgb.size}"
    # All filled pixels equal the input depth (±1 mm rounding).
    assert np.all(np.abs(nonzero.astype(int) - 1000) <= 1)


def test_zero_depth_pixels_stay_zero():
    """IR pixels with depth=0 must not produce any RGB output."""
    reg = _identity_reg()
    depth_ir = np.zeros((reg.ir_height, reg.ir_width), dtype=np.uint16)
    depth_rgb = align_depth_to_rgb(depth_ir, reg)
    assert depth_rgb.shape == (reg.rgb_height, reg.rgb_width)
    assert int(depth_rgb.sum()) == 0


def test_sparse_pixels_preserved():
    """Two isolated IR depths produce exactly two filled RGB pixels (identity reg)."""
    reg = _identity_reg(W=64, H=48)
    depth_ir = np.zeros((reg.ir_height, reg.ir_width), dtype=np.uint16)
    depth_ir[0, 0] = 1000
    depth_ir[reg.ir_height - 1, reg.ir_width - 1] = 1000
    depth_rgb = align_depth_to_rgb(depth_ir, reg)
    assert (depth_rgb > 0).sum() == 2


def test_z_buffer_keeps_closer_point():
    """When two IR pixels collide on the same RGB pixel, the smaller depth wins.

    Force a collision with K_ir = 4 × K_rgb (same resolution), so adjacent IR
    pixels round to the same RGB pixel."""
    W = H = 16
    K_rgb = np.array([[60.0, 0.0, W / 2.0],
                      [0.0,  60.0, H / 2.0],
                      [0.0,  0.0,  1.0]], dtype=np.float64)
    K_ir = K_rgb.copy()
    K_ir[0, 0] *= 4.0   # 4× focal length → adjacent IR pixels collapse to one RGB pixel
    K_ir[1, 1] *= 4.0
    reg = CameraRegistration(
        K_rgb=K_rgb, K_ir=K_ir,
        R_ir_to_rgb=np.eye(3, dtype=np.float64),
        t_ir_to_rgb=np.zeros(3, dtype=np.float64),
        rgb_width=W, rgb_height=H, ir_width=W, ir_height=H,
    )
    depth_ir = np.zeros((H, W), dtype=np.uint16)
    # IR pixels (8, 8) and (8, 9) both project to RGB pixel (8, 8):
    # u_rgb = 0.25*(u_ir - 8) + 8 → 8.0 and 8.25 (both round to 8).
    depth_ir[8, 8] = 1500
    depth_ir[8, 9] = 800
    depth_rgb = align_depth_to_rgb(depth_ir, reg)
    nonzero = depth_rgb[depth_rgb > 0]
    # Exactly one filled pixel; z-buffer kept the closer (800) value.
    assert nonzero.size == 1
    assert int(nonzero[0]) == 800
    assert int(depth_rgb[8, 8]) == 800