import numpy as np
import pytest

from gri_utils.constants import OMEGA_EARTH
from gri_utils.conversion import ecef_to_teme, teme_to_ecef
from gri_utils.conversion.julian import JD_J2000, JD_UNIX_EPOCH
from gri_utils.conversion.xyz_teme_vel import ecef_to_teme_vel, teme_to_ecef_vel

# Earth radius at equator (meters)
ER_M = 6378137.0

# Expected surface velocity at equator due to Earth rotation
EQUATOR_SURFACE_SPEED = OMEGA_EARTH * ER_M

# Unix timestamp for J2000.0 epoch
UNIX_J2000 = (JD_J2000 - JD_UNIX_EPOCH) * 86400.0


def test_stationary_ecef_has_velocity_in_teme() -> None:
    """A point stationary in ECEF should have velocity in TEME."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 0.0, 0.0])

    vel_teme = ecef_to_teme_vel(pos_ecef, vel_ecef, 0.0)

    # Should have ~465 m/s velocity in TEME
    speed_teme = np.linalg.norm(vel_teme)
    assert np.isclose(speed_teme, EQUATOR_SURFACE_SPEED, rtol=1e-6)


def test_north_pole_stationary_stays_stationary() -> None:
    """A point stationary at north pole has no velocity in TEME."""
    pos_ecef = np.array([0.0, 0.0, 6356752.3142])
    vel_ecef = np.array([0.0, 0.0, 0.0])

    vel_teme = ecef_to_teme_vel(pos_ecef, vel_ecef, 0.0)

    assert np.allclose(vel_teme, [0.0, 0.0, 0.0], atol=1e-10)


def test_z_velocity_unchanged() -> None:
    """Z component of velocity should be unchanged."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 0.0, 100.0])

    vel_teme = ecef_to_teme_vel(pos_ecef, vel_ecef, 0.0)

    assert np.isclose(vel_teme[2], 100.0)


def test_roundtrip_velocity() -> None:
    """Test ECEF -> TEME -> ECEF velocity roundtrip."""
    pos_ecef = np.array([1e7, 2e6, 3e6])
    vel_ecef = np.array([100.0, 200.0, 300.0])

    for unix_s in [0.0, 86400.0, UNIX_J2000]:
        pos_teme = ecef_to_teme(pos_ecef, unix_s)
        vel_teme = ecef_to_teme_vel(pos_ecef, vel_ecef, unix_s)
        vel_ecef_back = teme_to_ecef_vel(pos_teme, vel_teme, unix_s)

        assert np.allclose(vel_ecef_back, vel_ecef, rtol=1e-12), (
            f"Round-trip failed at t={unix_s}"
        )


def test_roundtrip_teme_to_ecef_velocity() -> None:
    """Test TEME -> ECEF -> TEME velocity roundtrip."""
    pos_teme = np.array([1e7, 2e6, 3e6])
    vel_teme = np.array([100.0, 200.0, 300.0])
    unix_s = UNIX_J2000

    pos_ecef = teme_to_ecef(pos_teme, unix_s)
    vel_ecef = teme_to_ecef_vel(pos_teme, vel_teme, unix_s)
    vel_teme_back = ecef_to_teme_vel(pos_ecef, vel_ecef, unix_s)

    assert np.allclose(vel_teme_back, vel_teme, rtol=1e-12)


def test_sgp4_typical_velocity() -> None:
    """Test with typical SGP4 satellite state (LEO satellite)."""
    # Typical LEO satellite: ~7000 km altitude, ~7.5 km/s orbital velocity
    r = 7000e3  # meters
    v_orbital = 7500.0  # m/s

    # Satellite at x-axis in TEME, moving in +y direction
    pos_teme = np.array([r, 0.0, 0.0])
    vel_teme = np.array([0.0, v_orbital, 0.0])

    vel_ecef = teme_to_ecef_vel(pos_teme, vel_teme, 0.0)

    # The key test is roundtrip consistency
    pos_ecef = teme_to_ecef(pos_teme, 0.0)
    vel_teme_back = ecef_to_teme_vel(pos_ecef, vel_ecef, 0.0)
    assert np.allclose(vel_teme_back, vel_teme, rtol=1e-12)

    # ECEF speed should differ from TEME speed due to Earth rotation
    speed_ecef = np.linalg.norm(vel_ecef)
    speed_teme = np.linalg.norm(vel_teme)
    assert not np.isclose(speed_ecef, speed_teme, rtol=0.01), (
        "ECEF and TEME speeds should differ"
    )


