# ruff: noqa: T201
"""Visualizations for hybrid TDOA+FDOA+AOA geolocation Monte Carlo validation.

Demonstrates combined geolocation using TDOA, FDOA, and AOA:
- TDOA (Time Difference of Arrival): Range-difference from HEO satellites
- FDOA (Frequency Difference of Arrival): Doppler-based from moving LEO satellites
- AOA (Angle of Arrival): Angular bearings from LEO satellites

For FDOA, the target is assumed stationary (velocity=0) and the LEO collectors
have realistic orbital velocities (~7.5 km/s).

Shows comparison of:
- Individual observables (TDOA, FDOA, AOA) with similar ~1000m RMS
- Full combined TDOA+FDOA+AOA (6 TDOA pairs + 6 FDOA pairs + 3 AOA)
- Degenerate combined (2 TDOA + 2 FDOA + 2 AOA)
- Minimal combined (1 TDOA + 1 FDOA + 2 AOA)

Noise levels calibrated for comparable RMS errors (~1000m):
- TDOA: 1 μs
- FDOA: 5 Hz (at 1 GHz carrier)
- AOA: 0.5 mrad

Run with:
    python -m gri_sandbox.plots.hybrid_fdoa_monte_carlo
"""

from __future__ import annotations

import numpy as np
import plotly.graph_objects as go
from gri_geosim.geolocation import (
    aoa_locator,
    tdoa_aoa_fdoa_locator,
    tdoa_locator,
)
from gri_utils.observables import get_aoa_quat, get_fdoa, get_tdoa
from scipy.spatial.transform import Rotation
from scipy.stats import gaussian_kde

EARTH_RADIUS = 6378137.0
LEO_ALTITUDE = 550_000.0  # 550 km LEO altitude (Starlink-like)
HEO_ALTITUDE = 20_000_000.0  # 20,000 km HEO altitude (GPS-like)
LEO_VELOCITY = 7500.0  # Approximate LEO orbital velocity (m/s)

DARK_TEMPLATE = "plotly_dark"


def _leo_velocity(position: np.ndarray, inclination_deg: float) -> np.ndarray:
    """Compute LEO orbital velocity perpendicular to position vector."""
    r_hat = position / np.linalg.norm(position)
    incl = np.radians(inclination_deg)

    if abs(r_hat[2]) < 0.9:  # noqa: PLR2004
        perp = np.cross(r_hat, np.array([0.0, 0.0, 1.0]))
    else:
        perp = np.cross(r_hat, np.array([1.0, 0.0, 0.0]))
    perp = perp / np.linalg.norm(perp)

    v_hat = perp * np.cos(incl) + np.cross(r_hat, perp) * np.sin(incl)
    return v_hat * LEO_VELOCITY


def _satellite_position(
    lat_deg: float,
    lon_deg: float,
    altitude: float,
) -> np.ndarray:
    """Compute satellite position from lat/lon over Earth at given altitude."""
    lat = np.radians(lat_deg)
    lon = np.radians(lon_deg)
    r = EARTH_RADIUS + altitude

    return np.array(
        [
            r * np.cos(lat) * np.cos(lon),
            r * np.cos(lat) * np.sin(lon),
            r * np.sin(lat),
        ],
    )


def _ground_target_position(lat_deg: float, lon_deg: float) -> np.ndarray:
    """Compute ground target position from lat/lon."""
    lat = np.radians(lat_deg)
    lon = np.radians(lon_deg)

    return np.array(
        [
            EARTH_RADIUS * np.cos(lat) * np.cos(lon),
            EARTH_RADIUS * np.cos(lat) * np.sin(lon),
            EARTH_RADIUS * np.sin(lat),
        ],
    )


