# ruff: noqa: T201
"""Sparse geolocation with observable fill-ins.

Run with:
    python -m gri_sandbox.kalman_tracker_demo.demo_obs_only

Shows what happens when you get a good Ell fix every Nth step,
with observable fill-ins between fixes. Uses a LEO + HEO collector
pair with realistic circular orbits (5-second observation cadence).

The LEO sweeps west-to-east across the sky while the HEO sits
nearly stationary to the northwest, providing continuously changing
TDOA/FDOA geometry over the ~6-minute LEO pass.

Compares Coast / TDOA / FDOA / TDOA+FDOA at two Ell cadences.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass

import numpy as np
import plotly.graph_objects as go
from gri_convolve.kalman import IMM, ConstantVelocity
from gri_convolve.kalman.motion import MotionModel
from gri_obs import FdoaObs, TdoaObs
from gri_pos import Pos
from gri_utils.conversion import (
    eci_to_ecef,
    eci_to_ecef_vel,
    get_aer,
    translate_aer,
)
from gri_utils.distance import slant_range
from gri_utils.observables import get_fdoa, get_tdoa
from gri_utils.orbit import OrbitalElements, elements_to_state, from_circular
from plotly.subplots import make_subplots

from gri_sandbox.kalman_tracker_demo._common import X0, make_ell

_DT = 5.0  # 5-second observation cadence
_N_WARMUP = 4  # initial Ell-only steps to establish velocity
_N_STEPS = 75 + _N_WARMUP  # warmup + LEO pass
_TDOA_STD = 1e-7  # 100 ns
_FDOA_STD = 0.5  # 0.5 Hz
_FREQ_HZ = 1e9  # 1 GHz
_ELL_SIGMA = 500.0  # Ell fix quality (1-sigma meters)
_MIN_EL_DEG = 5.0  # minimum elevation for observable collection

# Orbit epoch (arbitrary)
_UNIX_T0 = 1_704_067_200.0

_STRATEGIES = ["Coast", "TDOA", "FDOA", "BOTH"]

ModelFactory = Callable[[], MotionModel]


@dataclass
class _StepResult:
    """Per-step tracking data."""

    time: float
    pos_error_m: float
    std_e: float
    std_n: float
    std_u: float
    is_ell_step: bool
    has_obs: bool  # whether observables were available (LEO visible)


# ---------------------------------------------------------------------------
# Orbit setup
# ---------------------------------------------------------------------------


def _setup_orbits() -> tuple:
    """Create LEO and HEO orbital elements.

    Returns:
        (leo_elements, heo_elements)
    """
    # LEO: circular orbit at 550km, starting at az=290, el=20
    # from_circular creates a prograde orbit through this ECEF point
    leo_range = slant_range(X0, 290.0, 20.0, 550_000.0)
    leo_ecef = translate_aer(X0, np.array([290.0, 20.0, leo_range]))
    leo_elem = from_circular(leo_ecef, unix_s=_UNIX_T0)

    # HEO: circular orbit at 35,000km alt, to the NW at el=45
    # Near apogee speed (~2.4 km/s), barely moves during the pass
    heo_range = slant_range(X0, 315.0, 45.0, 35_000_000.0)
    heo_ecef = translate_aer(X0, np.array([315.0, 45.0, heo_range]))
    heo_elem = from_circular(heo_ecef, unix_s=_UNIX_T0)

    return leo_elem, heo_elem


def _get_collector_state(
    elements: OrbitalElements,
    unix_s: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Get ECEF position and velocity for an orbit at a given time."""
    result = elements_to_state(elements, unix_s)
    # Scalar unix_s always returns a single OrbitState
    if isinstance(result, list):  # pragma: no cover
        msg = "Expected single OrbitState, got list"
        raise TypeError(msg)
    state = result
    pos_ecef = eci_to_ecef(state.position_eci_m, unix_s)
    vel_ecef = eci_to_ecef_vel(
        state.position_eci_m,
        state.velocity_eci_ms,
        unix_s,
    )
    return np.asarray(pos_ecef, dtype=np.float64), np.asarray(
        vel_ecef,
        dtype=np.float64,
    )


# ---------------------------------------------------------------------------
# Observable factories
# ---------------------------------------------------------------------------


