"""Tests for from_states (orbit fitting from multiple observations)."""

import numpy as np
import pytest

from gri_utils.orbit import (
    OrbitalElements,
    elements_to_state,
    from_states,
    orbital_properties,
)


def _generate_observations(
    elements: OrbitalElements,
    n_obs: int,
    fraction_of_orbit: float = 0.5,
    *,
    include_velocities: bool = False,
) -> tuple:
    """Generate synthetic observations from known orbit."""
    props = orbital_properties(elements)
    period = props.period_s

    # Generate times spanning a fraction of the orbit
    times = np.linspace(0, period * fraction_of_orbit, n_obs)

    # Get states at those times (always returns list when given array)
    states = elements_to_state(elements, times)
    assert isinstance(states, list)

    positions = np.array([s.position_eci_m for s in states])
    velocities = (
        np.array([s.velocity_eci_ms for s in states]) if include_velocities else None
    )

    return positions, times, velocities


def test_fit_circular_orbit() -> None:
    """Test fitting a circular orbit from observations."""
    true_elements = OrbitalElements(
        semi_major_axis_m=7_000_000,
        eccentricity=0.0,
        inclination_rad=0.5,
        raan_rad=1.0,
        arg_periapsis_rad=0.0,
        mean_anomaly_rad=0.0,
        epoch_unix_s=0.0,
    )

    positions, times, _ = _generate_observations(true_elements, n_obs=10)

    fitted = from_states(positions, times)

    # Semi-major axis should be recovered
    assert fitted.semi_major_axis_m == pytest.approx(
        true_elements.semi_major_axis_m,
        rel=1e-4,
    )
    # Eccentricity should be near zero
    assert fitted.eccentricity == pytest.approx(0.0, abs=1e-4)
    # Inclination should be recovered
    assert fitted.inclination_rad == pytest.approx(
        true_elements.inclination_rad,
        rel=1e-3,
    )


def test_fit_elliptical_orbit() -> None:
    """Test fitting an elliptical orbit from observations."""
    true_elements = OrbitalElements(
        semi_major_axis_m=8_000_000,
        eccentricity=0.1,
        inclination_rad=0.8,
        raan_rad=0.5,
        arg_periapsis_rad=1.0,
        mean_anomaly_rad=0.0,
        epoch_unix_s=0.0,
    )

    positions, times, _ = _generate_observations(true_elements, n_obs=15)

    fitted = from_states(positions, times)

    assert fitted.semi_major_axis_m == pytest.approx(
        true_elements.semi_major_axis_m,
        rel=1e-3,
    )
    assert fitted.eccentricity == pytest.approx(
        true_elements.eccentricity,
        abs=1e-3,
    )


def test_fit_with_velocities() -> None:
    """Test that including velocities improves fit."""
    true_elements = OrbitalElements(
        semi_major_axis_m=7_500_000,
        eccentricity=0.05,
        inclination_rad=0.6,
        raan_rad=0.3,
        arg_periapsis_rad=0.8,
        mean_anomaly_rad=0.0,
        epoch_unix_s=0.0,
    )

    positions, times, velocities = _generate_observations(
        true_elements,
        n_obs=5,
        include_velocities=True,
    )

    # Fit with velocities - should be very accurate even with few points
    fitted = from_states(positions, times, velocities)

    assert fitted.semi_major_axis_m == pytest.approx(
        true_elements.semi_major_axis_m,
        rel=1e-5,
    )


def test_minimum_observations() -> None:
    """Test with minimum 3 observations."""
    true_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,
    )

    positions, times, _ = _generate_observations(true_elements, n_obs=3)

    # Should work with exactly 3 observations
    fitted = from_states(positions, times)
    assert fitted.semi_major_axis_m > 0


def test_too_few_observations_raises() -> None:
    """Test that fewer than 3 observations raises error."""
    positions = np.array([[7e6, 0, 0], [0, 7e6, 0]])
    times = np.array([0.0, 1000.0])

    with pytest.raises(ValueError, match="at least 3"):
        from_states(positions, times)


def test_shape_mismatch_raises() -> None:
    """Test that mismatched shapes raise error."""
    positions = np.array([[7e6, 0, 0], [0, 7e6, 0], [6e6, 3e6, 0]])
    times = np.array([0.0, 1000.0])  # Only 2 times for 3 positions

    with pytest.raises(ValueError, match="doesn't match"):
        from_states(positions, times)


def test_velocity_shape_mismatch_raises() -> None:
    """Test that velocity shape mismatch raises error."""
    positions = np.array([[7e6, 0, 0], [0, 7e6, 0], [6e6, 3e6, 0]])
    times = np.array([0.0, 1000.0, 2000.0])
    velocities = np.array([[0, 7e3, 0], [7e3, 0, 0]])  # Only 2 velocities

    with pytest.raises(ValueError, match="doesn't match"):
        from_states(positions, times, velocities)


def test_epoch_set_to_first_time() -> None:
    """Test that epoch is set to first observation time."""
    true_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=1000.0,  # Non-zero epoch
    )

    positions, times, _ = _generate_observations(true_elements, n_obs=5)
    times = times + 1000.0  # Shift times to start at 1000

    fitted = from_states(positions, times)

    assert fitted.epoch_unix_s == times[0]


def test_propagated_states_match_observations() -> None:
    """Test that fitted orbit propagates back to observations."""
    true_elements = OrbitalElements(
        semi_major_axis_m=7_000_000,
        eccentricity=0.02,
        inclination_rad=0.4,
        raan_rad=0.2,
        arg_periapsis_rad=0.6,
        mean_anomaly_rad=0.0,
        epoch_unix_s=0.0,
    )

    positions, times, _ = _generate_observations(true_elements, n_obs=10)

    fitted = from_states(positions, times)

    # Propagate fitted orbit back to observation times
    fitted_states = elements_to_state(fitted, times)
    assert isinstance(fitted_states, list)
    fitted_positions = np.array([s.position_eci_m for s in fitted_states])

    # Residuals should be small (< 1 km)
    residuals = np.linalg.norm(fitted_positions - positions, axis=1)
    assert np.max(residuals) < 1000.0


def test_noisy_observations() -> None:
    """Test fitting with slightly noisy observations."""
    true_elements = OrbitalElements(
        semi_major_axis_m=7_500_000,
        eccentricity=0.01,
        inclination_rad=0.5,
        raan_rad=0.0,
        arg_periapsis_rad=0.0,
        mean_anomaly_rad=0.0,
        epoch_unix_s=0.0,
    )

    positions, times, _ = _generate_observations(true_elements, n_obs=20)

    # Add small noise (100 meters = 0.1 km)
    rng = np.random.default_rng(42)
    noise = rng.normal(0, 100, positions.shape)
    noisy_positions = positions + noise

    fitted = from_states(noisy_positions, times)

    # Should still recover approximate orbit
    assert fitted.semi_major_axis_m == pytest.approx(
        true_elements.semi_major_axis_m,
        rel=1e-3,
    )
