import numpy as np
import pytest
from numpy.testing import assert_allclose

from iactsim.optics._cpu_transforms import (
    pointing_dir,
    local_to_pointing_rotation,
    local_to_telescope_rotation
)

class TestPointingDirection:
    """Tests for `pointing_dir` function."""

    @pytest.mark.parametrize("alt, az, expected_vec", [
        (0, 0,    [1, 0, 0]),  # North -> Local X
        (0, 90,   [0, -1, 0]), # East -> Local -Y
        (0, 180,  [-1, 0, 0]), # South -> Local -X
        (0, 270,  [0, 1, 0]),  # West -> Local +Y
        (90, 0,   [0, 0, 1]),  # Zenith -> Local Z
    ])

    def test_cardinal_directions(self, alt, az, expected_vec):
        """Verify standard cardinal points match the North-West-Up logic."""
        vec = pointing_dir(alt, az)
        assert_allclose(vec, expected_vec, atol=1e-7)

    def test_unit_norm(self):
        """Ensure the output is always a unit vector."""
        for alt in [0, 30, 45, 90]:
            for az in [0, 45, 90, 180, 270]:
                vec = pointing_dir(alt, az)
                assert_allclose(np.linalg.norm(vec), 1.0, atol=1e-7)


class TestRotationMatrices:
    """Tests for rotation matrix generation functions."""

    def test_default_is_identity(self):
        """
        Verify that if no custom_rotation is provided, the function 
        behaves exactly like local_to_pointing_rotation.
        """
        alt, az = 45, 45
        R_default = local_to_telescope_rotation(alt, az)
        R_pointing = local_to_pointing_rotation(alt, az)
        
        assert_allclose(R_default, R_pointing, atol=1e-15,
                        err_msg="Default behavior should be identity (match pointing rotation)")

    def test_astri_rotation(self):
        """
        Verify the specific ASTRI refrence system:
        Target frame when pointing North: 
          - X = West
          - Y = Up
          - Z = North
        """
        # The ASTRI custom transformation matrix pointing North
        astri_matrix = np.array([
            [ 0, 1, 0], # New X (West)
            [-1, 0, 0], # New Y (Up)
            [ 0, 0, 1]  # New Z (North)
        ])
        
        # Test at Reference Position (North, Horizon)
        alt, az = 0, 0
        R = local_to_telescope_rotation(alt, az, custom_rotation=astri_matrix)

        # Local North (1, 0, 0) should become telescope Z (0, 0, 1)
        assert_allclose(R @ [1, 0, 0], [0, 0, 1], atol=1e-7)

        # Local West (0, 1, 0) should become telescope X (1, 0, 0)
        assert_allclose(R @ [0, 1, 0], [1, 0, 0], atol=1e-7)

        # Local Up (0, 0, 1) should become telescope Y (0, 1, 0)
        assert_allclose(R @ [0, 0, 1], [0, 1, 0], atol=1e-7)

    def test_orthogonality(self):
        """Rotation matrices must be orthogonal."""
        alt, az = 30, 60
        R = local_to_telescope_rotation(alt, az) # Uses default identity
        
        identity = np.eye(3)
        assert_allclose(R @ R.T, identity, atol=1e-7)
        assert_allclose(np.linalg.det(R), 1.0, atol=1e-7)

    def test_pointing_frame_orientation_default(self):
        """
        Verify the orientation of the default output (pointing frame).
        
        At Alt=0, Az=0 (pointing North):
        - Z axis points North (local X)
        - X axis points Down (local -Z)
        - Y axis points West (local Y)
        """
        R = local_to_telescope_rotation(0, 0) # Default
        
        # Local North (1, 0, 0) -> pointing Z (0, 0, 1)
        assert_allclose(R @ [1, 0, 0], [0, 0, 1], atol=1e-7)

        # Local down (0, 0, -1) -> pointing X (1, 0, 0)
        assert_allclose(R @ [0, 0, -1], [1, 0, 0], atol=1e-7)

        # Local West (0, 1, 0) -> pointing Y (0, 1, 0)
        assert_allclose(R @ [0, 1, 0], [0, 1, 0], atol=1e-7)