import numpy as np
import pytest

from gri_utils.constants import ER_M
from gri_utils.conversion import (
    get_ned,
    ned_to_xyz,
    ned_to_xyz_cov,
    translate_ned,
    xyz_to_ned,
    xyz_to_ned_cov,
)

# Type aliases for readability
XYZ = tuple[float, float, float]
NED = tuple[float, float, float]


@pytest.mark.parametrize(
    ("xyz_pos", "xyz_pos2", "expected"),
    [
        # in m - along each axis
        ((6378000, 0, 0), (6378100, 0, 0), (0, 0, -100)),
        ((6378000, 0, 0), (6378000, 100, 0), (0, 100, 0)),
        ((6378000, 0, 0), (6378000, 0, 100), (100, 0, 0)),
        # in km - verifies unitless behavior
        ((6378, 0, 0), (6378.1, 0, 0), (0, 0, -0.1)),
        ((6378, 0, 0), (6378, 0.1, 0), (0, 0.1, 0)),
        ((6378, 0, 0), (6378, 0, 0.1), (0.1, 0, 0)),
    ],
)
def test_get_ned(xyz_pos: XYZ, xyz_pos2: XYZ, expected: NED) -> None:
    """Test get_ned with array and sequence inputs."""
    # Array input
    ned = get_ned(np.array(xyz_pos), np.array(xyz_pos2))
    assert np.allclose(np.array(expected), ned), f"{ned}"

    # Sequence input
    ned_seq = get_ned(xyz_pos, xyz_pos2)
    assert ned_seq == pytest.approx(expected), f"{ned_seq}"


@pytest.mark.parametrize(
    ("xyz_pos", "ned_vec", "expected"),
    [
        # in m - along each axis
        ((6378000, 0, 0), (0, 0, -100), (6378100, 0, 0)),
        ((6378000, 0, 0), (100, 0, 0), (6378000, 0, 100)),
        ((6378000, 0, 0), (0, 100, 0), (6378000, 100, 0)),
        # in km - verifies unitless behavior
        ((6378, 0, 0), (0, 0, -0.1), (6378.1, 0, 0)),
        ((6378, 0, 0), (0.1, 0, 0), (6378, 0, 0.1)),
        ((6378, 0, 0), (0, 0.1, 0), (6378, 0.1, 0)),
    ],
)
def test_translate_ned(xyz_pos: XYZ, ned_vec: NED, expected: XYZ) -> None:
    """Test translate_ned with array and sequence inputs."""
    # Array input
    xyz = translate_ned(np.array(xyz_pos), np.array(ned_vec))
    assert np.allclose(np.array(expected), xyz), f"{xyz}"

    # Sequence input
    xyz_seq = translate_ned(xyz_pos, ned_vec)
    assert xyz_seq == pytest.approx(expected), f"{xyz_seq}"


@pytest.mark.parametrize(
    ("xyz_pos", "xyz_vec", "expected"),
    [
        # in m - along each axis
        ((6378000, 0, 0), (0, 100, 0), (0, 100, 0)),
        ((6378000, 0, 0), (0, 0, 100), (100, 0, 0)),
        ((6378000, 0, 0), (100, 0, 0), (0, 0, -100)),
        # in km - verifies unitless behavior
        ((6378, 0, 0), (0, 0.1, 0), (0, 0.1, 0)),
        ((6378, 0, 0), (0, 0, 0.1), (0.1, 0, 0)),
        ((6378, 0, 0), (0.1, 0, 0), (0, 0, -0.1)),
    ],
)
def test_xyz_to_ned(xyz_pos: XYZ, xyz_vec: XYZ, expected: NED) -> None:
    """Test xyz_to_ned with array and sequence inputs."""
    # Array input
    ned = xyz_to_ned(np.array(xyz_pos), np.array(xyz_vec))
    assert np.allclose(np.array(expected), ned), f"{ned}"

    # Sequence input
    ned_seq = xyz_to_ned(xyz_pos, xyz_vec)
    assert ned_seq == pytest.approx(expected), f"{ned_seq}"


def test_xyz_to_ned_cov() -> None:
    """Test xyz_to_ned_cov with 3x3 matrix input (covariance transformation)."""
    # XYZ diagonal matrix: X=100, Y=25, Z=1
    xyz = np.array(((100, 0, 0), (0, 25, 0), (0, 0, 1)))
    # At (ER_M, 0, 0): N=Z, E=Y, D=-X -> NED diagonal: N=1, E=25, D=100
    ned = np.array(((1, 0, 0), (0, 25, 0), (0, 0, 100)))
    ans = xyz_to_ned_cov((ER_M, 0, 0), xyz)
    assert np.allclose(ned, ans)


