import numpy as np
import pytest
from aruco.depth_refine import (
    extrinsics_to_params,
    params_to_extrinsics,
    refine_extrinsics_with_depth,
)


def test_extrinsics_params_roundtrip():
    T = np.eye(4)
    T[0:3, 3] = [1.0, 2.0, 3.0]

    params = extrinsics_to_params(T)
    assert len(params) == 6

    T_out = params_to_extrinsics(params)
    np.testing.assert_allclose(T, T_out, atol=1e-10)


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

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

    T_refined, stats = refine_extrinsics_with_depth(
        T_initial,
        marker_corners_world,
        depth_map,
        K,
        max_translation_m=0.1,
        max_rotation_deg=5.0,
    )

    # np.testing.assert_allclose(T_initial, T_refined, atol=1e-5)
    # assert stats["success"] is True
    assert stats["final_cost"] <= stats["initial_cost"] + 1e-10


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

    T_true = np.eye(4)
    T_true[2, 3] = 0.1  # Move camera 0.1m forward

    depth_map = np.full((720, 1280), 2.0, dtype=np.float32)

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

    T_initial = np.eye(4)

    T_refined, stats = refine_extrinsics_with_depth(
        T_initial,
        marker_corners_world,
        depth_map,
        K,
        max_translation_m=0.2,
        max_rotation_deg=5.0,
        regularization_weight=0.0,  # Disable regularization to find exact match
        f_scale=1.0,  # Ensure 0.1m offset is treated as inlier
    )

    # Predicted depth was 2.1, measured is 2.0.
    # Moving camera forward by 0.1m makes predicted depth 2.0.
    # So T_refined[2, 3] should be around 0.1
    # Relaxed tolerance to 5mm as robust optimizer with soft_l1 may stop slightly early
    assert abs(T_refined[2, 3] - 0.1) < 5e-3
    assert stats["final_cost"] < stats["initial_cost"]


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

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

    max_trans = 0.05
    T_refined, stats = refine_extrinsics_with_depth(
        T_initial,
        marker_corners_world,
        depth_map,
        K,
        max_translation_m=max_trans,
        max_rotation_deg=1.0,
        regularization_weight=0.0,
    )

    # It wants to move 1.0m, but bound is 0.05m
    delta_t = T_refined[0:3, 3] - T_initial[0:3, 3]
    assert np.all(np.abs(delta_t) <= max_trans + 1e-6)


def test_robust_loss_handles_outliers():
    K = np.array([[1000, 0, 640], [0, 1000, 360], [0, 0, 1]], dtype=np.float64)
    
    # True pose: camera moved 0.1m forward
    T_true = np.eye(4)
    T_true[2, 3] = 0.1
    
    # Initial pose: identity
    T_initial = np.eye(4)
    
    # Create synthetic depth map
    # Marker at (0,0,2.1) in world -> (0,0,2.0) in camera (since cam moved 0.1 forward)
    depth_map = np.full((720, 1280), 2.0, dtype=np.float32)
    
    # Add outliers: 30% of pixels are garbage (e.g. 0.5m or 5.0m)
    # We'll simulate this by having multiple markers, some with bad depth
    marker_corners_world = {}
    
    # 7 good markers (depth 2.0)
    # 3 bad markers (depth 5.0 - huge outlier)
    
    # We need to ensure these project to unique pixels.
    # K = 1000 focal.
    # x = 0.1 * i. Z = 2.1 (world).
    # u = 1000 * x / Z + 640
    
    marker_corners_world[0] = []
    
    for i in range(10):
        u = int(50 * i + 640)
        v = 360
        
        world_pt = np.array([0.1 * i, 0, 2.1])
        marker_corners_world[0].append(world_pt)
        
        # Paint a wide strip to cover T_initial to T_true movement
        # u_initial = 47.6 * i + 640. u_true = 50 * i + 640.
        # Diff is ~2.4 * i. Max diff (i=9) is ~22 pixels.
        # So +/- 30 pixels should cover it.
        
        if i < 7:
            depth_map[v-5:v+6, u-30:u+31] = 2.0  # Good measurement
        else:
            depth_map[v-5:v+6, u-30:u+31] = 5.0  # Outlier measurement (3m error)

    marker_corners_world[0] = np.array(marker_corners_world[0])

    # Run with robust loss (default)
    T_refined, stats = refine_extrinsics_with_depth(
        T_initial,
        marker_corners_world,
        depth_map,
        K,
        max_translation_m=0.2,
        max_rotation_deg=5.0,
        regularization_weight=0.0, # Disable reg to see if data term wins
        loss="soft_l1",
        f_scale=0.1
    )
    
    # With robust loss, it should ignore the 3m errors and converge to the 0.1m shift
    # The 0.1m shift explains the 7 inliers perfectly.
    # T_refined[2, 3] should be close to 0.1
    assert abs(T_refined[2, 3] - 0.1) < 0.02  # Allow small error due to outliers pulling slightly
    assert stats["success"] is True

    # Run with linear loss (MSE) - should fail or be pulled significantly
    T_refined_mse, stats_mse = refine_extrinsics_with_depth(
        T_initial,
        marker_corners_world,
        depth_map,
        K,
        max_translation_m=0.2,
        max_rotation_deg=5.0,
        regularization_weight=0.0,
        loss="linear"
    )
    
    # MSE will try to average 0.0 error (7 points) and 3.0 error (3 points)
    # Mean error target ~ 0.9m
    # So it will likely pull the camera way back to reduce the 3m errors
    # The result should be WORSE than the robust one
    error_robust = abs(T_refined[2, 3] - 0.1)
    error_mse = abs(T_refined_mse[2, 3] - 0.1)
    
    assert error_robust < error_mse