# ruff: noqa: T201
"""Observable fill-in demo: comparing TDOA, FDOA, and TDOA+FDOA.

Run with:
    python -m gri_sandbox.kalman_tracker_demo.demo_observable

Four trackers start from the same 4 Ell updates (CV track), then
diverge for 4 steps with different fill-in strategies:

    1. Coast    -- no updates (prediction only)
    2. TDOA     -- 1 TDOA per step
    3. FDOA     -- 1 FDOA per step
    4. BOTH     -- 1 TDOA + 1 FDOA per step (same collector pair)

Shows per-axis ENU std devs and covariance volume ratio vs coasting.
All strategies receive identical observations (same noise realization).
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import dataclass
from typing import TYPE_CHECKING

import numpy as np
import plotly.graph_objects as go
from gri_convolve.kalman import IMM, ConstantVelocity, CoordinatedTurn
from gri_convolve.kalman.motion import MotionModel
from gri_obs import FdoaObs, TdoaObs
from gri_pos import Pos
from gri_utils.observables import get_fdoa, get_tdoa
from plotly.subplots import make_subplots

if TYPE_CHECKING:
    from gri_ell import Ell

from gri_sandbox.kalman_tracker_demo._common import DT, SIGMA, X0, make_ell

_N_ELL = 4  # shared Ell warm-up
_N_FILL = 4  # fill-in steps
_N_TOTAL = _N_ELL + _N_FILL

# Two LEO satellites at ~550km altitude, viewed from the emitter at X0.
# At X0 = [6.4e6, 0, 0] (equator, prime meridian): E=+Y, N=+Z, U=+X.
#
# LEO 1: near the southern horizon (az ~190, el ~20)
_LEO1 = X0 + np.array([200_000.0, -100_000.0, -500_000.0])
# LEO 2: near the eastern horizon (az ~80, el ~15)
_LEO2 = X0 + np.array([150_000.0, 550_000.0, 100_000.0])

# LEO velocities (~7.5 km/s orbital speed)
_LEO1_VEL = np.array([0.0, 0.0, -7500.0])  # heading south
_LEO2_VEL = np.array([0.0, 7500.0, 0.0])  # heading east

_FREQ_HZ = 1e9  # 1 GHz carrier for FDOA
_TDOA_STD = 1e-7  # 100 ns
_FDOA_STD = 0.5  # 0.5 Hz

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

# Factory callable that returns a fresh MotionModel instance
ModelFactory = Callable[[], MotionModel]


@dataclass
class _StepInfo:
    """Collected uncertainty data for one step."""

    sma_95: float
    std_e: float
    std_n: float
    std_u: float
    cov_det: float


def _ell_to_info(ell: Ell) -> _StepInfo:
    """Extract uncertainty info from an Ell."""
    cov = ell.cov  # ENU, 1-sigma
    return _StepInfo(
        sma_95=ell.ellipse.sma_95,
        std_e=float(np.sqrt(cov[0, 0])),
        std_n=float(np.sqrt(cov[1, 1])),
        std_u=float(np.sqrt(cov[2, 2])),
        cov_det=float(np.linalg.det(cov)),
    )


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


def _make_fdoa(
    emitter_xyz: np.ndarray,
    emitter_vel: np.ndarray,
    rng: np.random.Generator,
    time: float,
) -> FdoaObs:
    """Create a noisy FdoaObs from the LEO pair."""
    true_fdoa = float(
        get_fdoa(
            emitter_xyz,
            emitter_vel,
            _LEO1,
            _LEO2,
            _FREQ_HZ,
            collector1_vel=_LEO1_VEL,
            collector2_vel=_LEO2_VEL,
        ),
    )
    return FdoaObs(
        time=time,
        value=true_fdoa + rng.normal(0, _FDOA_STD),
        std_dev=_FDOA_STD,
        collector1=Pos.XYZ(_LEO1),
        collector2=Pos.XYZ(_LEO2),
        collector1_vel=_LEO1_VEL,
        collector2_vel=_LEO2_VEL,
        frequency_hz=_FREQ_HZ,
        sensor_id="leo_fdoa",
    )


@dataclass
class _ScenarioResult:
    """Final-step data for plotting."""

    title: str
    start: _StepInfo
    final: dict[str, _StepInfo]


def _run_scenario(  # noqa: C901, PLR0912, PLR0915
    truth: np.ndarray,
    truth_vel: np.ndarray,
    rng_seed: int,
    model_factory: ModelFactory,
    title: str,
) -> _ScenarioResult:
    """Run one scenario: shared Ell warm-up, then 4 diverging strategies."""
    times = [i * DT for i in range(_N_TOTAL)]

    # Phase 1: shared Ell warm-up (save noisy positions for reuse)
    rng_warmup = np.random.default_rng(rng_seed)
    noisy_positions = [truth[i] + rng_warmup.normal(0, SIGMA, 3) for i in range(_N_ELL)]
    boot_ells = [make_ell(p) for p in noisy_positions]

    # Phase 2: build observations once (shared across all strategies)
    obs_per_step: list[tuple[TdoaObs, FdoaObs]] = []
    for i in range(_N_ELL, _N_TOTAL):
        rng_step = np.random.default_rng(rng_seed + i + 1)
        tdoa = _make_tdoa(truth[i], rng_step, times[i])
        fdoa = _make_fdoa(truth[i], truth_vel[i], rng_step, times[i])
        obs_per_step.append((tdoa, fdoa))

    # Phase 3: run each strategy with its own tracker (fresh model)
    infos: dict[str, list[_StepInfo]] = {}

    for strategy in _STRATEGIES:
        tracker = IMM([model_factory()])
        step_infos: list[_StepInfo] = []

        # Warm-up (identical for all)
        for i in range(_N_ELL):
            tracker.update(boot_ells[i], timestamp_s=times[i])
            step_infos.append(_ell_to_info(tracker.result.state))

        # Fill-in (strategy-dependent)
        for j, i in enumerate(range(_N_ELL, _N_TOTAL)):
            t = times[i]
            tdoa, fdoa = obs_per_step[j]

            if strategy == "Coast":
                pred_ell = tracker.predict(timestamp_s=t)
                step_infos.append(_ell_to_info(pred_ell))
            elif strategy == "TDOA":
                tracker.update_observable(tdoa, timestamp_s=t)
                step_infos.append(_ell_to_info(tracker.result.state))
            elif strategy == "FDOA":
                tracker.update_observable(fdoa, timestamp_s=t)
                step_infos.append(_ell_to_info(tracker.result.state))
            elif strategy == "BOTH":
                tracker.update_observable(tdoa, timestamp_s=t)
                tracker.update_observable(fdoa, timestamp_s=t)
                step_infos.append(_ell_to_info(tracker.result.state))

        infos[strategy] = step_infos

    # -- Print comparison table --
    w = 113
    print(f"\n{'=' * w}")
    print(f"  {title}")
    print(
        f"  Ell({_N_ELL}) warm-up, then {_N_FILL} steps of: "
        f"Coast / TDOA / FDOA / TDOA+FDOA",
    )
    print(f"{'=' * w}")

    print(
        f"{'':>4}  {'':>5}  "
        f"{'--- Coast ---':^21}  "
        f"{'--- TDOA ---':^21}  "
        f"{'--- FDOA ---':^21}  "
        f"{'-- TDOA+FDOA --':^21}",
    )
    print(
        f"{'Step':>4}  {'Time':>5}  "
        f"{'E':>5} {'N':>5} {'U':>5} {'SMA':>5}  "
        f"{'E':>5} {'N':>5} {'U':>5} {'Vol':>5}  "
        f"{'E':>5} {'N':>5} {'U':>5} {'Vol':>5}  "
        f"{'E':>5} {'N':>5} {'U':>5} {'Vol':>5}",
    )
    print(f"{'-' * w}")

    coast_info = infos["Coast"]
    for i in range(_N_TOTAL):
        t = times[i]
        is_fill = i >= _N_ELL
        ci = coast_info[i]

        if not is_fill:
            print(
                f"{i:4d}  {t:5.0f}  "
                f"{ci.std_e:5.0f} {ci.std_n:5.0f} {ci.std_u:5.0f} "
                f"{ci.sma_95:5.0f}  "
                f"{'':>5} {'':>5} {'':>5} {'':>5}  "
                f"{'':>5} {'':>5} {'':>5} {'':>5}  "
                f"{'':>5} {'':>5} {'':>5} {'':>5}  "
                f"Ell",
            )
        else:
            parts = [f"{i:4d}  {t:5.0f}  "]
            parts.append(
                f"{ci.std_e:5.0f} {ci.std_n:5.0f} {ci.std_u:5.0f} {ci.sma_95:5.0f}  ",
            )
            for strategy in ["TDOA", "FDOA", "BOTH"]:
                si = infos[strategy][i]
                if si.cov_det > 0:
                    vr = (ci.cov_det / si.cov_det) ** 0.5
                else:
                    vr = float("inf")
                parts.append(
                    f"{si.std_e:5.0f} {si.std_n:5.0f} {si.std_u:5.0f} {vr:4.0f}x  ",
                )
            parts.append("<--")
            print("".join(parts))

    print("\nAll std devs in meters (1-sigma ENU). Vol ratio vs coast.")
    print("Same 2 LEO collectors for all obs types:")
    print("  LEO 1: ~az 190, el 20 (southern), heading south")
    print("  LEO 2: ~az 80, el 15 (eastern), heading east")
    print(
        f"TDOA: {_TDOA_STD * 1e9:.0f} ns | "
        f"FDOA: {_FDOA_STD} Hz @ {_FREQ_HZ / 1e9:.0f} GHz | "
        f"Ell: {SIGMA:.0f} m isotropic",
    )

    return _ScenarioResult(
        title=title,
        start=infos["Coast"][_N_ELL - 1],
        final={s: infos[s][-1] for s in _STRATEGIES},
    )


def _plot_comparison(cv: _ScenarioResult, ct: _ScenarioResult) -> None:
    """Bar chart comparing E/N/U std devs across strategies for CV and CT."""
    categories = ["Start", "Coast", "TDOA", "FDOA", "TDOA+FDOA"]
    axis_names = ["East", "North", "Up"]
    axis_colors = ["#636EFA", "#00CC96", "#FFA15A"]

    fig = make_subplots(
        rows=1,
        cols=2,
        subplot_titles=[cv.title, ct.title],
        horizontal_spacing=0.08,
    )

    for col, scenario in enumerate([cv, ct], start=1):
        vals = {
            "Start": [
                scenario.start.std_e,
                scenario.start.std_n,
                scenario.start.std_u,
            ],
            "Coast": [
                scenario.final["Coast"].std_e,
                scenario.final["Coast"].std_n,
                scenario.final["Coast"].std_u,
            ],
            "TDOA": [
                scenario.final["TDOA"].std_e,
                scenario.final["TDOA"].std_n,
                scenario.final["TDOA"].std_u,
            ],
            "FDOA": [
                scenario.final["FDOA"].std_e,
                scenario.final["FDOA"].std_n,
                scenario.final["FDOA"].std_u,
            ],
            "TDOA+FDOA": [
                scenario.final["BOTH"].std_e,
                scenario.final["BOTH"].std_n,
                scenario.final["BOTH"].std_u,
            ],
        }

        for axis_idx, (axis_name, color) in enumerate(
            zip(axis_names, axis_colors, strict=True),
        ):
            y = [vals[cat][axis_idx] for cat in categories]
            fig.add_trace(
                go.Bar(
                    name=axis_name,
                    x=categories,
                    y=y,
                    marker_color=color,
                    showlegend=(col == 1),
                    legendgroup=axis_name,
                    text=[f"{v:.0f}" for v in y],
                    textposition="outside",
                    textfont_size=10,
                ),
                row=1,
                col=col,
            )

    fig.update_layout(
        template="plotly_dark",
        title=(
            "Position Uncertainty After 4 Fill-in Steps "
            "(1-sigma ENU, meters)"
            "<br><sup>Same 2 LEO collectors | "
            f"TDOA: {_TDOA_STD * 1e9:.0f} ns | "
            f"FDOA: {_FDOA_STD} Hz @ {_FREQ_HZ / 1e9:.0f} GHz | "
            f"Ell: {SIGMA:.0f} m isotropic</sup>"
        ),
        barmode="group",
        bargap=0.2,
        bargroupgap=0.05,
        height=500,
        width=1100,
    )

    fig.update_yaxes(title_text="1-sigma std dev (m)", row=1, col=1)
    fig.update_yaxes(title_text="1-sigma std dev (m)", row=1, col=2)

    fig.show()


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


def _generate_cv_truth() -> tuple[np.ndarray, np.ndarray]:
    """Straight-line truth at ~14 m/s NE.

    At X0 = [6.4e6, 0, 0]: E=+Y, N=+Z, U=+X in ECEF.
    Velocity [0, 10, 10] ECEF = [10, 10, 0] ENU = ~14 m/s NE.
    """
    vel = np.array([0.0, 10.0, 10.0])
    positions = np.array([X0 + vel * i * DT for i in range(_N_TOTAL)])
    velocities = np.tile(vel, (_N_TOTAL, 1))
    return positions, velocities


def _generate_ct_truth() -> tuple[np.ndarray, np.ndarray]:
    """Turning truth at ~14 m/s in the E-N plane, omega=0.003 rad/s.

    Starting heading is East, turning toward North.
    """
    speed = 14.0
    omega = 0.003
    positions = []
    velocities = []
    for i in range(_N_TOTAL):
        angle = omega * i * DT
        dy = speed / omega * np.sin(angle)
        dz = speed / omega * (1 - np.cos(angle))
        positions.append(X0 + np.array([0.0, dy, dz]))
        vy = speed * np.cos(angle)
        vz = speed * np.sin(angle)
        velocities.append(np.array([0.0, vy, vz]))
    return np.array(positions), np.array(velocities)


def main() -> None:
    """Run observable fill-in demos for CV and CT tracks."""
    cv_pos, cv_vel = _generate_cv_truth()
    ct_pos, ct_vel = _generate_ct_truth()

    cv_result = _run_scenario(
        truth=cv_pos,
        truth_vel=cv_vel,
        rng_seed=42,
        model_factory=lambda: ConstantVelocity(process_noise_m_s2=1.0),
        title="Constant Velocity (~14 m/s NE)",
    )

    ct_result = _run_scenario(
        truth=ct_pos,
        truth_vel=ct_vel,
        rng_seed=58,
        model_factory=lambda: CoordinatedTurn(
            process_noise_m_s2=0.5,
            process_noise_omega_rad_s2=1e-5,
        ),
        title="Coordinated Turn (~14 m/s, E->N)",
    )

    _plot_comparison(cv_result, ct_result)


if __name__ == "__main__":
    main()