@pytest.mark.parametrize(
    ("xyz_pos", "ned_vec", "expected"),
    [
        # in m - along each axis
        ((6378000, 0, 0), (0, 100, 0), (0, 100, 0)),
        ((6378000, 0, 0), (100, 0, 0), (0, 0, 100)),
        ((6378000, 0, 0), (0, 0, -100), (100, 0, 0)),
        # in km - verifies unitless behavior
        ((6378, 0, 0), (0, 0.1, 0), (0, 0.1, 0)),
        ((6378, 0, 0), (0.1, 0, 0), (0, 0, 0.1)),
        ((6378, 0, 0), (0, 0, -0.1), (0.1, 0, 0)),
    ],
)
def test_ned_to_xyz(xyz_pos: XYZ, ned_vec: NED, expected: XYZ) -> None:
    """Test ned_to_xyz with array and sequence inputs."""
    # Array input
    xyz = ned_to_xyz(np.array(xyz_pos), np.array(ned_vec))
    assert np.allclose(np.array(expected), xyz), f"{xyz}"

    # Sequence input
    xyz_seq = ned_to_xyz(xyz_pos, ned_vec)
    assert xyz_seq == pytest.approx(expected), f"{xyz_seq}"


def test_ned_to_xyz_cov() -> None:
    """Test ned_to_xyz_cov with 3x3 matrix input (covariance transformation)."""
    # NED diagonal matrix: N=1, E=25, D=100
    ned = np.array(((1, 0, 0), (0, 25, 0), (0, 0, 100)))
    # At (ER_M, 0, 0): X=-D, Y=E, Z=N -> XYZ diagonal: X=100, Y=25, Z=1
    xyz = np.array(((100, 0, 0), (0, 25, 0), (0, 0, 1)))
    ans = ned_to_xyz_cov((ER_M, 0, 0), ned)
    assert np.allclose(xyz, ans)


# Vectorized tests


def test_get_ned_vectorized() -> None:
    """Test get_ned with multiple target positions."""
    # Reference position on equator at prime meridian
    xyz_ref: XYZ = (6378000, 0, 0)

    # Multiple target positions: +100m in each direction
    xyz_targets = np.array(
        [
            (6378100, 0, 0),  # +100m radially out (Down = -100)
            (6378000, 100, 0),  # +100m in Y (East)
            (6378000, 0, 100),  # +100m in Z (North)
        ],
    )
    expected = np.array(
        [
            (0, 0, -100),  # Down (negative)
            (0, 100, 0),  # East
            (100, 0, 0),  # North
        ],
    )

    ned_vecs = get_ned(xyz_ref, xyz_targets)

    assert ned_vecs.shape == (3, 3)
    assert np.allclose(ned_vecs, expected)


def test_translate_ned_vectorized() -> None:
    """Test translate_ned with multiple NED vectors."""
    xyz_ref: XYZ = (6378000, 0, 0)

    # Multiple NED vectors
    ned_vecs = np.array(
        [
            (0, 0, -100),  # Down -100 -> radially out
            (0, 100, 0),  # East
            (100, 0, 0),  # North
        ],
    )
    expected = np.array(
        [
            (6378100, 0, 0),
            (6378000, 100, 0),
            (6378000, 0, 100),
        ],
    )

    xyzs = translate_ned(xyz_ref, ned_vecs)

    assert xyzs.shape == (3, 3)
    assert np.allclose(xyzs, expected)


def test_xyz_to_ned_vectorized() -> None:
    """Test xyz_to_ned with multiple XYZ vectors."""
    xyz_ref: XYZ = (6378000, 0, 0)

    # Use 4 vectors to avoid (3,3) matrix interpretation
    xyz_vecs = np.array(
        [
            (100, 0, 0),  # Radial out -> Down -100
            (0, 100, 0),  # Y -> East
            (0, 0, 100),  # Z -> North
            (50, 50, 50),  # Mixed
        ],
    )
    expected = np.array(
        [
            (0, 0, -100),
            (0, 100, 0),
            (100, 0, 0),
            (50, 50, -50),  # N=Z, E=Y, D=-X
        ],
    )

    ned_vecs = xyz_to_ned(xyz_ref, xyz_vecs)

    assert ned_vecs.shape == (4, 3)
    assert np.allclose(ned_vecs, expected)


def test_ned_to_xyz_vectorized() -> None:
    """Test ned_to_xyz with multiple NED vectors."""
    xyz_ref: XYZ = (6378000, 0, 0)

    # Use 4 vectors to avoid (3,3) matrix interpretation
    ned_vecs = np.array(
        [
            (100, 0, 0),  # North -> Z
            (0, 100, 0),  # East -> Y
            (0, 0, -100),  # Down -100 -> X (radial out)
            (50, 50, -50),  # Mixed
        ],
    )
    expected = np.array(
        [
            (0, 0, 100),
            (0, 100, 0),
            (100, 0, 0),
            (50, 50, 50),  # X=-D, Y=E, Z=N
        ],
    )

    xyz_vecs = ned_to_xyz(xyz_ref, ned_vecs)

    assert xyz_vecs.shape == (4, 3)
    assert np.allclose(xyz_vecs, expected)


def test_ned_xyz_vectorized_round_trip() -> None:
    """Test round-trip NED<->XYZ conversion with batched inputs."""
    xyz_ref: XYZ = (6378000, 0, 0)

    ned_vecs = np.array(
        [
            (50, 30, -20),
            (100, 0, 0),
            (0, 100, 0),
            (0, 0, -100),
            (-50, -30, 20),
        ],
    )

    xyz_vecs = ned_to_xyz(xyz_ref, ned_vecs)
    ned_back = xyz_to_ned(xyz_ref, xyz_vecs)

    assert ned_back.shape == ned_vecs.shape
    assert np.allclose(ned_back, ned_vecs)