def _make_tdoa(
    emitter_xyz: np.ndarray,
    leo_xyz: np.ndarray,
    heo_xyz: np.ndarray,
    rng: np.random.Generator,
    time: float,
) -> TdoaObs:
    """Create a noisy TdoaObs from LEO+HEO pair."""
    true_tdoa = float(get_tdoa(emitter_xyz, leo_xyz, heo_xyz))
    return TdoaObs(
        time=time,
        value=true_tdoa + rng.normal(0, _TDOA_STD),
        std_dev=_TDOA_STD,
        collector1=Pos.XYZ(leo_xyz),
        collector2=Pos.XYZ(heo_xyz),
        sensor_id="leo_heo_tdoa",
    )


def _make_fdoa(  # noqa: PLR0913
    emitter_xyz: np.ndarray,
    emitter_vel: np.ndarray,
    leo_xyz: np.ndarray,
    heo_xyz: np.ndarray,
    leo_vel: np.ndarray,
    heo_vel: np.ndarray,
    rng: np.random.Generator,
    time: float,
) -> FdoaObs:
    """Create a noisy FdoaObs from LEO+HEO pair."""
    true_fdoa = float(
        get_fdoa(
            emitter_xyz,
            emitter_vel,
            leo_xyz,
            heo_xyz,
            _FREQ_HZ,
            collector1_vel=leo_vel,
            collector2_vel=heo_vel,
        ),
    )
    return FdoaObs(
        time=time,
        value=true_fdoa + rng.normal(0, _FDOA_STD),
        std_dev=_FDOA_STD,
        collector1=Pos.XYZ(leo_xyz),
        collector2=Pos.XYZ(heo_xyz),
        collector1_vel=leo_vel,
        collector2_vel=heo_vel,
        frequency_hz=_FREQ_HZ,
        sensor_id="leo_heo_fdoa",
    )


# ---------------------------------------------------------------------------
# Scenario runner
# ---------------------------------------------------------------------------


