"""Tests for elements_to_state (Keplerian orbit propagation)."""

import math

import numpy as np
import pytest

from gri_utils.constants import GM_EARTH
from gri_utils.orbit import OrbitalElements, elements_to_state


def _circular_velocity(radius_m: float) -> float:
    """Compute circular orbit velocity magnitude."""
    return math.sqrt(GM_EARTH / radius_m)


# Reference orbit: circular equatorial at 7000 km radius
_CIRCULAR_ELEMENTS = OrbitalElements(
    semi_major_axis_m=7_000_000,
    eccentricity=0.0,
    inclination_rad=0.0,
    raan_rad=0.0,
    arg_periapsis_rad=0.0,
    mean_anomaly_rad=0.0,
    epoch_unix_s=0.0,
)


def test_circular_orbit_at_epoch() -> None:
    """Test propagation of circular orbit at epoch returns periapsis."""
    state = elements_to_state(_CIRCULAR_ELEMENTS, 0.0)
    assert not isinstance(state, list)

    # At M=0 for circular orbit, should be at x-axis
    assert state.position_eci_m[0] == pytest.approx(7_000_000, rel=1e-6)
    assert state.position_eci_m[1] == pytest.approx(0.0, abs=1e-6)
    assert state.position_eci_m[2] == pytest.approx(0.0, abs=1e-6)

    # Velocity should be in y-direction (perpendicular to x)
    expected_v = _circular_velocity(7_000_000)
    assert state.velocity_eci_ms[0] == pytest.approx(0.0, abs=1e-6)
    assert state.velocity_eci_ms[1] == pytest.approx(expected_v, rel=1e-6)
    assert state.velocity_eci_ms[2] == pytest.approx(0.0, abs=1e-6)

    # Acceleration should be toward center (negative x)
    expected_a = -GM_EARTH / (7_000_000**2)
    assert state.acceleration_eci_ms2[0] == pytest.approx(expected_a, rel=1e-6)


def test_circular_orbit_quarter_period() -> None:
    """Test that quarter period moves 90 degrees around orbit."""
    # Compute period
    a = _CIRCULAR_ELEMENTS.semi_major_axis_m
    period = 2 * math.pi * math.sqrt(a**3 / GM_EARTH)
    quarter_period = period / 4

    state = elements_to_state(_CIRCULAR_ELEMENTS, quarter_period)
    assert not isinstance(state, list)

    # After 1/4 period, should be at 90 degrees (y-axis for equatorial prograde)
    assert state.position_eci_m[0] == pytest.approx(0.0, abs=1e-3)
    assert state.position_eci_m[1] == pytest.approx(7_000_000, rel=1e-6)
    assert state.position_eci_m[2] == pytest.approx(0.0, abs=1e-6)


def test_circular_orbit_half_period() -> None:
    """Test that half period puts satellite on opposite side."""
    a = _CIRCULAR_ELEMENTS.semi_major_axis_m
    period = 2 * math.pi * math.sqrt(a**3 / GM_EARTH)
    half_period = period / 2

    state = elements_to_state(_CIRCULAR_ELEMENTS, half_period)
    assert not isinstance(state, list)

    # After 1/2 period, should be at negative x-axis
    assert state.position_eci_m[0] == pytest.approx(-7_000_000, rel=1e-6)
    assert state.position_eci_m[1] == pytest.approx(0.0, abs=1e-3)
    assert state.position_eci_m[2] == pytest.approx(0.0, abs=1e-6)


def test_circular_orbit_full_period() -> None:
    """Test that full period returns to starting position."""
    a = _CIRCULAR_ELEMENTS.semi_major_axis_m
    period = 2 * math.pi * math.sqrt(a**3 / GM_EARTH)

    state_0 = elements_to_state(_CIRCULAR_ELEMENTS, 0.0)
    state_p = elements_to_state(_CIRCULAR_ELEMENTS, period)
    assert not isinstance(state_0, list)
    assert not isinstance(state_p, list)

    np.testing.assert_allclose(
        state_p.position_eci_m,
        state_0.position_eci_m,
        rtol=1e-6,
    )
    np.testing.assert_allclose(
        state_p.velocity_eci_ms,
        state_0.velocity_eci_ms,
        rtol=1e-6,
    )


