import pytest
import numpy as np
from unittest.mock import MagicMock, patch
import sys
from pathlib import Path

# Add py_workspace to path
sys.path.append(str(Path(__file__).parent.parent))

from calibrate_extrinsics import apply_depth_verify_refine_postprocess


@pytest.fixture
def mock_dependencies():
    with (
        patch("calibrate_extrinsics.verify_extrinsics_with_depth") as mock_verify,
        patch("calibrate_extrinsics.refine_extrinsics_with_depth") as mock_refine,
        patch("calibrate_extrinsics.click.echo") as mock_echo,
    ):
        # Setup mock return values
        mock_verify_res = MagicMock()
        mock_verify_res.rmse = 0.05
        mock_verify_res.mean_abs = 0.04
        mock_verify_res.median = 0.03
        mock_verify_res.depth_normalized_rmse = 0.02
        mock_verify_res.n_valid = 100
        mock_verify_res.n_total = 120
        mock_verify_res.residuals = []
        mock_verify.return_value = mock_verify_res

        mock_refine.return_value = (np.eye(4), {"success": True})

        yield mock_verify, mock_refine, mock_echo


def test_pool_size_1_equivalence(mock_dependencies):
    """
    Regression test: Ensure pool_size=1 behaves exactly like the old single-frame path.
    """
    mock_verify, _, _ = mock_dependencies

    serial = "123456"
    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}

    # Create a frame with specific depth values
    depth_map = np.ones((10, 10)) * 2.0
    conf_map = np.zeros((10, 10))

    frame_mock = MagicMock()
    frame_mock.depth_map = depth_map
    frame_mock.confidence_map = conf_map

    vf = {
        "frame": frame_mock,
        "ids": np.array([[1]]),
        "corners": np.zeros((1, 4, 2)),
        "score": 100.0,
    }

    # Structure for new implementation: list of frames
    verification_frames = {serial: [vf]}
    marker_geometry = {1: np.zeros((4, 3))}
    camera_matrices = {serial: np.eye(3)}

    # Run with pool_size=1
    apply_depth_verify_refine_postprocess(
        results=results,
        verification_frames=verification_frames,
        marker_geometry=marker_geometry,
        camera_matrices=camera_matrices,
        verify_depth=True,
        refine_depth=False,
        use_confidence_weights=False,
        depth_confidence_threshold=50,
        depth_pool_size=1,
    )

    # Verify that verify_extrinsics_with_depth was called with the exact depth map from the frame
    args, _ = mock_verify.call_args
    passed_depth_map = args[2]

    np.testing.assert_array_equal(passed_depth_map, depth_map)
    assert passed_depth_map is depth_map


def test_pool_size_5_integration(mock_dependencies):
    """
    Test that pool_size > 1 actually calls pooling and uses the result.
    """
    mock_verify, _, mock_echo = mock_dependencies

    serial = "123456"
    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}

    # Create 3 frames with different depth values
    # Frame 1: 2.0m
    # Frame 2: 2.2m
    # Frame 3: 1.8m
    # Median should be 2.0m

    frames = []
    for d in [2.0, 2.2, 1.8]:
        f = MagicMock()
        f.depth_map = np.ones((10, 10)) * d
        f.confidence_map = np.zeros((10, 10))
        frames.append(f)

    vfs = []
    for i, f in enumerate(frames):
        vfs.append(
            {
                "frame": f,
                "ids": np.array([[1]]),
                "corners": np.zeros((1, 4, 2)),
                "score": 100.0 - i,
            }
        )

    verification_frames = {serial: vfs}
    marker_geometry = {1: np.zeros((4, 3))}
    camera_matrices = {serial: np.eye(3)}

    # Run with pool_size=3
    apply_depth_verify_refine_postprocess(
        results=results,
        verification_frames=verification_frames,
        marker_geometry=marker_geometry,
        camera_matrices=camera_matrices,
        verify_depth=True,
        refine_depth=False,
        use_confidence_weights=False,
        depth_confidence_threshold=50,
        depth_pool_size=3,
    )

    # Check that "Using pooled depth" was logged
    any_pooled = any(
        "Using pooled depth" in str(call.args[0]) for call in mock_echo.call_args_list
    )
    assert any_pooled

    # Check that the depth map passed to verify is the median (2.0)
    args, _ = mock_verify.call_args
    passed_depth_map = args[2]

    expected_median = np.ones((10, 10)) * 2.0
    np.testing.assert_allclose(passed_depth_map, expected_median)

    # Verify metadata was added
    assert "depth_pool" in results[serial]
    assert results[serial]["depth_pool"]["pooled"] is True
    assert results[serial]["depth_pool"]["pool_size_actual"] == 3


