"""Tests for distortion and undistortion operations."""

import numpy as np
import pytest

import deltacamera
import deltacamera.distortion
import deltacamera.validity
from conftest import (
    BROWN_CONRADY_COEFFS,
    BROWN_CONRADY_5_COEFFS,
    FISHEYE_COEFFS,
    extend_distortion_coeffs,
    sample_inside_polygon,
    make_camera_with_distortion,
)


# =============================================================================
# Brown-Conrady distortion roundtrip tests
# =============================================================================

class TestBrownConradyExRoundtrip:
    """Test that undistort(distort(p)) ≈ p for Brown-Conrady model."""

    @pytest.mark.parametrize("d", BROWN_CONRADY_COEFFS)
    def test_roundtrip_12_param(self, d):
        """Roundtrip with 12-parameter model."""
        camera = make_camera_with_distortion(d)
        # get_valid_distortion_region returns (undistorted_region, distorted_region)
        valid_region_ud, _ = deltacamera.validity.get_valid_distortion_region(
            camera, cartesian=True, n_vertices=128
        )

        # Check if boundary is at r=5 default (no Jacobian=0 crossing)
        # If so, use smaller region to avoid extreme distortion values
        boundary_r = np.max(np.linalg.norm(valid_region_ud, axis=1))
        if boundary_r > 4.5:
            # Practical limit, not real boundary - sample from smaller region
            # Use uniform in [-0.3, 0.3] to ensure bounded radius
            points = (np.random.rand(500, 2).astype(np.float32) - 0.5) * 0.6
        else:
            points = sample_inside_polygon(valid_region_ud, n=500, margin=0.1)

        distorted = deltacamera.distortion.distort_points(
            points, d, check_validity=False
        )
        recovered = deltacamera.distortion.undistort_points(
            distorted, d, check_validity=False
        )

        np.testing.assert_allclose(recovered, points, rtol=1e-4, atol=1e-6)

    @pytest.mark.parametrize("d", BROWN_CONRADY_5_COEFFS)
    def test_roundtrip_5_param(self, d):
        """Roundtrip with 5-parameter model (OpenCV default)."""
        d12 = extend_distortion_coeffs(d, 12)
        camera = make_camera_with_distortion(d12)

        # get_valid_distortion_region returns (undistorted_region, distorted_region)
        valid_region_ud, _ = deltacamera.validity.get_valid_distortion_region(
            camera, cartesian=True, n_vertices=128
        )

        # Check if boundary is at r=5 default (no Jacobian=0 crossing)
        boundary_r = np.max(np.linalg.norm(valid_region_ud, axis=1))
        if boundary_r > 4.5:
            # Practical limit - sample from smaller region
            points = (np.random.rand(500, 2).astype(np.float32) - 0.5) * 0.6
        else:
            points = sample_inside_polygon(valid_region_ud, n=500, margin=0.1)

        distorted = deltacamera.distortion.distort_points(
            points, d12, check_validity=False
        )
        recovered = deltacamera.distortion.undistort_points(
            distorted, d12, check_validity=False
        )

        np.testing.assert_allclose(recovered, points, rtol=1e-4, atol=1e-6)

    def test_roundtrip_at_origin(self):
        """Point at origin should be unchanged."""
        d = BROWN_CONRADY_COEFFS[0]
        origin = np.array([[0, 0]], np.float32)

        distorted = deltacamera.distortion.distort_points(
            origin, d, check_validity=False
        )
        np.testing.assert_allclose(distorted, origin, atol=1e-10)

        recovered = deltacamera.distortion.undistort_points(
            distorted, d, check_validity=False
        )
        np.testing.assert_allclose(recovered, origin, atol=1e-10)

    def test_roundtrip_small_radius(self):
        """Points very close to center should have minimal distortion."""
        d = BROWN_CONRADY_COEFFS[0]
        small_points = np.random.randn(100, 2).astype(np.float32) * 0.01

        distorted = deltacamera.distortion.distort_points(
            small_points, d, check_validity=False
        )
        # Small points should barely move
        np.testing.assert_allclose(distorted, small_points, rtol=0.01, atol=1e-4)


# =============================================================================
# Fisheye distortion roundtrip tests
# =============================================================================