def test_eccentric_orbit_apsis_distances() -> None:
    """Test that eccentric orbit reaches correct apogee and perigee."""
    a = 10_000_000  # 10,000 km semi-major axis
    e = 0.2  # eccentricity 0.2

    elements = OrbitalElements(
        semi_major_axis_m=a,
        eccentricity=e,
        inclination_rad=0.0,
        raan_rad=0.0,
        arg_periapsis_rad=0.0,
        mean_anomaly_rad=0.0,
        epoch_unix_s=0.0,
    )

    # At epoch (M=0), satellite is at periapsis
    state_peri = elements_to_state(elements, 0.0)
    assert not isinstance(state_peri, list)
    r_peri = np.linalg.norm(state_peri.position_eci_m)
    expected_peri = a * (1 - e)
    assert r_peri == pytest.approx(expected_peri, rel=1e-6)

    # At M=pi, satellite is at apoapsis
    # Need to find time when M=pi
    n = math.sqrt(GM_EARTH / a**3)  # mean motion
    t_apo = math.pi / n

    state_apo = elements_to_state(elements, t_apo)
    assert not isinstance(state_apo, list)
    r_apo = np.linalg.norm(state_apo.position_eci_m)
    expected_apo = a * (1 + e)
    assert r_apo == pytest.approx(expected_apo, rel=1e-6)


def test_inclined_orbit_z_component() -> None:
    """Test that inclined orbit reaches expected z-height."""
    inclination = math.radians(30)  # 30 degree inclination

    elements = OrbitalElements(
        semi_major_axis_m=7_000_000,
        eccentricity=0.0,
        inclination_rad=inclination,
        raan_rad=0.0,
        arg_periapsis_rad=0.0,
        mean_anomaly_rad=math.pi / 2,  # Start at 90 degrees true anomaly
        epoch_unix_s=0.0,
    )

    state = elements_to_state(elements, 0.0)
    assert not isinstance(state, list)

    # At 90 degrees true anomaly with i=30, should have z = r * sin(i)
    r = np.linalg.norm(state.position_eci_m)
    max_z = r * math.sin(inclination)
    assert abs(state.position_eci_m[2]) == pytest.approx(max_z, rel=1e-6)


def test_multiple_times_returns_list() -> None:
    """Test that array of times returns list of states."""
    times = [0.0, 1000.0, 2000.0]
    states = elements_to_state(_CIRCULAR_ELEMENTS, times)

    assert isinstance(states, list)
    assert len(states) == 3
    for state in states:
        assert state.position_eci_m.shape == (3,)


def test_single_time_returns_scalar() -> None:
    """Test that single time returns single state (not list)."""
    state = elements_to_state(_CIRCULAR_ELEMENTS, 0.0)

    # Should not be a list, should be a single OrbitState
    assert not isinstance(state, list)
    assert state.position_eci_m.shape == (3,)


def test_acceleration_magnitude() -> None:
    """Test that acceleration magnitude follows inverse square law."""
    state = elements_to_state(_CIRCULAR_ELEMENTS, 0.0)
    assert not isinstance(state, list)

    r = np.linalg.norm(state.position_eci_m)
    a = np.linalg.norm(state.acceleration_eci_ms2)

    expected_a = GM_EARTH / r**2
    assert a == pytest.approx(expected_a, rel=1e-6)


def test_acceleration_direction() -> None:
    """Test that acceleration points toward Earth center."""
    state = elements_to_state(_CIRCULAR_ELEMENTS, 0.0)
    assert not isinstance(state, list)

    # Acceleration should be opposite to position direction
    r_unit = state.position_eci_m / np.linalg.norm(state.position_eci_m)
    a_unit = state.acceleration_eci_ms2 / np.linalg.norm(state.acceleration_eci_ms2)

    np.testing.assert_allclose(a_unit, -r_unit, rtol=1e-10)


def test_vis_viva_equation() -> None:
    """Test that velocity magnitude satisfies vis-viva equation."""
    a = 10_000_000
    e = 0.3

    elements = OrbitalElements(
        semi_major_axis_m=a,
        eccentricity=e,
        inclination_rad=0.5,
        raan_rad=1.0,
        arg_periapsis_rad=0.5,
        mean_anomaly_rad=1.0,
        epoch_unix_s=0.0,
    )

    state = elements_to_state(elements, 0.0)
    assert not isinstance(state, list)

    r = np.linalg.norm(state.position_eci_m)
    v = np.linalg.norm(state.velocity_eci_ms)

    # Vis-viva: v^2 = mu * (2/r - 1/a)
    v_expected = math.sqrt(GM_EARTH * (2 / r - 1 / a))
    assert v == pytest.approx(v_expected, rel=1e-6)