def _create_geometry() -> dict:
    """Create geometry for full TDOA+FDOA+AOA scenario.

    Returns:
        Dictionary with all geometry parameters.
    """
    # Target on Earth's surface (stationary)
    target_xyz = _ground_target_position(lat_deg=30.0, lon_deg=0.0)
    target_vel = np.zeros(3)

    # 4 TDOA HEO satellites -> 6 pairs
    tdoa_satellites = np.array(
        [
            _satellite_position(45.0, -5.0, HEO_ALTITUDE),
            _satellite_position(-20.0, -15.0, HEO_ALTITUDE),
            _satellite_position(30.0, 20.0, HEO_ALTITUDE),
            _satellite_position(-35.0, 5.0, HEO_ALTITUDE),
        ],
    )
    pairs = [(0, 1), (0, 2), (0, 3), (1, 2), (1, 3), (2, 3)]
    tdoa_collector1_xyz = np.array([tdoa_satellites[i] for i, _ in pairs])
    tdoa_collector2_xyz = np.array([tdoa_satellites[j] for _, j in pairs])

    # 4 FDOA LEO satellites with orbital velocities -> 6 pairs
    fdoa_sat_configs = [
        (40.0, -10.0, 0.0),
        (25.0, 15.0, 30.0),
        (35.0, 5.0, 60.0),
        (-20.0, -5.0, 90.0),
    ]
    fdoa_satellites = np.array(
        [
            _satellite_position(lat, lon, LEO_ALTITUDE)
            for lat, lon, _ in fdoa_sat_configs
        ],
    )
    fdoa_velocities = np.array(
        [
            _leo_velocity(_satellite_position(lat, lon, LEO_ALTITUDE), incl)
            for lat, lon, incl in fdoa_sat_configs
        ],
    )
    fdoa_collector1_xyz = np.array([fdoa_satellites[i] for i, _ in pairs])
    fdoa_collector2_xyz = np.array([fdoa_satellites[j] for _, j in pairs])
    fdoa_collector1_vel = np.array([fdoa_velocities[i] for i, _ in pairs])
    fdoa_collector2_vel = np.array([fdoa_velocities[j] for _, j in pairs])

    # 3 AOA LEO satellites
    aoa_positions = [(40.0, -10.0), (25.0, 15.0), (35.0, 5.0)]
    aoa_collector_xyz = np.array(
        [_satellite_position(lat, lon, LEO_ALTITUDE) for lat, lon in aoa_positions],
    )
    aoa_rotations = Rotation.identity(3)

    return {
        "target_xyz": target_xyz,
        "target_vel": target_vel,
        # TDOA (full)
        "tdoa_collector1_xyz": tdoa_collector1_xyz,
        "tdoa_collector2_xyz": tdoa_collector2_xyz,
        # FDOA (full)
        "fdoa_collector1_xyz": fdoa_collector1_xyz,
        "fdoa_collector2_xyz": fdoa_collector2_xyz,
        "fdoa_collector1_vel": fdoa_collector1_vel,
        "fdoa_collector2_vel": fdoa_collector2_vel,
        # AOA (full)
        "aoa_collector_xyz": aoa_collector_xyz,
        "aoa_rotations": aoa_rotations,
    }


