import numpy as np
import pytest
from aruco.depth_verify import (
    project_point_to_pixel,
    compute_depth_residual,
    compute_marker_corner_residuals,
    verify_extrinsics_with_depth,
    DepthVerificationResult,
)


def test_project_point_to_pixel():
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)

    # Point directly in front at 2m
    P_cam = np.array([0, 0, 2.0])
    u, v = project_point_to_pixel(P_cam, K)
    assert u == 640
    assert v == 360

    # Point offset
    P_cam = np.array([0.2, -0.1, 2.0])
    # u = 1000 * 0.2 / 2.0 + 640 = 100 + 640 = 740
    # v = 1000 * -0.1 / 2.0 + 360 = -50 + 360 = 310
    u, v = project_point_to_pixel(P_cam, K)
    assert u == 740
    assert v == 310

    # Point behind camera
    P_cam = np.array([0, 0, -1.0])
    u, v = project_point_to_pixel(P_cam, K)
    assert u is None
    assert v is None


def test_compute_depth_residual():
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
    T_world_cam = np.eye(4)  # Camera at origin, looking along Z

    # Create a synthetic depth map (100x100)
    depth_map = np.full((720, 1280), 2.0, dtype=np.float32)

    # Point at (0, 0, 2) in world/cam coords
    P_world = np.array([0, 0, 2.0])

    # Perfect case
    residual = compute_depth_residual(P_world, T_world_cam, depth_map, K, window_size=1)
    assert residual is not None
    assert abs(residual) < 1e-6

    # Offset case (measured is 2.1, predicted is 2.0)
    depth_map[360, 640] = 2.1
    residual = compute_depth_residual(P_world, T_world_cam, depth_map, K, window_size=1)
    assert residual is not None
    assert abs(residual - 0.1) < 1e-6

    # Invalid depth case
    depth_map[360, 640] = np.nan
    residual = compute_depth_residual(P_world, T_world_cam, depth_map, K, window_size=1)
    assert residual is None

    # Window size case
    depth_map[358:363, 638:643] = 2.2
    residual = compute_depth_residual(P_world, T_world_cam, depth_map, K, window_size=5)
    assert residual is not None
    assert abs(residual - 0.2) < 1e-6


def test_compute_marker_corner_residuals():
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
    T_world_cam = np.eye(4)
    depth_map = np.full((720, 1280), 2.0, dtype=np.float32)

    marker_corners_world = {
        1: np.array([[0, 0, 2.0], [0.1, 0, 2.0], [0.1, 0.1, 2.0], [0, 0.1, 2.0]])
    }

    residuals = compute_marker_corner_residuals(
        T_world_cam, marker_corners_world, depth_map, K
    )
    assert len(residuals) == 4
    for marker_id, corner_idx, res in residuals:
        assert marker_id == 1
        assert abs(res) < 1e-6


def test_verify_extrinsics_with_depth():
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
    T_world_cam = np.eye(4)
    depth_map = np.full((720, 1280), 2.0, dtype=np.float32)

    # Add some noise/offset - fill the 5x5 window because compute_marker_corner_residuals uses window_size=5
    depth_map[358:363, 638:643] = 2.1

    marker_corners_world = {1: np.array([[0, 0, 2.0]])}

    result = verify_extrinsics_with_depth(
        T_world_cam, marker_corners_world, depth_map, K
    )
    assert isinstance(result, DepthVerificationResult)
    assert result.n_valid == 1
    assert result.n_total == 1
    assert abs(result.rmse - 0.1) < 1e-6
    assert abs(result.mean_abs - 0.1) < 1e-6
    assert abs(result.median - 0.1) < 1e-6
    # depth_normalized_rmse = 0.1 / 2.0 = 0.05
    assert abs(result.depth_normalized_rmse - 0.05) < 1e-6