# ruff: noqa: T201, S101, PLR2004, E501, C901
"""Compare interpolation methods for orbital ephemeris data.

Run with:
    python -m gri_sandbox.interpolator_comparison.main

Generates dense "truth" ephemeris from SGP4 propagation of real TLEs (which
includes J2, drag, and other perturbations). Subsamples at various intervals,
then each interpolation method reconstructs intermediate states from the
subsampled data and is compared against the dense truth.

Methods compared:
    - Hermite (cubic Hermite spline using position + velocity)
    - Cubic Spline (cubic spline using position only)
    - Linear (piecewise linear using position only)
    - Keplerian (two-body propagation from first state vector)
    - SGP4 (re-propagation from original TLE -- spacing-independent baseline)

Orbit regimes compared (real TLEs):
    - LEO: ISS (ZARYA), ~400 km, i=51.6 deg
    - MEO: GPS BIIF-2 (PRN 01), ~20,200 km, i=55 deg
    - GEO: GALAXY 30, geostationary
    - HEO: MOLNIYA 1-91, e=0.74, i=63.4 deg

Outputs a summary table and Plotly interactive plots (dark theme) of
position/velocity error vs sample interval for each orbit and method.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
import plotly.graph_objects as go
from gri_utils.orbit import (
    hermite_interpolate,
)
from plotly.subplots import make_subplots
from scipy.interpolate import CubicSpline, interp1d
from sgp4.api import Satrec

if TYPE_CHECKING:
    from collections.abc import Callable

    from numpy.typing import NDArray

    # Method signature: (knot_t, knot_pos, knot_vel, query_t, **kwargs) -> (pos, vel)
    InterpolationMethod = Callable[
        ...,
        tuple[NDArray[np.floating], NDArray[np.floating]],
    ]

logger = logging.getLogger(__name__)


# ---- Orbit definitions using real TLEs ---- #


@dataclass(frozen=True)
class OrbitFamily:
    """Defines an orbit regime for comparison using a real TLE."""

    name: str
    tle_line1: str
    tle_line2: str
    duration_s: float  # How long to simulate


_ORBIT_FAMILIES = [
    OrbitFamily(
        name="LEO",
        # ISS (ZARYA) -- epoch 2024-01-01
        tle_line1="1 25544U 98067A   24001.50000000  .00016717  00000-0  10270-3 0  9025",
        tle_line2="2 25544  51.6400  10.5000 0006000  90.0000 270.0000 15.49000000000000",
        duration_s=5560.0,  # ~1 orbit
    ),
    OrbitFamily(
        name="MEO",
        # GPS BIIF-2 (PRN 01) -- epoch 2024-01-01
        tle_line1="1 37753U 11036A   24001.50000000  .00000010  00000-0  00000-0 0  9999",
        tle_line2="2 37753  55.0000 120.0000 0050000   0.0000   0.0000  2.00565000000000",
        duration_s=7200.0,  # 2 hours of a ~12 hr orbit
    ),
    OrbitFamily(
        name="GEO",
        # GALAXY 30 (geostationary) -- epoch 2024-01-01
        tle_line1="1 26900U 01046A   24001.50000000  .00000100  00000-0  00000-0 0  9994",
        tle_line2="2 26900   0.1000   0.0000 0001000   0.0000   0.0000  1.00274000000000",
        duration_s=7200.0,  # 2 hours of a ~24 hr orbit
    ),
    OrbitFamily(
        name="HEO",
        # Molniya 1-91 (highly eccentric) -- epoch 2024-01-01
        tle_line1="1 25485U 98054A   24001.50000000  .00000100  00000-0  10000-3 0  9992",
        tle_line2="2 25485  63.4000 200.0000 7400000 270.0000   0.0000  2.00600000000000",
        duration_s=7200.0,  # 2 hours of a ~12 hr orbit
    ),
]

# Sample intervals to test (seconds)
_INTERVALS_S = [10.0, 30.0, 60.0, 120.0, 300.0, 600.0]

# Plotly color palette (D3 category)
_COLORS = {
    "Hermite": "#1f77b4",
    "Cubic Spline": "#ff7f0e",
    "Linear": "#d62728",
    "Keplerian": "#2ca02c",
    "SGP4": "#9467bd",
}

_SYMBOLS = {
    "Hermite": "circle",
    "Cubic Spline": "square",
    "Linear": "diamond",
    "Keplerian": "triangle-up",
    "SGP4": "star",
}


# ---- SGP4 truth generation ---- #


def _sgp4_propagate(
    satrec: Satrec,
    epoch_jd: float,
    epoch_fr: float,
    dt_offsets_s: NDArray[np.floating],
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Propagate SGP4 at time offsets from epoch.

    Args:
        satrec: SGP4 satellite record.
        epoch_jd: Julian date integer part of reference epoch.
        epoch_fr: Julian date fractional part of reference epoch.
        dt_offsets_s: Time offsets from epoch in seconds.

    Returns:
        Tuple of (positions_teme_m, velocities_teme_ms) arrays, shape (n, 3).
    """
    positions = np.empty((len(dt_offsets_s), 3))
    velocities = np.empty((len(dt_offsets_s), 3))

    for i, dt in enumerate(dt_offsets_s):
        fr = epoch_fr + dt / 86400.0
        jd = epoch_jd + int(fr)
        fr = fr - int(fr)

        error, r_km, v_kms = satrec.sgp4(jd, fr)
        if error != 0:
            msg = f"SGP4 error {error} at dt={dt:.1f}s"
            raise RuntimeError(msg)

        positions[i] = np.array(r_km) * 1000.0
        velocities[i] = np.array(v_kms) * 1000.0

    return positions, velocities