class TestFisheyeRoundtrip:
    """Test that undistort(distort(p)) ≈ p for fisheye model."""

    @pytest.mark.parametrize("d", FISHEYE_COEFFS)
    def test_roundtrip(self, d):
        """Basic roundtrip test."""
        ru_valid, _ = deltacamera.validity.fisheye_valid_r_max(d)

        # Handle infinite valid radius (well-behaved distortion)
        ru_max = ru_valid if np.isfinite(ru_valid) else 5.0

        # Sample points inside valid region
        r = np.random.uniform(0, ru_max * 0.8, 500).astype(np.float32)
        theta = np.random.uniform(-np.pi, np.pi, 500).astype(np.float32)
        points = np.stack([r * np.cos(theta), r * np.sin(theta)], axis=1)

        distorted = deltacamera.distortion.distort_points_fisheye(
            points, d, check_validity=False
        )
        recovered = deltacamera.distortion.undistort_points_fisheye(
            distorted, d, check_validity=False
        )

        # Fisheye undistortion has lower precision due to iterative algorithm
        # Allow ~3% relative error for strong distortions
        np.testing.assert_allclose(recovered, points, rtol=3e-2, atol=1e-4)

    def test_roundtrip_at_origin(self):
        """Point at origin should be unchanged."""
        d = FISHEYE_COEFFS[0]
        origin = np.array([[0, 0]], np.float32)

        distorted = deltacamera.distortion.distort_points_fisheye(
            origin, d, check_validity=False
        )
        np.testing.assert_allclose(distorted, origin, atol=1e-10)


# =============================================================================
# Validity checking tests
# =============================================================================

class TestDistortionValidityChecking:
    """Test that validity checking works correctly."""

    def test_valid_points_not_nan(self):
        """Points inside valid region should not produce NaN."""
        d = BROWN_CONRADY_COEFFS[0]
        camera = make_camera_with_distortion(d)
        # get_valid_distortion_region returns (undistorted_region, distorted_region)
        valid_region_ud, _ = deltacamera.validity.get_valid_distortion_region(
            camera, cartesian=True, n_vertices=128
        )
        points = sample_inside_polygon(valid_region_ud, n=200, margin=0.15)

        result = deltacamera.distortion.distort_points(
            points, d, check_validity=True
        )

        assert not np.any(np.isnan(result)), "Valid points should not produce NaN"

    def test_invalid_points_are_nan(self):
        """Points outside valid region should produce NaN."""
        d = BROWN_CONRADY_COEFFS[0]
        camera = make_camera_with_distortion(d)
        # get_valid_distortion_region returns (undistorted_region, distorted_region)
        valid_region_ud, _ = deltacamera.validity.get_valid_distortion_region(
            camera, cartesian=True, n_vertices=128
        )

        # Scale points to be outside valid region
        outside_points = valid_region_ud[::4] * 1.5  # Subsample and scale

        result = deltacamera.distortion.distort_points(
            outside_points, d, check_validity=True
        )

        nan_count = np.sum(np.isnan(result[:, 0]))
        # Most points should be NaN
        assert nan_count > len(outside_points) * 0.5, \
            f"Expected most points to be NaN, got {nan_count}/{len(outside_points)}"


# =============================================================================
# Zero distortion tests
# =============================================================================

class TestZeroDistortion:
    """Test behavior with zero distortion coefficients."""

    def test_zero_distortion_is_identity(self):
        """Zero distortion coefficients should not change points."""
        d = np.zeros(12, np.float32)
        points = np.random.randn(100, 2).astype(np.float32)

        distorted = deltacamera.distortion.distort_points(
            points, d, check_validity=False
        )

        np.testing.assert_allclose(distorted, points, atol=1e-10)

    def test_zero_fisheye_roundtrip(self):
        """Zero fisheye coefficients roundtrip should work."""
        d = np.zeros(4, np.float32)
        points = np.random.randn(100, 2).astype(np.float32) * 0.5

        distorted = deltacamera.distortion.distort_points_fisheye(
            points, d, check_validity=False
        )
        recovered = deltacamera.distortion.undistort_points_fisheye(
            distorted, d, check_validity=False
        )

        np.testing.assert_allclose(recovered, points, rtol=1e-4, atol=1e-6)


# =============================================================================
# Edge cases
# =============================================================================

class TestDistortionEdgeCases:
    """Test edge cases and boundary conditions."""

    def test_empty_points(self):
        """Empty point arrays should not crash."""
        d = BROWN_CONRADY_COEFFS[0]
        empty = np.zeros((0, 2), np.float32)

        result = deltacamera.distortion.distort_points(empty, d, check_validity=False)
        assert result.shape == (0, 2)

    def test_single_point(self):
        """Single point should work."""
        d = BROWN_CONRADY_COEFFS[0]
        point = np.array([[0.1, 0.1]], np.float32)

        distorted = deltacamera.distortion.distort_points(
            point, d, check_validity=False
        )
        recovered = deltacamera.distortion.undistort_points(
            distorted, d, check_validity=False
        )

        np.testing.assert_allclose(recovered, point, rtol=1e-4)

    def test_symmetric_distortion(self):
        """Symmetric points should distort symmetrically (no tangential)."""
        d = extend_distortion_coeffs(np.array([-0.2, 0.1, 0, 0, 0], np.float32))
        points = np.array([
            [0.3, 0.0],
            [-0.3, 0.0],
            [0.0, 0.3],
            [0.0, -0.3],
        ], np.float32)

        distorted = deltacamera.distortion.distort_points(
            points, d, check_validity=False
        )

        # All points at same radius should distort to same radius
        r_distorted = np.linalg.norm(distorted, axis=1)
        np.testing.assert_allclose(r_distorted, r_distorted[0], rtol=1e-6)