def run_monte_carlo(  # noqa: PLR0913
    geometry: dict,
    tdoa_sigma: float = 1e-6,
    fdoa_sigma: float = 5.0,
    aoa_sigma: float = 0.0005,
    freq_hz: float = 1e9,
    n_trials: int = 500,
    seed: int = 42,
) -> dict:
    """Run Monte Carlo for all scenarios.

    Args:
        geometry: Dictionary with collector positions, velocities, and target.
        tdoa_sigma: TDOA measurement noise (1-sigma, seconds).
        fdoa_sigma: FDOA measurement noise (1-sigma, Hz).
        aoa_sigma: AOA measurement noise (1-sigma, radians).
        freq_hz: Carrier frequency in Hz for FDOA.
        n_trials: Number of Monte Carlo trials.
        seed: Random seed for reproducibility.

    Returns:
        Dictionary with errors for each solver configuration.
    """
    rng = np.random.default_rng(seed)

    target_true = geometry["target_xyz"]
    target_vel = geometry["target_vel"]

    # Full TDOA setup (6 pairs)
    tdoa_c1 = geometry["tdoa_collector1_xyz"]
    tdoa_c2 = geometry["tdoa_collector2_xyz"]
    n_tdoa_full = len(tdoa_c1)
    true_tdoa_full = np.array(
        [get_tdoa(target_true, tdoa_c1[i], tdoa_c2[i]) for i in range(n_tdoa_full)],
    )

    # Full FDOA setup (6 pairs)
    fdoa_c1 = geometry["fdoa_collector1_xyz"]
    fdoa_c2 = geometry["fdoa_collector2_xyz"]
    fdoa_v1 = geometry["fdoa_collector1_vel"]
    fdoa_v2 = geometry["fdoa_collector2_vel"]
    n_fdoa_full = len(fdoa_c1)
    true_fdoa_full = np.array(
        [
            get_fdoa(
                target_true,
                target_vel,
                fdoa_c1[i],
                fdoa_c2[i],
                freq_hz,
                collector1_vel=fdoa_v1[i],
                collector2_vel=fdoa_v2[i],
            )
            for i in range(n_fdoa_full)
        ],
    )

    # Full AOA setup (3 collectors)
    aoa_xyz = geometry["aoa_collector_xyz"]
    aoa_rot = geometry["aoa_rotations"]
    n_aoa_full = len(aoa_xyz)
    true_aoa_full = get_aoa_quat(aoa_xyz, target_true, aoa_rot)
    aoa_cov = np.eye(2) * aoa_sigma**2

    # Degenerate subsets
    n_tdoa_min = 1
    n_fdoa_min = 1
    n_aoa_min = 2

    n_tdoa_mid = 2
    n_fdoa_mid = 2
    n_aoa_mid = 2

    results = {
        "tdoa_only": [],
        "fdoa_only": [],
        "aoa_only": [],
        "hybrid_full": [],
        "hybrid_2_2_2": [],
        "hybrid_1_1_2": [],
    }

    for _ in range(n_trials):
        # Generate noise for full sets
        noise_tdoa = rng.standard_normal(n_tdoa_full) * tdoa_sigma
        noise_fdoa = rng.standard_normal(n_fdoa_full) * fdoa_sigma
        noise_aoa = rng.standard_normal((n_aoa_full, 2)) * aoa_sigma

        noisy_tdoa_full = true_tdoa_full + noise_tdoa
        noisy_fdoa_full = true_fdoa_full + noise_fdoa
        noisy_aoa_full = true_aoa_full + noise_aoa

        # TDOA-only (6 pairs)
        result = tdoa_locator(
            tdoa_c1,
            tdoa_c2,
            noisy_tdoa_full,
            np.full(n_tdoa_full, tdoa_sigma**2),
        )
        results["tdoa_only"].append(np.linalg.norm(result.position - target_true))

        # FDOA-only (6 pairs)
        result = tdoa_aoa_fdoa_locator(
            fdoa_collector1_xyz=fdoa_c1,
            fdoa_collector2_xyz=fdoa_c2,
            fdoa_measurements=noisy_fdoa_full,
            fdoa_variances=np.full(n_fdoa_full, fdoa_sigma**2),
            freq_hz=freq_hz,
            fdoa_collector1_vel=fdoa_v1,
            fdoa_collector2_vel=fdoa_v2,
            fix_velocity=target_vel,
        )
        results["fdoa_only"].append(np.linalg.norm(result.position - target_true))

        # AOA-only (3 collectors)
        result = aoa_locator(
            aoa_xyz,
            aoa_rot,
            noisy_aoa_full,
            np.tile(aoa_cov, (n_aoa_full, 1, 1)),
        )
        results["aoa_only"].append(np.linalg.norm(result.position - target_true))

        # Full hybrid (6 TDOA + 6 FDOA + 3 AOA)
        result = tdoa_aoa_fdoa_locator(
            fdoa_collector1_xyz=fdoa_c1,
            fdoa_collector2_xyz=fdoa_c2,
            fdoa_measurements=noisy_fdoa_full,
            fdoa_variances=np.full(n_fdoa_full, fdoa_sigma**2),
            freq_hz=freq_hz,
            fdoa_collector1_vel=fdoa_v1,
            fdoa_collector2_vel=fdoa_v2,
            tdoa_collector1_xyz=tdoa_c1,
            tdoa_collector2_xyz=tdoa_c2,
            tdoa_measurements=noisy_tdoa_full,
            tdoa_variances=np.full(n_tdoa_full, tdoa_sigma**2),
            aoa_collector_xyz=aoa_xyz,
            aoa_rotations=aoa_rot,
            aoa_measurements=noisy_aoa_full,
            aoa_covariances=np.tile(aoa_cov, (n_aoa_full, 1, 1)),
            fix_velocity=target_vel,
        )
        results["hybrid_full"].append(np.linalg.norm(result.position - target_true))

        # Mid hybrid (2 TDOA + 2 FDOA + 2 AOA)
        result = tdoa_aoa_fdoa_locator(
            fdoa_collector1_xyz=fdoa_c1[:n_fdoa_mid],
            fdoa_collector2_xyz=fdoa_c2[:n_fdoa_mid],
            fdoa_measurements=noisy_fdoa_full[:n_fdoa_mid],
            fdoa_variances=np.full(n_fdoa_mid, fdoa_sigma**2),
            freq_hz=freq_hz,
            fdoa_collector1_vel=fdoa_v1[:n_fdoa_mid],
            fdoa_collector2_vel=fdoa_v2[:n_fdoa_mid],
            tdoa_collector1_xyz=tdoa_c1[:n_tdoa_mid],
            tdoa_collector2_xyz=tdoa_c2[:n_tdoa_mid],
            tdoa_measurements=noisy_tdoa_full[:n_tdoa_mid],
            tdoa_variances=np.full(n_tdoa_mid, tdoa_sigma**2),
            aoa_collector_xyz=aoa_xyz[:n_aoa_mid],
            aoa_rotations=Rotation.from_quat(aoa_rot.as_quat()[:n_aoa_mid]),
            aoa_measurements=noisy_aoa_full[:n_aoa_mid],
            aoa_covariances=np.tile(aoa_cov, (n_aoa_mid, 1, 1)),
            fix_velocity=target_vel,
        )
        results["hybrid_2_2_2"].append(np.linalg.norm(result.position - target_true))

        # Minimal hybrid (1 TDOA + 1 FDOA + 2 AOA)
        result = tdoa_aoa_fdoa_locator(
            fdoa_collector1_xyz=fdoa_c1[:n_fdoa_min],
            fdoa_collector2_xyz=fdoa_c2[:n_fdoa_min],
            fdoa_measurements=noisy_fdoa_full[:n_fdoa_min],
            fdoa_variances=np.full(n_fdoa_min, fdoa_sigma**2),
            freq_hz=freq_hz,
            fdoa_collector1_vel=fdoa_v1[:n_fdoa_min],
            fdoa_collector2_vel=fdoa_v2[:n_fdoa_min],
            tdoa_collector1_xyz=tdoa_c1[:n_tdoa_min],
            tdoa_collector2_xyz=tdoa_c2[:n_tdoa_min],
            tdoa_measurements=noisy_tdoa_full[:n_tdoa_min],
            tdoa_variances=np.full(n_tdoa_min, tdoa_sigma**2),
            aoa_collector_xyz=aoa_xyz[:n_aoa_min],
            aoa_rotations=Rotation.from_quat(aoa_rot.as_quat()[:n_aoa_min]),
            aoa_measurements=noisy_aoa_full[:n_aoa_min],
            aoa_covariances=np.tile(aoa_cov, (n_aoa_min, 1, 1)),
            fix_velocity=target_vel,
        )
        results["hybrid_1_1_2"].append(np.linalg.norm(result.position - target_true))

    # Convert to arrays
    return {key: np.array(val) for key, val in results.items()}


