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_acc import ecef_to_eci_acc, eci_to_ecef_acc
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

# Centrifugal acceleration at equator surface
# a = ω² * r ≈ 0.0339 m/s²
CENTRIFUGAL_EQUATOR = OMEGA_EARTH**2 * ER_M

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


def test_stationary_ecef_has_acceleration_in_eci() -> None:
    """A point at rest in ECEF experiences acceleration in ECI.

    A stationary point in ECEF is actually moving in a circle in ECI due to
    Earth's rotation. This circular motion requires centripetal acceleration.
    Additionally, the transport velocity creates a Coriolis term.

    The total acceleration includes:
    - Centrifugal: ω x (ω x r) = -ω²r (pointing inward, toward axis)
    - Coriolis: 2ω x v_eci (due to transport velocity)
    """
    pos_ecef = np.array([ER_M, 0.0, 0.0])  # Equator
    vel_ecef = np.array([0.0, 0.0, 0.0])  # Stationary
    acc_ecef = np.array([0.0, 0.0, 0.0])  # No acceleration in ECEF

    acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, 0.0)

    # Should have non-zero acceleration in X-Y plane
    acc_magnitude = np.linalg.norm(acc_eci[:2])
    assert acc_magnitude > 0, "Should have acceleration in ECI"

    # The acceleration includes both centrifugal and Coriolis terms
    # For a point stationary in ECEF:
    # - Centrifugal: -ω²r ≈ -0.034 m/s² (inward)
    # - Coriolis: 2ω x v_eci where v_eci is the transport velocity
    #   v_eci ≈ 465 m/s, so 2ω x v ≈ 0.068 m/s² (perpendicular to v)
    # Total is on order of 0.1 m/s², which matches observed value

    # Z component should be zero
    assert np.isclose(acc_eci[2], 0.0, atol=1e-10)


def test_north_pole_no_centrifugal() -> None:
    """A point at north pole has no centrifugal acceleration."""
    pos_ecef = np.array([0.0, 0.0, 6356752.3142])  # North pole
    vel_ecef = np.array([0.0, 0.0, 0.0])
    acc_ecef = np.array([0.0, 0.0, 0.0])

    acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, 0.0)

    # On rotation axis, centrifugal and Coriolis terms are zero
    assert np.allclose(acc_eci, [0.0, 0.0, 0.0], atol=1e-10)


def test_z_acceleration_unchanged() -> None:
    """Z component of acceleration passes through (modified only by rotation)."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 0.0, 0.0])
    acc_ecef = np.array([0.0, 0.0, 10.0])  # Pure Z acceleration

    acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, 0.0)

    # Z component unchanged (Coriolis and centrifugal don't affect Z)
    assert np.isclose(acc_eci[2], 10.0)


def test_coriolis_effect() -> None:
    """Test that Coriolis acceleration is present for moving objects."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 1000.0, 0.0])  # Moving east at 1 km/s
    acc_ecef = np.array([0.0, 0.0, 0.0])

    acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, 0.0)

    # Coriolis: 2ω x v, with v in ECI
    # For v in +Y direction in ECI, Coriolis is in +X direction
    # Combined with centrifugal, should have significant X-Y acceleration

    # The magnitude should be larger than just centrifugal due to Coriolis
    acc_xy = np.sqrt(acc_eci[0] ** 2 + acc_eci[1] ** 2)
    assert acc_xy > CENTRIFUGAL_EQUATOR


def test_roundtrip_acceleration() -> None:
    """Test ECEF -> ECI -> ECEF acceleration roundtrip."""
    pos_ecef = np.array([1e7, 2e6, 3e6])
    vel_ecef = np.array([100.0, 200.0, 300.0])
    acc_ecef = np.array([1.0, 2.0, 3.0])

    for unix_s in [0.0, 86400.0, UNIX_J2000]:
        pos_eci = ecef_to_eci(pos_ecef, unix_s)
        vel_eci = ecef_to_eci_vel(pos_ecef, vel_ecef, unix_s)
        acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, unix_s)

        acc_ecef_back = eci_to_ecef_acc(pos_eci, vel_eci, acc_eci, unix_s)

        assert np.allclose(acc_ecef_back, acc_ecef, rtol=1e-10), (
            f"Round-trip failed at t={unix_s}"
        )


def test_roundtrip_eci_to_ecef_acceleration() -> None:
    """Test ECI -> ECEF -> ECI acceleration roundtrip."""
    pos_eci = np.array([1e7, 2e6, 3e6])
    vel_eci = np.array([100.0, 200.0, 300.0])
    acc_eci = np.array([1.0, 2.0, 3.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)
    acc_ecef = eci_to_ecef_acc(pos_eci, vel_eci, acc_eci, unix_s)

    acc_eci_back = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, unix_s)

    assert np.allclose(acc_eci_back, acc_eci, rtol=1e-10)