def _generate_sgp4_truth(
    family: OrbitFamily,
    dt_s: float = 1.0,
) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating], Satrec]:
    """Generate dense truth from SGP4 propagation.

    Args:
        family: Orbit family with TLE data.
        dt_s: Truth sampling interval in seconds.

    Returns:
        Tuple of (times_s, positions_teme_m, velocities_teme_ms, satrec).
    """
    satrec = Satrec.twoline2rv(family.tle_line1, family.tle_line2)
    epoch_jd = satrec.jdsatepoch
    epoch_fr = satrec.jdsatepochF

    times = np.arange(0.0, family.duration_s, dt_s)
    positions, velocities = _sgp4_propagate(satrec, epoch_jd, epoch_fr, times)

    return times, positions, velocities, satrec


def _subsample(
    times: NDArray[np.floating],
    positions: NDArray[np.floating],
    velocities: NDArray[np.floating],
    interval_s: float,
) -> tuple[NDArray[np.floating], NDArray[np.floating], NDArray[np.floating]]:
    """Subsample arrays at a given interval.

    Args:
        times: Dense times.
        positions: Dense positions.
        velocities: Dense velocities.
        interval_s: Desired sample interval in seconds.

    Returns:
        Subsampled (times, positions, velocities).
    """
    indices = np.arange(0, len(times), int(interval_s))
    if indices[-1] != len(times) - 1:
        indices = np.append(indices, len(times) - 1)
    return times[indices], positions[indices], velocities[indices]


# ---- Interpolation methods ---- #


@dataclass
class InterpolationResult:
    """Result of an interpolation comparison."""

    method: str
    interval_s: float
    pos_errors_m: NDArray[np.floating]
    vel_errors_ms: NDArray[np.floating]
    elapsed_s: float
    n_queries: int


