"""Create OrbitalElements by fitting to multiple state observations.

This module provides orbit determination from multiple observations of position
(and optionally velocity) at different times. It uses least-squares fitting to
find the orbital elements that best match all observations.

This is useful when you have multiple radar or optical observations of a
satellite and want to determine its orbit.

References:
    - Vallado, D. "Fundamentals of Astrodynamics and Applications", 4th ed.
    - Tapley, Schutz, Born "Statistical Orbit Determination"
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from scipy.optimize import least_squares

from .elements_to_state import elements_to_state
from .orbital_elements import OrbitalElements
from .state_to_elements import position_velocity_to_elements

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def from_states(  # noqa: C901, PLR0915
    positions_eci_m: NDArray[np.floating] | Sequence,
    times_unix_s: NDArray[np.floating] | Sequence[float],
    velocities_eci_ms: NDArray[np.floating] | Sequence | None = None,
) -> OrbitalElements:
    """Fit orbital elements to multiple position (and optionally velocity) observations.

    Uses unweighted least-squares optimization to find the orbital elements that
    minimize the residuals between predicted and observed positions/velocities.

    The fitting process:
    1. Uses the first observation with velocity (or estimated velocity) to get
       an initial guess for orbital elements
    2. Optimizes the 6 orbital elements to minimize position/velocity residuals
    3. Returns the best-fit elements with epoch at the first observation time

    Args:
        positions_eci_m: Array of position observations in ECI frame, meters.
            Shape (N, 3) where N is the number of observations.
        times_unix_s: Unix timestamps for each observation in seconds.
            Shape (N,) matching the number of positions.
        velocities_eci_ms: Optional array of velocity observations in ECI frame,
            m/s. Shape (N, 3) or None. If provided, velocities are included in
            the fit. If None, only positions are used.

    Returns:
        OrbitalElements that best fit the observations. The epoch is set to
        the time of the first observation.

    Raises:
        ValueError: If inputs have inconsistent shapes, or if fewer than 3
            position observations are provided, or if fitting fails.

    Example:
        >>> import numpy as np
        >>> from gri_utils.orbit import from_states, elements_to_state
        >>> # Simulate observations by propagating a known orbit
        >>> # ... (generate test positions/times)
        >>> # Fit orbital elements to the observations
        >>> elements = from_states(positions, times)
    """
    positions = np.asarray(positions_eci_m, dtype=np.float64)
    times = np.asarray(times_unix_s, dtype=np.float64)

    if positions.ndim == 1:
        positions = positions.reshape(1, -1)

    n_obs = positions.shape[0]

    if n_obs < 3:  # noqa: PLR2004
        msg = f"Need at least 3 observations, got {n_obs}"
        raise ValueError(msg)

    if positions.shape[1] != 3:  # noqa: PLR2004
        msg = f"Positions must have shape (N, 3), got {positions.shape}"
        raise ValueError(msg)

    if len(times) != n_obs:
        msg = f"Times length {len(times)} doesn't match positions count {n_obs}"
        raise ValueError(msg)

    use_velocities = velocities_eci_ms is not None
    if use_velocities:
        velocities = np.asarray(velocities_eci_ms, dtype=np.float64)
        if velocities.ndim == 1:
            velocities = velocities.reshape(1, -1)
        if velocities.shape != positions.shape:
            msg = f"Velocities shape {velocities.shape} doesn't match positions {positions.shape}"  # noqa: E501
            raise ValueError(msg)
    else:
        velocities = None

    # Get initial guess from first two observations
    # Estimate velocity from position difference if not provided
    epoch = times[0]

    if use_velocities and velocities is not None:
        initial_velocity = velocities[0]
    else:
        # Estimate velocity from first two positions
        dt = times[1] - times[0]
        if abs(dt) < 1e-6:  # noqa: PLR2004
            msg = "First two observations have identical times"
            raise ValueError(msg)
        initial_velocity = (positions[1] - positions[0]) / dt

    # Get initial orbital elements
    initial_elements = position_velocity_to_elements(
        positions[0],
        initial_velocity,
        epoch,
    )

    # Pack elements into parameter array for optimization
    # [a, e, i, raan, omega, M0]
    x0 = np.array(
        [
            initial_elements.semi_major_axis_m,
            initial_elements.eccentricity,
            initial_elements.inclination_rad,
            initial_elements.raan_rad,
            initial_elements.arg_periapsis_rad,
            initial_elements.mean_anomaly_rad,
        ],
    )

    # Scale factors for better conditioning
    # Semi-major axis is in meters (millions), others are radians (order 1)
    scales = np.array([1e6, 1.0, 1.0, 1.0, 1.0, 1.0])

    def residuals(x: NDArray[np.floating]) -> NDArray[np.floating]:
        """Compute residuals between predicted and observed states."""
        # Unpack and unscale parameters
        a = x[0] * scales[0]
        e = x[1]
        i = x[2]
        raan = x[3]
        omega = x[4]
        m0 = x[5]

        # Ensure valid parameters
        if a <= 0 or e < 0 or e >= 1:
            return np.full(n_obs * 3 if not use_velocities else n_obs * 6, 1e10)

        elements = OrbitalElements(
            semi_major_axis_m=a,
            eccentricity=e,
            inclination_rad=i,
            raan_rad=raan,
            arg_periapsis_rad=omega,
            mean_anomaly_rad=m0,
            epoch_unix_s=epoch,
        )

        # Compute predicted states (times is always an array here, so this returns a
        # list)
        predicted_states = elements_to_state(elements, times)
        if not isinstance(predicted_states, list):
            predicted_states = [predicted_states]

        # Compute residuals
        res_list = []
        for j in range(n_obs):
            state = predicted_states[j]
            # Position residual (in km for better scaling)
            pos_res = (state.position_eci_m - positions[j]) / 1000.0
            res_list.extend(pos_res.tolist())

            if use_velocities and velocities is not None:
                # Velocity residual (in m/s)
                vel_res = state.velocity_eci_ms - velocities[j]
                res_list.extend(vel_res.tolist())

        return np.array(res_list)

    # Scale initial guess
    x0_scaled = x0 / scales

    # Run optimization
    result = least_squares(
        residuals,
        x0_scaled,
        method="lm",  # Levenberg-Marquardt
        ftol=1e-10,
        xtol=1e-10,
        gtol=1e-10,
        max_nfev=1000,
    )

    if not result.success:
        msg = f"Orbit fitting failed: {result.message}"
        raise ValueError(msg)

    # Unpack result
    x_final = result.x * scales

    return OrbitalElements(
        semi_major_axis_m=x_final[0],
        eccentricity=max(0.0, min(x_final[1], 0.9999)),  # Clamp eccentricity
        inclination_rad=x_final[2] % (2.0 * np.pi),
        raan_rad=x_final[3] % (2.0 * np.pi),
        arg_periapsis_rad=x_final[4] % (2.0 * np.pi),
        mean_anomaly_rad=x_final[5] % (2.0 * np.pi),
        epoch_unix_s=epoch,
    )