def test_pool_fallback_insufficient_valid(mock_dependencies):
    """
    Test fallback to single frame when pooled result has too few valid points.
    """
    mock_verify, _, mock_echo = mock_dependencies

    serial = "123456"
    results = {serial: {"pose": "1 0 0 0 0 1 0 0 0 0 1 0 0 0 0 1"}}

    # Frame 1: Good depth
    f1 = MagicMock()
    f1.depth_map = np.ones((10, 10)) * 2.0
    f1.confidence_map = np.zeros((10, 10))

    # Frame 2: NaN depth (simulating misalignment or noise)
    f2 = MagicMock()
    f2.depth_map = np.full((10, 10), np.nan)
    f2.confidence_map = np.zeros((10, 10))

    # Frame 3: NaN depth
    f3 = MagicMock()
    f3.depth_map = np.full((10, 10), np.nan)
    f3.confidence_map = np.zeros((10, 10))

    # With median pooling, if >50% are NaN, result is NaN (standard median behavior with NaNs usually propagates or ignores)
    # Our pool_depth_maps uses nanmedian, which ignores NaNs.
    # But if we have [2.0, NaN, NaN], median of [2.0] is 2.0.
    # Wait, let's make it so they are valid but inconsistent to cause variance?
    # Or just force the pooled result to be bad by making them all different and sparse?

    # Let's use the fact that we can patch pool_depth_maps in the test!
    with patch("calibrate_extrinsics.pool_depth_maps") as mock_pool:
        # Return empty/invalid map
        mock_pool.return_value = (
            np.zeros((10, 10)),
            None,
        )  # Zeros are invalid depth (<=0)

        # Frame 1: Valid on left half
        d1 = np.full((10, 10), np.nan)
        d1[:, :5] = 2.0
        f1.depth_map = d1
        f1.confidence_map = np.zeros((10, 10))

        # Frame 2: Valid on right half
        d2 = np.full((10, 10), np.nan)
        d2[:, 5:] = 2.0
        f2.depth_map = d2
        f2.confidence_map = np.zeros((10, 10))

        vfs = [
            {
                "frame": f1,
                "ids": np.array([[1]]),
                "corners": np.zeros((1, 4, 2)),
                "score": 100,
            },
            {
                "frame": f2,
                "ids": np.array([[1]]),
                "corners": np.zeros((1, 4, 2)),
                "score": 90,
            },
        ]

        verification_frames = {serial: vfs}
        marker_geometry = {1: np.zeros((4, 3))}
        camera_matrices = {serial: np.eye(3)}

        apply_depth_verify_refine_postprocess(
            results=results,
            verification_frames=verification_frames,
            marker_geometry=marker_geometry,
            camera_matrices=camera_matrices,
            verify_depth=True,
            refine_depth=False,
            use_confidence_weights=False,
            depth_confidence_threshold=50,
            depth_pool_size=2,
        )

        # Check for fallback message
        any_fallback = any(
            "Falling back to best single frame" in str(call.args[0])
            for call in mock_echo.call_args_list
        )
        assert any_fallback

        # Verify we used the best frame (f1)
        args, _ = mock_verify.call_args
        passed_depth_map = args[2]
        assert passed_depth_map is d1

        # Verify metadata
        assert results[serial]["depth_pool"]["pooled"] is False
        assert (
            results[serial]["depth_pool"]["fallback_reason"]
            == "insufficient_valid_points"
        )