def test_vectorized_single_time() -> None:
    """Test with batched (N, 3) inputs 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))
    accelerations = np.zeros((3, 3))

    acc_eci = ecef_to_eci_acc(positions, velocities, accelerations, 0.0)

    assert acc_eci.shape == (3, 3)

    # Verify each result individually
    for i in range(3):
        acc_single = ecef_to_eci_acc(
            positions[i],
            velocities[i],
            accelerations[i],
            0.0,
        )
        assert np.allclose(acc_eci[i], acc_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],
        ],
    )
    accelerations = np.array(
        [
            [1.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
        ],
    )
    times = np.array([0.0, 3600.0, 7200.0])

    acc_eci = ecef_to_eci_acc(positions, velocities, accelerations, times)

    assert acc_eci.shape == (3, 3)

    # Results should differ due to different times
    assert not np.allclose(acc_eci[0], acc_eci[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],
        ],
    )
    accelerations_ecef = np.array(
        [
            [1.0, 2.0, 3.0],
            [-0.5, 1.0, -2.0],
            [0.0, 0.0, 5.0],
        ],
    )
    times = np.array([0.0, 86400.0, UNIX_J2000])

    positions_eci = ecef_to_eci(positions_ecef, times)
    velocities_eci = ecef_to_eci_vel(positions_ecef, velocities_ecef, times)
    acc_eci = ecef_to_eci_acc(
        positions_ecef,
        velocities_ecef,
        accelerations_ecef,
        times,
    )

    acc_ecef_back = eci_to_ecef_acc(positions_eci, velocities_eci, acc_eci, times)

    assert np.allclose(acc_ecef_back, accelerations_ecef, rtol=1e-10)


def test_sequence_input() -> None:
    """Test that sequence inputs (lists, tuples) work correctly."""
    pos = (ER_M, 0.0, 0.0)
    vel = [0.0, 0.0, 0.0]
    acc = [1.0, 2.0, 3.0]

    acc_eci = ecef_to_eci_acc(pos, vel, acc, 0.0)

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


@pytest.mark.parametrize(
    "unix_s",
    [0.0, 86400.0, -86400.0, UNIX_J2000, 1e9],
)
def test_centrifugal_direction(unix_s: float) -> None:
    """Test that centrifugal acceleration points radially outward."""
    pos_ecef = np.array([ER_M, 0.0, 0.0])
    vel_ecef = np.array([0.0, 0.0, 0.0])
    acc_ecef = np.array([0.0, 0.0, 0.0])

    acc_eci = ecef_to_eci_acc(pos_ecef, vel_ecef, acc_ecef, unix_s)
    pos_eci = ecef_to_eci(pos_ecef, unix_s)

    # Centrifugal should point radially outward (same direction as position)
    # in the X-Y plane
    pos_xy = pos_eci[:2]
    acc_xy = acc_eci[:2]

    # Dot product should be negative (centrifugal points outward, away from axis)
    # Actually, ω x (ω x r) = -ω²r for r perpendicular to ω
    # This is -ω² times the position, so it points INWARD (toward axis)
    # Wait, that's wrong - centrifugal points OUTWARD
    # Let me reconsider: ω x (ω x r) where ω = [0,0,ω], r = [x,y,z]
    # ω x r = [-ωy, ωx, 0]
    # ω x (ω x r) = [0,0,ω] x [-ωy, ωx, 0] = [-ω²x, -ω²y, 0]
    # This IS pointing inward (opposite to r). But in the rotating frame,
    # centrifugal points outward. The formula gives the acceleration in ECI
    # that accounts for the rotation, so it's the inward centripetal acceleration.

    # For a stationary point in ECEF, the ECI sees it moving in a circle,
    # so it must have centripetal acceleration (inward).
    dot_xy = np.dot(pos_xy, acc_xy)
    assert dot_xy < 0, "Centripetal acceleration should point inward"


def test_broadcast_positions() -> None:
    """Test broadcasting different position shapes with velocity/acceleration."""
    # Single position, multiple velocities/accelerations
    pos = np.array([ER_M, 0.0, 0.0])
    velocities = np.array(
        [
            [0.0, 0.0, 0.0],
            [100.0, 0.0, 0.0],
            [0.0, 100.0, 0.0],
        ],
    )
    accelerations = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
        ],
    )

    acc_eci = ecef_to_eci_acc(pos, velocities, accelerations, 0.0)

    assert acc_eci.shape == (3, 3)

    for i in range(3):
        acc_single = ecef_to_eci_acc(pos, velocities[i], accelerations[i], 0.0)
        assert np.allclose(acc_eci[i], acc_single)