def _eval_hermite(
    knot_t: NDArray[np.floating],
    knot_pos: NDArray[np.floating],
    knot_vel: NDArray[np.floating],
    query_t: NDArray[np.floating],
    **_kwargs: object,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Evaluate Hermite interpolation at query times."""
    states = hermite_interpolate(knot_t, knot_pos, knot_vel, query_t)
    assert isinstance(states, list)
    pos = np.array([s.position_eci_m for s in states])
    vel = np.array([s.velocity_eci_ms for s in states])
    return pos, vel


def _eval_cubic_spline(
    knot_t: NDArray[np.floating],
    knot_pos: NDArray[np.floating],
    _knot_vel: NDArray[np.floating],
    query_t: NDArray[np.floating],
    **_kwargs: object,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Evaluate cubic spline (position-only) at query times."""
    spline = CubicSpline(knot_t, knot_pos)
    pos = np.asarray(spline(query_t), dtype=np.float64)
    vel = np.asarray(spline(query_t, 1), dtype=np.float64)
    return pos, vel


def _eval_linear(
    knot_t: NDArray[np.floating],
    knot_pos: NDArray[np.floating],
    _knot_vel: NDArray[np.floating],
    query_t: NDArray[np.floating],
    **_kwargs: object,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Evaluate linear interpolation at query times."""
    pos = np.column_stack(
        [interp1d(knot_t, knot_pos[:, i], kind="linear")(query_t) for i in range(3)],
    )
    dt = 0.01
    pos_fwd = np.column_stack(
        [
            interp1d(knot_t, knot_pos[:, i], kind="linear")(
                np.clip(query_t + dt, knot_t[0], knot_t[-1]),
            )
            for i in range(3)
        ],
    )
    return pos, (pos_fwd - pos) / dt


def _eval_keplerian(
    knot_t: NDArray[np.floating],
    knot_pos: NDArray[np.floating],
    knot_vel: NDArray[np.floating],
    query_t: NDArray[np.floating],
    **_kwargs: object,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Evaluate Keplerian two-body propagation from first state vector."""
    from gri_utils.orbit import (  # noqa: PLC0415
        elements_to_state,
        from_position_velocity,
    )

    elements = from_position_velocity(knot_pos[0], knot_vel[0], knot_t[0])
    states = elements_to_state(elements, query_t)
    assert isinstance(states, list)
    pos = np.array([s.position_eci_m for s in states])
    vel = np.array([s.velocity_eci_ms for s in states])
    return pos, vel


def _eval_sgp4(
    _knot_t: NDArray[np.floating],
    _knot_pos: NDArray[np.floating],
    _knot_vel: NDArray[np.floating],
    query_t: NDArray[np.floating],
    *,
    satrec: Satrec | None = None,
    epoch_jd: float = 0.0,
    epoch_fr: float = 0.0,
    **_kwargs: object,
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Evaluate SGP4 re-propagation from original TLE at query times."""
    assert satrec is not None
    return _sgp4_propagate(satrec, epoch_jd, epoch_fr, query_t)


_METHODS: list[tuple[str, InterpolationMethod]] = [
    ("Hermite", _eval_hermite),
    ("Cubic Spline", _eval_cubic_spline),
    ("Linear", _eval_linear),
    ("Keplerian", _eval_keplerian),
    ("SGP4", _eval_sgp4),
]


# ---- Main comparison ---- #


def _compute_errors(
    true_pos: NDArray[np.floating],
    true_vel: NDArray[np.floating],
    interp_pos: NDArray[np.floating],
    interp_vel: NDArray[np.floating],
) -> tuple[NDArray[np.floating], NDArray[np.floating]]:
    """Compute per-point position and velocity error norms."""
    pos_err = np.linalg.norm(interp_pos - true_pos, axis=1)
    vel_err = np.linalg.norm(interp_vel - true_vel, axis=1)
    return pos_err, vel_err


def run_comparison(family: OrbitFamily) -> list[InterpolationResult]:
    """Run interpolation comparison for one orbit family.

    Args:
        family: Orbit family to test.

    Returns:
        List of InterpolationResult for each method/interval combination.
    """
    print(f"\n{'=' * 70}")
    print(f"  {family.name}")
    print(f"  Duration: {family.duration_s / 60:.1f} min")
    print(f"{'=' * 70}")

    truth_t, truth_pos, truth_vel, satrec = _generate_sgp4_truth(family)
    epoch_jd = satrec.jdsatepoch
    epoch_fr = satrec.jdsatepochF

    results = []

    for interval_s in _INTERVALS_S:
        if interval_s > family.duration_s / 4:
            continue

        knot_t, knot_pos, knot_vel = _subsample(
            truth_t,
            truth_pos,
            truth_vel,
            interval_s,
        )

        # Query at non-knot truth times
        query_mask = np.ones(len(truth_t), dtype=bool)
        for kt in knot_t:
            query_mask &= np.abs(truth_t - kt) > 0.5
        query_t = truth_t[query_mask]
        query_pos = truth_pos[query_mask]
        query_vel = truth_vel[query_mask]

        if len(query_t) < 10:
            continue

        for method_name, method_func in _METHODS:
            t_start = time.perf_counter()
            try:
                interp_pos, interp_vel = method_func(
                    knot_t,
                    knot_pos,
                    knot_vel,
                    query_t,
                    satrec=satrec,
                    epoch_jd=epoch_jd,
                    epoch_fr=epoch_fr,
                )
                elapsed = time.perf_counter() - t_start

                pos_err, vel_err = _compute_errors(
                    query_pos,
                    query_vel,
                    interp_pos,
                    interp_vel,
                )
            except Exception:
                logger.exception(
                    "Failed: %s at %ss interval",
                    method_name,
                    interval_s,
                )
                continue

            results.append(
                InterpolationResult(
                    method=method_name,
                    interval_s=interval_s,
                    pos_errors_m=pos_err,
                    vel_errors_ms=vel_err,
                    elapsed_s=elapsed,
                    n_queries=len(query_t),
                ),
            )

    return results


def print_summary(family_name: str, results: list[InterpolationResult]) -> None:
    """Print a formatted summary table for one orbit family."""
    print(f"\n  {family_name} -- Position Error Summary")
    print(
        f"  {'Method':<16} {'Interval':>8} {'RMS (m)':>12} "
        f"{'Max (m)':>12} {'P95 (m)':>12} {'Time (ms)':>10}",
    )
    print(f"  {'-' * 74}")

    for r in sorted(results, key=lambda x: (x.interval_s, x.method)):
        rms = float(np.sqrt(np.mean(r.pos_errors_m**2)))
        max_err = float(np.max(r.pos_errors_m))
        p95 = float(np.percentile(r.pos_errors_m, 95))
        print(
            f"  {r.method:<16} {r.interval_s:>7.0f}s "
            f"{rms:>12.4f} {max_err:>12.4f} {p95:>12.4f} "
            f"{r.elapsed_s * 1000:>9.1f}",
        )

    print(f"\n  {family_name} -- Velocity Error Summary")
    print(
        f"  {'Method':<16} {'Interval':>8} {'RMS (m/s)':>12} "
        f"{'Max (m/s)':>12} {'P95 (m/s)':>12}",
    )
    print(f"  {'-' * 64}")

    for r in sorted(results, key=lambda x: (x.interval_s, x.method)):
        rms = float(np.sqrt(np.mean(r.vel_errors_ms**2)))
        max_err = float(np.max(r.vel_errors_ms))
        p95 = float(np.percentile(r.vel_errors_ms, 95))
        print(
            f"  {r.method:<16} {r.interval_s:>7.0f}s "
            f"{rms:>12.6f} {max_err:>12.6f} {p95:>12.6f}",
        )


def plot_results(all_results: dict[str, list[InterpolationResult]]) -> go.Figure:
    """Create Plotly figure with position error, velocity error, and throughput.

    Args:
        all_results: Dict mapping family name to results.

    Returns:
        Plotly Figure object.
    """
    families = list(all_results.keys())
    n = len(families)

    fig = make_subplots(
        rows=3,
        cols=n,
        subplot_titles=(
            [f"{f} -- Position" for f in families]
            + [f"{f} -- Velocity" for f in families]
            + [f"{f} -- Throughput" for f in families]
        ),
        vertical_spacing=0.08,
        horizontal_spacing=0.06,
    )

    for col, family_name in enumerate(families, start=1):
        results = all_results[family_name]
        methods = sorted({r.method for r in results})

        for method in methods:
            method_results = sorted(
                [r for r in results if r.method == method],
                key=lambda x: x.interval_s,
            )
            if not method_results:
                continue

            intervals = [r.interval_s for r in method_results]
            pos_rms = [
                float(np.sqrt(np.mean(r.pos_errors_m**2))) for r in method_results
            ]
            vel_rms = [
                float(np.sqrt(np.mean(r.vel_errors_ms**2))) for r in method_results
            ]
            us_per_query = [r.elapsed_s / r.n_queries * 1e6 for r in method_results]

            color = _COLORS.get(method, "#888888")
            symbol = _SYMBOLS.get(method, "circle")
            show_legend = col == 1

            # Position error (row 1) -- skip SGP4 (zero error, invisible on log)
            if method != "SGP4":
                fig.add_trace(
                    go.Scatter(
                        x=intervals,
                        y=pos_rms,
                        mode="lines+markers",
                        name=method,
                        legendgroup=method,
                        showlegend=show_legend,
                        line={"color": color, "width": 2},
                        marker={"symbol": symbol, "size": 8},
                        hovertemplate=(
                            f"{method}<br>"
                            "Interval: %{x}s<br>"
                            "RMS: %{y:.4f} m<extra></extra>"
                        ),
                    ),
                    row=1,
                    col=col,
                )

            # Velocity error (row 2) -- skip SGP4
            if method != "SGP4":
                fig.add_trace(
                    go.Scatter(
                        x=intervals,
                        y=vel_rms,
                        mode="lines+markers",
                        name=method,
                        legendgroup=method,
                        showlegend=False,
                        line={"color": color, "width": 2},
                        marker={"symbol": symbol, "size": 8},
                        hovertemplate=(
                            f"{method}<br>"
                            "Interval: %{x}s<br>"
                            "RMS: %{y:.6f} m/s<extra></extra>"
                        ),
                    ),
                    row=2,
                    col=col,
                )

            # Throughput (row 3) -- all methods including SGP4
            # Show legend for SGP4 here since it's excluded from error rows
            show_legend_throughput = method == "SGP4" and col == 1
            fig.add_trace(
                go.Scatter(
                    x=intervals,
                    y=us_per_query,
                    mode="lines+markers",
                    name=method,
                    legendgroup=method,
                    showlegend=show_legend_throughput,
                    line={"color": color, "width": 2},
                    marker={"symbol": symbol, "size": 8},
                    hovertemplate=(
                        f"{method}<br>"
                        "Interval: %{x}s<br>"
                        "%{y:.1f} us/query<extra></extra>"
                    ),
                ),
                row=3,
                col=col,
            )

    # Apply log x-axes to all rows
    for row in range(1, 4):
        for col in range(1, n + 1):
            fig.update_xaxes(
                type="log",
                row=row,
                col=col,
                gridcolor="rgba(255,255,255,0.1)",
            )

    # Position error: 1e-6 to 1e6 m
    for col in range(1, n + 1):
        fig.update_yaxes(
            type="log",
            range=[-6, 6],
            row=1,
            col=col,
            gridcolor="rgba(255,255,255,0.1)",
        )
    # Velocity error: 1e-3 to 1e4 m/s
    for col in range(1, n + 1):
        fig.update_yaxes(
            type="log",
            range=[-3, 4],
            row=2,
            col=col,
            gridcolor="rgba(255,255,255,0.1)",
        )
    # Throughput: log scale, auto range
    for col in range(1, n + 1):
        fig.update_yaxes(
            type="log",
            row=3,
            col=col,
            gridcolor="rgba(255,255,255,0.1)",
        )

    # X-axis labels only on bottom row
    for col in range(1, n + 1):
        fig.update_xaxes(title_text="Sample Interval (s)", row=3, col=col)

    fig.update_yaxes(title_text="Position RMS Error (m)", row=1, col=1)
    fig.update_yaxes(title_text="Velocity RMS Error (m/s)", row=2, col=1)
    fig.update_yaxes(title_text="Time per Query (us)", row=3, col=1)

    fig.update_layout(
        template="plotly_dark",
        title={
            "text": "Ephemeris Interpolator Fidelity vs Sample Interval",
            "font": {"size": 18},
        },
        height=1100,
        width=350 * n,
        legend={
            "font": {"size": 12},
            "bgcolor": "rgba(0,0,0,0.5)",
        },
    )

    return fig


def main() -> None:
    """Run the full interpolator comparison."""
    print("Ephemeris Interpolator Comparison")
    print("Truth: SGP4 propagation from real TLEs at 1-second resolution")
    print(
        "Methods: Hermite, Cubic Spline, Linear, Keplerian (two-body), "
        "SGP4 (re-propagation)",
    )

    all_results: dict[str, list[InterpolationResult]] = {}

    for family in _ORBIT_FAMILIES:
        results = run_comparison(family)
        all_results[family.name] = results
        print_summary(family.name, results)

    fig = plot_results(all_results)

    out_dir = Path(__file__).parent
    html_path = out_dir / "interpolator_comparison.html"
    png_path = out_dir / "interpolator_comparison.png"

    fig.write_html(str(html_path))
    print(f"\nInteractive plot: {html_path}")

    try:
        fig.write_image(str(png_path), scale=2)
        print(f"Static image: {png_path}")
    except (ImportError, ValueError):
        print("(kaleido not installed -- skipping PNG export)")
    fig.show()


if __name__ == "__main__":
    main()
