import numpy as np
import pytest

from gri_utils.constants import OMEGA_EARTH
from gri_utils.conversion import ecef_to_eci, eci_to_ecef
from gri_utils.conversion.julian import JD_J2000, JD_UNIX_EPOCH
from gri_utils.conversion.xyz_eci_vel import ecef_to_eci_vel, eci_to_ecef_vel

# Earth radius at equator (meters)
ER_M = 6378137.0

# Expected surface velocity at equator due to Earth rotation
# v = ω x r = ω * r_equator ≈ 465.1 m/s
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_eci() -> None:
    """A point stationary in ECEF should have velocity in ECI due to Earth rotation."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])  # Equator, prime meridian
    vel_ecef = np.array([0.0, 0.0, 0.0])  # Stationary in ECEF

    vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, 0.0)

    # Should have ~465 m/s velocity in ECI (Earth surface velocity)
    speed_eci = np.linalg.norm(vel_eci)
    assert np.isclose(speed_eci, EQUATOR_SURFACE_SPEED, rtol=1e-6), (
        f"Expected ~{EQUATOR_SURFACE_SPEED:.1f} m/s, got {speed_eci:.1f} m/s"
    )

    # The key check is the speed magnitude (perpendicularity is frame-dependent)


def test_north_pole_stationary_stays_stationary() -> None:
    """A point stationary at north pole in ECEF should be stationary in ECI."""
    pos_ecef = np.array([0.0, 0.0, 6356752.3142])  # North pole
    vel_ecef = np.array([0.0, 0.0, 0.0])

    vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, 0.0)

    # North pole is on rotation axis, so no transport velocity
    assert np.allclose(vel_eci, [0.0, 0.0, 0.0], atol=1e-10)


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

    vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, 0.0)

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


def test_roundtrip_velocity() -> None:
    """Test ECEF -> ECI -> 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]:
        # Convert position to ECI for the reverse transform
        pos_eci = ecef_to_eci(pos_ecef, unix_s)
        vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, unix_s)
        vel_ecef_back = eci_to_ecef_vel(pos_eci, vel_eci, unix_s)

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


def test_roundtrip_eci_to_ecef_velocity() -> None:
    """Test ECI -> ECEF -> ECI velocity roundtrip."""
    pos_eci = np.array([1e7, 2e6, 3e6])
    vel_eci = np.array([100.0, 200.0, 300.0])
    unix_s = UNIX_J2000

    pos_ecef = eci_to_ecef(pos_eci, unix_s)
    vel_ecef = eci_to_ecef_vel(pos_eci, vel_eci, unix_s)
    vel_eci_back = ecef_to_eci_vel(pos_ecef, vel_ecef, unix_s)

    assert np.allclose(vel_eci_back, vel_eci, rtol=1e-12)


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],  # North pole
        ],
    )
    velocities = np.array(
        [
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 0.0],
        ],
    )

    vel_eci = ecef_to_eci_vel(positions, velocities, 0.0)

    assert vel_eci.shape == (3, 3)

    # Verify each result individually
    for i in range(3):
        vel_single = ecef_to_eci_vel(positions[i], velocities[i], 0.0)
        assert np.allclose(vel_eci[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_eci = ecef_to_eci_vel(positions, velocities, times)

    assert vel_eci.shape == (3, 3)

    # Each result should differ due to different times
    assert not np.allclose(vel_eci[0], vel_eci[1])
    assert not np.allclose(vel_eci[1], vel_eci[2])


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_eci = ecef_to_eci(positions_ecef, times)
    vel_eci = ecef_to_eci_vel(positions_ecef, velocities_ecef, times)
    vel_ecef_back = eci_to_ecef_vel(positions_eci, vel_eci, times)

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


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_eci = ecef_to_eci_vel(pos, velocities, 0.0)

    assert vel_eci.shape == (3, 3)

    # Each velocity vector should transform differently
    # but the transport term (ω x r) is the same for all
    for i in range(3):
        vel_single = ecef_to_eci_vel(pos, velocities[i], 0.0)
        assert np.allclose(vel_eci[i], vel_single)


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

    vel_eci = ecef_to_eci_vel(pos_tuple, vel_list, 0.0)

    assert isinstance(vel_eci, np.ndarray)
    assert vel_eci.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 (eastward)."""
    # Point on equator, stationary in ECEF
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 0.0, 0.0])

    vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, unix_s)

    # Convert position to ECI to check perpendicularity
    pos_eci = ecef_to_eci(pos_ecef, unix_s)

    # Velocity should be perpendicular to position in X-Y plane
    # (tangent to circular motion)
    dot_xy = pos_eci[0] * vel_eci[0] + pos_eci[1] * vel_eci[1]
    assert np.isclose(dot_xy, 0.0, atol=1e-6), (
        f"Velocity should be perpendicular to position: dot={dot_xy}"
    )

    # Cross product pos x vel should point in +Z direction (counterclockwise rotation)
    cross_z = pos_eci[0] * vel_eci[1] - pos_eci[1] * vel_eci[0]
    assert cross_z > 0, "Transport velocity should be eastward (counterclockwise)"


def test_high_velocity_satellite() -> None:
    """Test with typical LEO satellite velocity (~7.5 km/s)."""
    # Satellite at ~400 km altitude, moving tangentially
    r = ER_M + 400e3
    pos_ecef = np.array([r, 0.0, 0.0])

    # Orbital velocity for circular orbit: v = sqrt(GM/r) ≈ 7.67 km/s
    # In ECEF, subtract Earth rotation: ~7.67 - 0.47 ≈ 7.2 km/s
    v_orbital = 7200.0  # m/s, approximate
    vel_ecef = np.array([0.0, v_orbital, 0.0])  # Moving in +Y (east)

    vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, 0.0)

    # ECI velocity should be higher than ECEF (add Earth rotation)
    speed_eci = np.linalg.norm(vel_eci)
    speed_ecef = np.linalg.norm(vel_ecef)

    # Transport velocity at this altitude: ω * r ≈ 494 m/s
    transport_speed = OMEGA_EARTH * r

    # ECI speed should be greater than ECEF speed for prograde orbit
    # (Earth rotation adds to orbital velocity)
    assert speed_eci > speed_ecef, "ECI speed should be greater than ECEF speed"

    # The difference in speed should be on order of transport velocity
    speed_diff = speed_eci - speed_ecef
    assert speed_diff > 0, "Speed should increase in ECI"
    assert speed_diff < 2 * transport_speed, "Speed increase should be bounded"

    # Most importantly, verify roundtrip
    pos_eci = ecef_to_eci(pos_ecef, 0.0)
    vel_ecef_back = eci_to_ecef_vel(pos_eci, vel_eci, 0.0)
    assert np.allclose(vel_ecef_back, vel_ecef, rtol=1e-12)