def _run_scenario(  # noqa: C901, PLR0912, PLR0913, PLR0915
    truth_pos: np.ndarray,
    truth_vel: np.ndarray,
    rng_seed: int,
    model_factory: ModelFactory,
    ell_every_n: int,
    title: str,
) -> dict[str, list[_StepResult]]:
    """Run sparse-fix tracking with observable fill-ins."""
    leo_elem, heo_elem = _setup_orbits()
    times = [i * _DT for i in range(_N_STEPS)]

    # Pre-compute collector states and visibility
    leo_states: list[tuple[np.ndarray, np.ndarray]] = []
    heo_states: list[tuple[np.ndarray, np.ndarray]] = []
    leo_visible: list[bool] = []

    for i in range(_N_STEPS):
        unix_t = _UNIX_T0 + times[i]
        lp, lv = _get_collector_state(leo_elem, unix_t)
        hp, hv = _get_collector_state(heo_elem, unix_t)
        leo_states.append((lp, lv))
        heo_states.append((hp, hv))
        aer = get_aer(X0, lp)
        leo_visible.append(bool(aer[1] > _MIN_EL_DEG))

    # Pre-build observations (shared across strategies)
    rng_obs = np.random.default_rng(rng_seed)
    ell_noise = [rng_obs.normal(0, _ELL_SIGMA, 3) for _ in range(_N_STEPS)]

    obs_per_step: list[tuple[TdoaObs | None, FdoaObs | None]] = []
    for i in range(_N_STEPS):
        if leo_visible[i]:
            rng_step = np.random.default_rng(rng_seed + 100 * (i + 1))
            lp, lv = leo_states[i]
            hp, hv = heo_states[i]
            tdoa = _make_tdoa(truth_pos[i], lp, hp, rng_step, times[i])
            fdoa = _make_fdoa(
                truth_pos[i],
                truth_vel[i],
                lp,
                hp,
                lv,
                hv,
                rng_step,
                times[i],
            )
            obs_per_step.append((tdoa, fdoa))
        else:
            obs_per_step.append((None, None))

    # Run each strategy independently
    all_results: dict[str, list[_StepResult]] = {}

    for strategy in _STRATEGIES:
        tracker = IMM([model_factory()])
        results: list[_StepResult] = []

        for i in range(_N_STEPS):
            t = times[i]
            # Warm-up: Ell every step. After warm-up: Ell every N steps.
            is_ell = True if i < _N_WARMUP else (i - _N_WARMUP) % ell_every_n == 0
            tdoa, fdoa = obs_per_step[i]
            has_obs = tdoa is not None

            if is_ell:
                noisy_pos = truth_pos[i] + ell_noise[i]
                tracker.update(make_ell(noisy_pos, _ELL_SIGMA), timestamp_s=t)
            elif has_obs and tdoa is not None and fdoa is not None:
                if strategy == "Coast":
                    pass
                elif strategy == "TDOA":
                    tracker.update_observable(tdoa, timestamp_s=t)
                elif strategy == "FDOA":
                    tracker.update_observable(fdoa, timestamp_s=t)
                elif strategy == "BOTH":
                    tracker.update_observable(tdoa, timestamp_s=t)
                    tracker.update_observable(fdoa, timestamp_s=t)

            # Record state
            if tracker.is_initialized:
                if not is_ell and strategy == "Coast" and has_obs:
                    ell = tracker.predict(timestamp_s=t)
                elif not is_ell and not has_obs and strategy != "Coast":
                    # LEO below horizon, no obs available
                    ell = tracker.predict(timestamp_s=t)
                else:
                    ell = tracker.result.state

                est_xyz = np.asarray(ell.xyz)
                cov = ell.cov
                results.append(
                    _StepResult(
                        time=t,
                        pos_error_m=float(np.linalg.norm(est_xyz - truth_pos[i])),
                        std_e=float(np.sqrt(cov[0, 0])),
                        std_n=float(np.sqrt(cov[1, 1])),
                        std_u=float(np.sqrt(cov[2, 2])),
                        is_ell_step=is_ell,
                        has_obs=has_obs,
                    ),
                )

        all_results[strategy] = results

    # Print summary table (every 5th step to keep it readable)
    w = 105
    print(f"\n{'=' * w}")
    print(f"  {title}")
    print(
        f"  {_N_WARMUP} Ell warm-up, then Ell every {ell_every_n} steps "
        f"({ell_every_n * _DT:.0f}s) | Obs: {_DT:.0f}s | LEO+HEO",
    )
    print(f"{'=' * w}")
    print(
        f"{'':>3} {'':>5} {'':>3}  "
        f"{'-- Coast --':^17}  "
        f"{'--- TDOA ---':^17}  "
        f"{'--- FDOA ---':^17}  "
        f"{'- TDOA+FDOA -':^17}",
    )
    print(
        f"{'i':>3} {'t':>5} {'':>3}  "
        f"{'Err':>6} {'E':>5} {'N':>5}  "
        f"{'Err':>6} {'E':>5} {'N':>5}  "
        f"{'Err':>6} {'E':>5} {'N':>5}  "
        f"{'Err':>6} {'E':>5} {'N':>5}",
    )
    print(f"{'-' * w}")

    for i in range(len(all_results["Coast"])):
        r_coast = all_results["Coast"][i]
        # Print every 5th step, plus Ell steps and first/last
        if (
            i % 5 != 0
            and not r_coast.is_ell_step
            and i != len(all_results["Coast"]) - 1
        ):
            continue

        tag = "Ell" if r_coast.is_ell_step else ("obs" if r_coast.has_obs else "---")
        parts = [f"{i:3d} {r_coast.time:5.0f} {tag:>3}  "]
        for strategy in _STRATEGIES:
            r = all_results[strategy][i]
            parts.append(
                f"{r.pos_error_m:6.0f} {r.std_e:5.0f} {r.std_n:5.0f}  ",
            )
        print("".join(parts))

    n_vis = sum(1 for v in leo_visible if v)
    print(f"\n  LEO visible for {n_vis} of {_N_STEPS} steps ({n_vis * _DT:.0f}s)")
    print(
        f"  Ell: {_ELL_SIGMA:.0f}m | TDOA: {_TDOA_STD * 1e9:.0f}ns | "
        f"FDOA: {_FDOA_STD}Hz @ {_FREQ_HZ / 1e9:.0f}GHz",
    )

    return all_results