def _kde_line(data: np.ndarray, n_points: int = 200) -> tuple[np.ndarray, np.ndarray]:
    """Compute KDE for smooth density estimation."""
    kde = gaussian_kde(data)
    x_min, x_max = data.min(), data.max()
    padding = (x_max - x_min) * 0.1
    x = np.linspace(max(0, x_min - padding), x_max + padding, n_points)
    return x, kde(x)


def plot_error_comparison(results: dict, n_trials: int) -> go.Figure:
    """Plot position error comparison for all scenarios.

    Args:
        results: Results from run_monte_carlo.
        n_trials: Number of trials (for title).

    Returns:
        Plotly figure with KDE error distributions.
    """
    fig = go.Figure()

    # Individual observables (dotted, similar accuracy)
    individual_configs = [
        ("tdoa_only", "TDOA (6 pairs)", "#EF553B"),
        ("fdoa_only", "FDOA (6 pairs)", "#00CC96"),
        ("aoa_only", "AOA (3 sats)", "#636EFA"),
    ]

    for key, label, color in individual_configs:
        errors = results[key]
        mean_err = np.mean(errors)
        x, density = _kde_line(errors)
        fig.add_trace(
            go.Scatter(
                x=x,
                y=density,
                mode="lines",
                name=f"{label} (μ={mean_err:.0f}m)",
                line={"color": color, "width": 2, "dash": "dot"},
            ),
        )

    # Hybrid full (solid, best accuracy)
    errors = results["hybrid_full"]
    mean_err = np.mean(errors)
    x, density = _kde_line(errors)
    fig.add_trace(
        go.Scatter(
            x=x,
            y=density,
            mode="lines",
            name=f"Full TDOA+FDOA+AOA (μ={mean_err:.0f}m)",
            line={"color": "#AB63FA", "width": 3},
        ),
    )

    # Hybrid 2+2+2 (long dash)
    errors = results["hybrid_2_2_2"]
    mean_err = np.mean(errors)
    x, density = _kde_line(errors)
    fig.add_trace(
        go.Scatter(
            x=x,
            y=density,
            mode="lines",
            name=f"Hybrid 2+2+2 (μ={mean_err:.0f}m)",
            line={"color": "#FFA15A", "width": 2.5, "dash": "longdash"},
        ),
    )

    # Hybrid 1+1+2 (short dash)
    errors = results["hybrid_1_1_2"]
    mean_err = np.mean(errors)
    x, density = _kde_line(errors)
    fig.add_trace(
        go.Scatter(
            x=x,
            y=density,
            mode="lines",
            name=f"Hybrid 1+1+2 (μ={mean_err:.0f}m)",
            line={"color": "#19D3F3", "width": 2.5, "dash": "dash"},
        ),
    )

    fig.update_layout(
        template=DARK_TEMPLATE,
        title=f"Position Errors: Individual vs Combined (n={n_trials})",
        xaxis_title="Position Error (m)",
        yaxis_title="Density",
        height=500,
        legend={
            "x": 0.98,
            "y": 0.98,
            "xanchor": "right",
            "yanchor": "top",
            "bgcolor": "rgba(0,0,0,0.5)",
        },
    )

    return fig


