import numpy as np
import pytest
from aruco.depth_pool import pool_depth_maps


def test_pool_depth_maps_empty():
    with pytest.raises(ValueError, match="depth_maps list cannot be empty"):
        _ = pool_depth_maps([])


def test_pool_depth_maps_shape_mismatch():
    dm1 = np.ones((10, 10))
    dm2 = np.ones((10, 11))
    with pytest.raises(ValueError, match="inconsistent shape"):
        _ = pool_depth_maps([dm1, dm2])


def test_pool_depth_maps_confidence_mismatch():
    dm1 = np.ones((10, 10))
    cm1 = np.ones((10, 10))
    with pytest.raises(ValueError, match="must match number of depth maps"):
        _ = pool_depth_maps([dm1], confidence_maps=[cm1, cm1])


def test_pool_depth_maps_confidence_shape_mismatch():
    dm1 = np.ones((10, 10))
    cm1 = np.ones((10, 11))
    with pytest.raises(ValueError, match="inconsistent shape"):
        _ = pool_depth_maps([dm1], confidence_maps=[cm1])


def test_pool_depth_maps_single_map():
    # N=1 returns masked copy behavior
    dm = np.array([[1.0, -1.0], [np.nan, 2.0]])
    pooled, conf = pool_depth_maps([dm])

    expected = np.array([[1.0, np.nan], [np.nan, 2.0]])
    np.testing.assert_allclose(pooled, expected)
    assert conf is None

    # Test min_valid_count > 1 for single map
    pooled, _ = pool_depth_maps([dm], min_valid_count=2)
    assert np.all(np.isnan(pooled))


def test_pool_depth_maps_median():
    # Median pooling with clean values
    dm1 = np.array([[1.0, 2.0], [3.0, 4.0]])
    dm2 = np.array([[1.2, 1.8], [3.2, 3.8]])
    dm3 = np.array([[0.8, 2.2], [2.8, 4.2]])

    pooled, _ = pool_depth_maps([dm1, dm2, dm3])

    # Median of [1.0, 1.2, 0.8] is 1.0
    # Median of [2.0, 1.8, 2.2] is 2.0
    # Median of [3.0, 3.2, 2.8] is 3.0
    # Median of [4.0, 3.8, 4.2] is 4.0
    expected = np.array([[1.0, 2.0], [3.0, 4.0]])
    np.testing.assert_allclose(pooled, expected)


def test_pool_depth_maps_invalid_handling():
    # NaN/invalid handling (non-finite or <=0)
    dm1 = np.array([[1.0, np.nan], [0.0, -1.0]])
    dm2 = np.array([[1.2, 2.0], [3.0, 4.0]])

    pooled, _ = pool_depth_maps([dm1, dm2])

    # (0,0): median(1.0, 1.2) = 1.1
    # (0,1): median(nan, 2.0) = 2.0
    # (1,0): median(0.0, 3.0) = 3.0 (0.0 is invalid)
    # (1,1): median(-1.0, 4.0) = 4.0 (-1.0 is invalid)
    expected = np.array([[1.1, 2.0], [3.0, 4.0]])
    np.testing.assert_allclose(pooled, expected)


def test_pool_depth_maps_confidence_gating():
    # Confidence gating (confidence > threshold excluded)
    dm1 = np.array([[1.0, 1.0], [1.0, 1.0]])
    dm2 = np.array([[2.0, 2.0], [2.0, 2.0]])

    cm1 = np.array([[10, 60], [10, 60]])
    cm2 = np.array([[60, 10], [10, 10]])

    # threshold = 50
    pooled, pooled_conf = pool_depth_maps(
        [dm1, dm2], confidence_maps=[cm1, cm2], confidence_thresh=50.0
    )

    # (0,0): dm1 valid (10), dm2 invalid (60) -> 1.0
    # (0,1): dm1 invalid (60), dm2 valid (10) -> 2.0
    # (1,0): dm1 valid (10), dm2 valid (10) -> 1.5
    # (1,1): dm1 invalid (60), dm2 valid (10) -> 2.0
    expected_depth = np.array([[1.0, 2.0], [1.5, 2.0]])
    expected_conf = np.array([[10, 10], [10, 10]])

    np.testing.assert_allclose(pooled, expected_depth)
    assert pooled_conf is not None
    np.testing.assert_allclose(pooled_conf, expected_conf)


def test_pool_depth_maps_all_invalid():
    # All invalid -> NaN outputs
    dm1 = np.array([[np.nan, 0.0], [-1.0, 1.0]])
    cm1 = np.array([[10, 10], [10, 100]])  # 100 > 50

    pooled, _ = pool_depth_maps([dm1], confidence_maps=[cm1], confidence_thresh=50.0)
    assert np.all(np.isnan(pooled))


def test_pool_depth_maps_min_valid_count():
    # min_valid_count enforcement
    dm1 = np.array([[1.0, 1.0], [1.0, 1.0]])
    dm2 = np.array([[2.0, 2.0], [np.nan, np.nan]])

    # min_valid_count = 2
    pooled, _ = pool_depth_maps([dm1, dm2], min_valid_count=2)

    # (0,0): 2 valid -> 1.5
    # (0,1): 2 valid -> 1.5
    # (1,0): 1 valid -> nan
    # (1,1): 1 valid -> nan
    expected = np.array([[1.5, 1.5], [np.nan, np.nan]])
    np.testing.assert_allclose(pooled, expected)


def test_pool_depth_maps_confidence_none():
    # confidence_maps None behavior
    dm1 = np.ones((2, 2))
    dm2 = np.ones((2, 2)) * 2

    pooled, conf = pool_depth_maps([dm1, dm2])
    assert conf is None
    np.testing.assert_allclose(pooled, np.ones((2, 2)) * 1.5)