def _plot_results(
    results_by_cadence: dict[int, dict[str, list[_StepResult]]],
) -> None:
    """Plot position error over time, one subplot per cadence."""
    colors = {
        "Coast": "#888888",
        "TDOA": "#636EFA",
        "FDOA": "#00CC96",
        "BOTH": "#FFA15A",
    }
    labels = {
        "Coast": "Coast",
        "TDOA": "TDOA",
        "FDOA": "FDOA",
        "BOTH": "TDOA+FDOA",
    }

    cadences = sorted(results_by_cadence.keys())
    fig = make_subplots(
        rows=1,
        cols=len(cadences),
        subplot_titles=[f"Ell every {n} steps ({n * _DT:.0f}s)" for n in cadences],
        horizontal_spacing=0.08,
    )

    for col, n in enumerate(cadences, start=1):
        results = results_by_cadence[n]
        for strategy in _STRATEGIES:
            steps = results[strategy]
            fig.add_trace(
                go.Scatter(
                    x=[r.time for r in steps],
                    y=[r.pos_error_m for r in steps],
                    mode="lines",
                    name=labels[strategy],
                    line={
                        "color": colors[strategy],
                        "width": 2,
                        "dash": "dash" if strategy == "Coast" else "solid",
                    },
                    legendgroup=strategy,
                    showlegend=(col == 1),
                ),
                row=1,
                col=col,
            )

        # Mark Ell fix steps
        ell_times = [r.time for r in results["Coast"] if r.is_ell_step]
        for et in ell_times:
            fig.add_vline(
                x=et,
                line_dash="dot",
                line_color="rgba(255,255,255,0.15)",
            )

    fig.update_layout(
        template="plotly_dark",
        title=(
            "Sparse Geolocation + Observable Fill-ins: "
            "Position Error (CV ~14 m/s NE)"
            f"<br><sup>LEO+HEO pair | "
            f"Ell: {_ELL_SIGMA:.0f}m | "
            f"TDOA: {_TDOA_STD * 1e9:.0f}ns | "
            f"FDOA: {_FDOA_STD}Hz | "
            f"Obs cadence: {_DT:.0f}s</sup>"
        ),
        height=450,
        width=1100,
    )
    for col in range(1, len(cadences) + 1):
        fig.update_xaxes(title_text="Time (s)", row=1, col=col)
        fig.update_yaxes(
            title_text="Position Error (m)",
            range=[0, 2000],
            row=1,
            col=col,
        )

    fig.show()


# ---------------------------------------------------------------------------
# Truth generator
# ---------------------------------------------------------------------------


def _generate_cv_truth() -> tuple[np.ndarray, np.ndarray]:
    """~14 m/s NE. E=+Y, N=+Z, U=+X at X0."""
    vel = np.array([0.0, 10.0, 10.0])
    positions = np.array([X0 + vel * i * _DT for i in range(_N_STEPS)])
    velocities = np.tile(vel, (_N_STEPS, 1))
    return positions, velocities


def main() -> None:
    """Run sparse geolocation demos at two cadences."""
    cv_pos, cv_vel = _generate_cv_truth()

    print(f"\n{'#' * 70}")
    print("  Sparse Geolocation + Observable Fill-ins (CV)")
    print(f"  {_N_STEPS} steps x {_DT:.0f}s = {_N_STEPS * _DT / 60:.1f} min")
    print("  LEO (550km, west-to-east pass) + HEO (35000km, NW, ~stationary)")
    print(f"{'#' * 70}")

    results_by_cadence: dict[int, dict[str, list[_StepResult]]] = {}

    for ell_every_n in [20, 10]:
        results = _run_scenario(
            truth_pos=cv_pos,
            truth_vel=cv_vel,
            rng_seed=42,
            model_factory=lambda: ConstantVelocity(process_noise_m_s2=1.0),
            ell_every_n=ell_every_n,
            title=f"CV (~14 m/s NE) -- Ell every {ell_every_n} steps "
            f"({ell_every_n * _DT:.0f}s)",
        )
        results_by_cadence[ell_every_n] = results

    _plot_results(results_by_cadence)


if __name__ == "__main__":
    main()