def main() -> None:
    """Run hybrid TDOA+FDOA+AOA Monte Carlo visualizations."""
    print("=" * 70)
    print("Hybrid TDOA+FDOA+AOA Monte Carlo Validation")
    print("=" * 70)

    n_trials = 1000
    geometry = _create_geometry()

    # Noise levels calibrated for ~1000m RMS
    tdoa_sigma = 1e-6  # 1 μs
    fdoa_sigma = 5.0  # 5 Hz at 1 GHz
    aoa_sigma = 0.0005  # 0.5 mrad
    freq_hz = 1e9

    print("\nNoise levels (calibrated for ~1000m RMS):")
    print(f"  TDOA: {tdoa_sigma * 1e6:.1f} μs")
    print(f"  FDOA: {fdoa_sigma:.1f} Hz at {freq_hz / 1e9:.0f} GHz")
    print(f"  AOA:  {aoa_sigma * 1000:.1f} mrad")
    print("  Target velocity: 0 (stationary)")

    # Run Monte Carlo
    print(f"\nRunning Monte Carlo ({n_trials} trials)...")
    results = run_monte_carlo(
        geometry,
        tdoa_sigma=tdoa_sigma,
        fdoa_sigma=fdoa_sigma,
        aoa_sigma=aoa_sigma,
        freq_hz=freq_hz,
        n_trials=n_trials,
    )

    # Print results
    print("\nRESULTS")
    print("-" * 50)
    labels = [
        ("tdoa_only", "TDOA only (6 HEO pairs)"),
        ("fdoa_only", "FDOA only (6 LEO pairs)"),
        ("aoa_only", "AOA only (3 LEO sats)"),
        ("hybrid_full", "Full TDOA+FDOA+AOA (6+6+3)"),
        ("hybrid_2_2_2", "Hybrid (2+2+2)"),
        ("hybrid_1_1_2", "Hybrid (1+1+2)"),
    ]
    for key, label in labels:
        mean_err = np.mean(results[key])
        print(f"  {label}: {mean_err:.0f}m")

    # Improvement summary
    print("\nIMPROVEMENT SUMMARY")
    print("-" * 50)
    baseline = np.mean(
        [
            np.mean(results["tdoa_only"]),
            np.mean(results["fdoa_only"]),
            np.mean(results["aoa_only"]),
        ],
    )
    full_mean = np.mean(results["hybrid_full"])
    mid_mean = np.mean(results["hybrid_2_2_2"])
    min_mean = np.mean(results["hybrid_1_1_2"])

    full_pct = (baseline - full_mean) / baseline * 100
    mid_pct = (baseline - mid_mean) / baseline * 100
    min_pct = (baseline - min_mean) / baseline * 100
    print(f"  Individual average: {baseline:.0f}m")
    print(f"  Full hybrid (6+6+3): {full_mean:.0f}m ({full_pct:.0f}% better)")
    print(f"  Hybrid (2+2+2): {mid_mean:.0f}m ({mid_pct:.0f}% better)")
    print(f"  Hybrid (1+1+2): {min_mean:.0f}m ({min_pct:.0f}% better)")

    # Generate plot
    print("\nGenerating plot...")
    fig = plot_error_comparison(results, n_trials)
    fig.show()


if __name__ == "__main__":
    main()
