"""Hermite interpolation of orbital state vectors.

Given discrete state vectors (position + velocity at known times), interpolate
position, velocity, and acceleration at arbitrary query times using cubic
Hermite splines. Unlike plain cubic splines which only match positions, Hermite
interpolation matches both position and velocity at each knot, producing
smoother and more physically consistent results between data points.

This is useful when you have ephemeris data from an external source (e.g., a
high-fidelity propagator, SP3 precise ephemeris, or observed track data) and
need to query states at intermediate times.

The implementation uses scipy.interpolate.CubicHermiteSpline, which constructs
a C1-continuous piecewise cubic polynomial matching both function values and
first derivatives at the knot points.

References:
    - Burden, Faires "Numerical Analysis" ch. 3 (Hermite interpolation)
    - scipy.interpolate.CubicHermiteSpline documentation
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
from scipy.interpolate import CubicHermiteSpline

from .orbit_state import OrbitState

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def hermite_interpolate(
    times_s: NDArray[np.floating] | Sequence[float],
    positions_m: NDArray[np.floating],
    velocities_ms: NDArray[np.floating],
    query_times_s: float | NDArray[np.floating] | Sequence[float],
) -> OrbitState | list[OrbitState]:
    """Interpolate orbital state using cubic Hermite splines.

    Constructs a piecewise cubic Hermite interpolant through the given
    (time, position, velocity) knots, then evaluates position, velocity,
    and acceleration at the requested query times.

    Velocity is the exact derivative of the interpolating polynomial.
    Acceleration is the second derivative of the polynomial (piecewise
    linear between knots).

    Args:
        times_s: Knot times in seconds. Must be strictly increasing.
            Shape: (n,) where n >= 2.
        positions_m: Position vectors at knot times, meters.
            Shape: (n, 3) in any consistent frame (e.g., ECI).
        velocities_ms: Velocity vectors at knot times, m/s.
            Shape: (n, 3) in the same frame as positions.
        query_times_s: Time(s) at which to evaluate the interpolant.
            Scalar for a single query, or array-like for multiple.
            Times outside the knot range are extrapolated.

    Returns:
        If query_times_s is scalar: Single OrbitState.
        If query_times_s is array-like: List of OrbitState objects.

    Raises:
        ValueError: If fewer than 2 knot points provided.
        ValueError: If knot times are not strictly increasing.
        ValueError: If positions/velocities shape doesn't match times.

    Example:
        >>> import numpy as np
        >>> times = np.array([0.0, 100.0, 200.0])
        >>> positions = np.array([
        ...     [7_000_000.0, 0.0, 0.0],
        ...     [6_999_000.0, 500_000.0, 0.0],
        ...     [6_996_000.0, 999_000.0, 0.0],
        ... ])
        >>> velocities = np.array([
        ...     [0.0, 7500.0, 0.0],
        ...     [-50.0, 7480.0, 0.0],
        ...     [-100.0, 7440.0, 0.0],
        ... ])
        >>> state = hermite_interpolate(times, positions, velocities, 50.0)
        >>> state.position_eci_m.shape
        (3,)
    """
    times = np.asarray(times_s, dtype=np.float64)
    pos = np.asarray(positions_m, dtype=np.float64)
    vel = np.asarray(velocities_ms, dtype=np.float64)

    _validate_inputs(times, pos, vel)

    # Build one CubicHermiteSpline for all 3 coordinates simultaneously.
    # CubicHermiteSpline accepts y shape (n, 3) and dy shape (n, 3).
    spline = CubicHermiteSpline(times, pos, vel)

    # Evaluate
    is_scalar = np.isscalar(query_times_s)
    query = np.atleast_1d(np.asarray(query_times_s, dtype=np.float64))

    interp_pos = spline(query)  # shape (m, 3)
    interp_vel = spline(query, 1)  # 1st derivative
    interp_acc = spline(query, 2)  # 2nd derivative

    results = [
        OrbitState(
            position_eci_m=interp_pos[i],
            velocity_eci_ms=interp_vel[i],
            acceleration_eci_ms2=interp_acc[i],
        )
        for i in range(len(query))
    ]

    if is_scalar:
        return results[0]
    return results


def _validate_inputs(
    times: NDArray[np.floating],
    positions: NDArray[np.floating],
    velocities: NDArray[np.floating],
) -> None:
    """Validate input arrays for hermite_interpolate.

    Args:
        times: Knot times array.
        positions: Position array.
        velocities: Velocity array.

    Raises:
        ValueError: If inputs are invalid.
    """
    if times.ndim != 1 or len(times) < 2:  # noqa: PLR2004
        msg = f"Need at least 2 knot times, got shape {times.shape}"
        raise ValueError(msg)

    if not np.all(np.diff(times) > 0):
        raise ValueError("Knot times must be strictly increasing")

    n = len(times)
    if positions.shape != (n, 3):
        msg = f"positions shape {positions.shape} doesn't match ({n}, 3)"
        raise ValueError(msg)

    if velocities.shape != (n, 3):
        msg = f"velocities shape {velocities.shape} doesn't match ({n}, 3)"
        raise ValueError(msg)