# =============================================================================
# OpenCV oracle tests
# =============================================================================

class TestDistortionAgainstOpenCV:
    """Verify distort_points matches OpenCV's projectPoints."""

    @staticmethod
    def opencv_distort(points_2d, distortion_coeffs):
        """Apply distortion using OpenCV as oracle.

        Args:
            points_2d: (N, 2) normalized undistorted points
            distortion_coeffs: distortion coefficients (5, 8, 12, or 14)

        Returns:
            (N, 2) distorted normalized points
        """
        import cv2

        # Convert 2D normalized coords to 3D at z=1
        points_3d = np.zeros((len(points_2d), 1, 3), dtype=np.float64)
        points_3d[:, 0, :2] = points_2d
        points_3d[:, 0, 2] = 1.0

        # Identity rotation/translation and camera matrix
        rvec = np.zeros(3, dtype=np.float64)
        tvec = np.zeros(3, dtype=np.float64)
        K = np.eye(3, dtype=np.float64)

        # OpenCV expects specific lengths, pad if needed
        d = np.asarray(distortion_coeffs, dtype=np.float64)

        result, _ = cv2.projectPoints(points_3d, rvec, tvec, K, d)
        return result[:, 0, :].astype(np.float32)

    @pytest.mark.parametrize("d", BROWN_CONRADY_5_COEFFS)
    def test_5_param_matches_opencv(self, d):
        """5-parameter distortion should match OpenCV exactly."""
        # Sample points in valid region (small radius to avoid extreme distortion)
        np.random.seed(42)
        points = (np.random.rand(200, 2).astype(np.float32) - 0.5) * 0.5

        # Our implementation
        d12 = extend_distortion_coeffs(d, 12)
        ours = deltacamera.distortion.distort_points(points, d12, check_validity=False)

        # OpenCV oracle
        opencv = self.opencv_distort(points, d)

        np.testing.assert_allclose(ours, opencv, rtol=1e-5, atol=1e-7)

    @pytest.mark.parametrize("d", BROWN_CONRADY_COEFFS)
    def test_12_param_matches_opencv(self, d):
        """12-parameter distortion should match OpenCV."""
        np.random.seed(42)
        points = (np.random.rand(200, 2).astype(np.float32) - 0.5) * 0.5

        ours = deltacamera.distortion.distort_points(points, d, check_validity=False)
        opencv = self.opencv_distort(points, d)

        np.testing.assert_allclose(ours, opencv, rtol=1e-5, atol=1e-7)

    def test_14_param_with_tilt_matches_opencv(self):
        """14-parameter distortion with tilt should match OpenCV."""
        np.random.seed(42)
        points = (np.random.rand(200, 2).astype(np.float32) - 0.5) * 0.4

        # 14-param with nonzero tilt
        d14 = np.zeros(14, dtype=np.float32)
        d14[0] = -0.15  # k1
        d14[1] = 0.05   # k2
        d14[2] = 0.0003  # p1
        d14[3] = -0.0002  # p2
        d14[4] = 0.01   # k3
        d14[8] = 0.0001  # s1
        d14[10] = 0.0001  # s3
        d14[12] = 0.02  # tau_x
        d14[13] = -0.015  # tau_y

        ours = deltacamera.distortion.distort_points(points, d14, check_validity=False)
        opencv = self.opencv_distort(points, d14)

        np.testing.assert_allclose(ours, opencv, rtol=1e-5, atol=1e-7)

    def test_origin_unchanged(self):
        """Origin should map to origin for any distortion."""
        d = BROWN_CONRADY_COEFFS[0]
        origin = np.array([[0.0, 0.0]], np.float32)

        ours = deltacamera.distortion.distort_points(origin, d, check_validity=False)
        opencv = self.opencv_distort(origin, d)

        np.testing.assert_allclose(ours, opencv, atol=1e-10)
        np.testing.assert_allclose(ours, origin, atol=1e-10)