def test_vectorized_single_time() -> None:
    """Test with batched (N, 3) position/velocity and single time."""
    positions = np.array(
        [
            [ER_M, 0.0, 0.0],
            [0.0, ER_M, 0.0],
            [0.0, 0.0, 6356752.3142],
        ],
    )
    velocities = np.zeros((3, 3))

    vel_teme = ecef_to_teme_vel(positions, velocities, 0.0)

    assert vel_teme.shape == (3, 3)

    for i in range(3):
        vel_single = ecef_to_teme_vel(positions[i], velocities[i], 0.0)
        assert np.allclose(vel_teme[i], vel_single)


def test_vectorized_multiple_times() -> None:
    """Test with batched inputs and array of times."""
    positions = np.array(
        [
            [ER_M, 0.0, 0.0],
            [ER_M, 0.0, 0.0],
            [ER_M, 0.0, 0.0],
        ],
    )
    velocities = np.array(
        [
            [100.0, 0.0, 0.0],
            [100.0, 0.0, 0.0],
            [100.0, 0.0, 0.0],
        ],
    )
    times = np.array([0.0, 3600.0, 7200.0])

    vel_teme = ecef_to_teme_vel(positions, velocities, times)

    assert vel_teme.shape == (3, 3)

    # Results should differ due to different times
    assert not np.allclose(vel_teme[0], vel_teme[1])


def test_vectorized_roundtrip() -> None:
    """Test vectorized round-trip transformation."""
    positions_ecef = np.array(
        [
            [1e7, 2e6, 3e6],
            [4e6, 5e6, 6e6],
            [-1e6, -2e6, 1e6],
        ],
    )
    velocities_ecef = np.array(
        [
            [100.0, 200.0, 300.0],
            [-50.0, 100.0, -200.0],
            [0.0, 0.0, 500.0],
        ],
    )
    times = np.array([0.0, 86400.0, UNIX_J2000])

    positions_teme = ecef_to_teme(positions_ecef, times)
    vel_teme = ecef_to_teme_vel(positions_ecef, velocities_ecef, times)
    vel_ecef_back = teme_to_ecef_vel(positions_teme, vel_teme, times)

    assert np.allclose(vel_ecef_back, velocities_ecef, rtol=1e-12)


def test_sequence_input() -> None:
    """Test that sequence inputs work correctly."""
    pos_tuple = (ER_M, 0.0, 0.0)
    vel_list = [100.0, 200.0, 300.0]

    vel_teme = ecef_to_teme_vel(pos_tuple, vel_list, 0.0)

    assert isinstance(vel_teme, np.ndarray)
    assert vel_teme.shape == (3,)


@pytest.mark.parametrize(
    "unix_s",
    [0.0, 86400.0, -86400.0, UNIX_J2000, 1e9],
)
def test_transport_velocity_direction(unix_s: float) -> None:
    """Test that transport velocity is in correct direction."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 0.0, 0.0])

    vel_teme = ecef_to_teme_vel(pos_ecef, vel_ecef, unix_s)
    pos_teme = ecef_to_teme(pos_ecef, unix_s)

    # Velocity should be perpendicular to position in X-Y plane
    dot_xy = pos_teme[0] * vel_teme[0] + pos_teme[1] * vel_teme[1]
    assert np.isclose(dot_xy, 0.0, atol=1e-6)

    # Cross product should be positive Z (counterclockwise)
    cross_z = pos_teme[0] * vel_teme[1] - pos_teme[1] * vel_teme[0]
    assert cross_z > 0


def test_broadcast_single_position_multiple_velocities() -> None:
    """Test broadcasting single position with multiple velocities."""
    pos = np.array([ER_M, 0.0, 0.0])
    velocities = np.array(
        [
            [100.0, 0.0, 0.0],
            [0.0, 100.0, 0.0],
            [0.0, 0.0, 100.0],
        ],
    )

    vel_teme = ecef_to_teme_vel(pos, velocities, 0.0)

    assert vel_teme.shape == (3, 3)

    for i in range(3):
        vel_single = ecef_to_teme_vel(pos, velocities[i], 0.0)
        assert np.allclose(vel_teme[i], vel_single)
