# ruff: noqa: T201
"""Demo: observable isochrones on flat and terrain surfaces.

Shows five use cases:
1. TDOA isochrones from 4 collectors on a satellite map
2. TDOA isochrones with constant-width sigma bands
3. TDOA + FDOA + AOA (no sigma) on a flat plane with elevated collectors
4. TDOA + FDOA + AOA with uncertainty bands/ellipses
5. Isochrones on a terrain-like surface with arbitrary altitude plane

Run: python -m gri_sandbox.isochrone_demo.main
"""

from __future__ import annotations

import numpy as np
from gri_plot import FigureMap, MapStyle
from gri_utils.constants import C
from gri_utils.conversion import lla_to_xyz, xyz_to_lla
from gri_utils.observables import (
    get_aoa_enu,
    get_aoa_residual,
    get_fdoa_residual,
    get_tdoa_residual,
)
from gri_utils.terrain import intersect_sheet, lla_to_xyz_sheet


def _make_scenario() -> tuple[np.ndarray, list[tuple[int, int, float]], np.ndarray]:
    """Create a 4-collector TDOA scenario.

    Returns:
        Tuple of (collectors_xyz, tdoa_pairs, emitter_xyz).
    """
    # 4 collectors in a ~50km diamond around Denver area
    collector_llas = np.array(
        [
            [39.3, -105.0, 0.0],  # North
            [38.7, -105.0, 0.0],  # South
            [39.0, -105.4, 0.0],  # West
            [39.0, -104.6, 0.0],  # East
        ],
    )
    collectors_xyz = np.array([lla_to_xyz(lla) for lla in collector_llas])

    # Emitter slightly off-center
    emitter_lla = np.array([39.05, -104.85, 0.0])
    emitter_xyz = lla_to_xyz(emitter_lla)

    # Compute TDOAs: use collector 0 (North) as reference
    ranges = np.linalg.norm(collectors_xyz - emitter_xyz, axis=1)
    tdoa_pairs = [
        (0, 1, float((ranges[0] - ranges[1]) / C)),  # North-South
        (0, 2, float((ranges[0] - ranges[2]) / C)),  # North-West
        (0, 3, float((ranges[0] - ranges[3]) / C)),  # North-East
    ]

    return collectors_xyz, tdoa_pairs, emitter_xyz


def _make_sheet(
    center_lat: float = 39.0,
    center_lon: float = -105.0,
    span_deg: float = 0.8,
    n_points: int = 500,
    altitude_m: float = 0.0,
) -> np.ndarray:
    """Create a flat XYZ sheet at constant altitude."""
    lats = np.linspace(
        center_lat - span_deg / 2,
        center_lat + span_deg / 2,
        n_points,
    )
    lons = np.linspace(
        center_lon - span_deg / 2,
        center_lon + span_deg / 2,
        n_points,
    )
    alts = np.full((n_points, n_points), altitude_m)
    return lla_to_xyz_sheet(lats, lons, alts)


def _make_terrain_sheet(
    center_lat: float = 39.0,
    center_lon: float = -105.0,
    span_deg: float = 0.8,
    n_points: int = 300,
) -> np.ndarray:
    """Create a terrain-like XYZ sheet with smoothed random elevation.

    Generates a pseudo-terrain surface by layering multiple octaves of
    smoothed random noise, producing rolling hills that look like DTED.

    Returns:
        XYZ sheet in ECEF meters, shape (n_points, n_points, 3).
    """
    rng = np.random.default_rng(42)

    alts = np.zeros((n_points, n_points))

    # Layer octaves of smoothed noise for terrain-like appearance
    for scale, amplitude in [(8, 1500), (16, 700), (32, 300), (64, 100)]:
        # Generate small random grid and upscale via interpolation
        small = rng.standard_normal((scale, scale))
        from scipy.ndimage import zoom  # noqa: PLC0415

        layer = zoom(small, n_points / scale, order=3)
        alts += layer[:n_points, :n_points] * amplitude  # type: ignore[reportOperatorIssue]

    # Shift so minimum is ~500m (realistic terrain base)
    alts -= alts.min()
    alts += 500.0

    lats = np.linspace(
        center_lat - span_deg / 2,
        center_lat + span_deg / 2,
        n_points,
    )
    lons = np.linspace(
        center_lon - span_deg / 2,
        center_lon + span_deg / 2,
        n_points,
    )
    return lla_to_xyz_sheet(lats, lons, alts)


# ---- Demo 1: TDOA on satellite map ----


def demo_tdoa_satellite() -> None:
    """TDOA isochrones on satellite map."""
    print("--- Demo 1: TDOA Isochrones (Satellite Map) ---")

    collectors_xyz, tdoa_pairs, emitter_xyz = _make_scenario()
    sheet = _make_sheet()

    fig = FigureMap(title="TDOA Isochrones (Satellite)")
    labels = ["TDOA N-S", "TDOA N-W", "TDOA N-E"]

    for i, (c1_idx, c2_idx, tdoa_s) in enumerate(tdoa_pairs):
        residual = get_tdoa_residual(
            collectors_xyz[c1_idx],
            collectors_xyz[c2_idx],
            tdoa_s,
        )
        contours = intersect_sheet(sheet, residual)
        print(f"  {labels[i]}: {len(contours)} contour(s)")
        fig.add_contours(contours, label=labels[i])

    # Add markers
    collector_labels = ["C1 (N)", "C2 (S)", "C3 (W)", "C4 (E)"]
    fig.add_markers(collectors_xyz, labels=collector_labels, color="white")

    emitter_lla = xyz_to_lla(emitter_xyz.reshape(1, 3))
    fig.add_markers(emitter_xyz, labels=["Emitter"], color="#ff0000", size=12)

    print(f"  Emitter at: {emitter_lla[0, 0]:.3f}, {emitter_lla[0, 1]:.3f}")
    fig.show()


# ---- Demo 2: TDOA with sigma bands ----


def demo_tdoa_bands() -> None:
    """TDOA isochrones with sigma uncertainty bands."""
    print("--- Demo 2: TDOA with Sigma Bands (Satellite) ---")

    collectors_xyz, tdoa_pairs, emitter_xyz = _make_scenario()
    sheet = _make_sheet()

    fig = FigureMap(title="TDOA Isochrones with Uncertainty Bands")
    labels = ["TDOA N-S", "TDOA N-W", "TDOA N-E"]

    # Assume ~1km uncertainty width for visualization
    band_width_m = 1000.0

    for i, (c1_idx, c2_idx, tdoa_s) in enumerate(tdoa_pairs):
        residual = get_tdoa_residual(
            collectors_xyz[c1_idx],
            collectors_xyz[c2_idx],
            tdoa_s,
        )
        contours = intersect_sheet(sheet, residual)

        # Band first (renders behind), then center line
        fig.add_band(contours, band_width_m)
        fig.add_contours(contours, label=labels[i])

    collector_labels = ["C1 (N)", "C2 (S)", "C3 (W)", "C4 (E)"]
    fig.add_markers(collectors_xyz, labels=collector_labels, color="white")
    fig.add_markers(emitter_xyz, labels=["Emitter"], color="#ff0000", size=12)

    fig.show()


# ---- Demo 3: TDOA + FDOA + AOA (no sigma) ----


def _make_mixed_scenario() -> dict:
    """Create a scenario with elevated collectors for TDOA/FDOA/AOA.

    Collectors are at altitude (aircraft-like). One pair has velocities
    for FDOA. One has an AOA measurement at az~45, el~45.

    Returns:
        Dict with all scenario elements.
    """
    # Three collectors at aircraft altitude (~10km)
    c1_lla = np.array([39.2, -105.2, 10000.0])  # NW, for TDOA ref
    c2_lla = np.array([38.8, -104.8, 10000.0])  # SE, for TDOA pair
    c3_lla = np.array([39.0, -105.5, 12000.0])  # W, for AOA

    c1_xyz = lla_to_xyz(c1_lla)
    c2_xyz = lla_to_xyz(c2_lla)
    c3_xyz = lla_to_xyz(c3_lla)
    collectors_xyz = np.array([c1_xyz, c2_xyz, c3_xyz])

    # Emitter on the ground
    emitter_lla = np.array([39.0, -105.0, 0.0])
    emitter_xyz = lla_to_xyz(emitter_lla)

    # TDOA between c1 and c2
    r1 = np.linalg.norm(c1_xyz - emitter_xyz)
    r2 = np.linalg.norm(c2_xyz - emitter_xyz)
    tdoa_s = (r1 - r2) / C

    # FDOA: collectors moving in different directions for interesting curve
    # c1 moving east at 200 m/s, c2 moving north at 200 m/s
    c1_vel_enu = np.array([200.0, 0.0, 0.0])  # East
    c2_vel_enu = np.array([0.0, 200.0, 0.0])  # North
    from gri_utils.conversion import enu_to_xyz  # noqa: PLC0415

    c1_vel_xyz = enu_to_xyz(c1_xyz, c1_vel_enu)
    c2_vel_xyz = enu_to_xyz(c2_xyz, c2_vel_enu)

    # Compute true FDOA
    from gri_utils.observables import get_fdoa  # noqa: PLC0415

    freq_hz = 1.0e9  # 1 GHz carrier
    fdoa_hz = float(
        get_fdoa(
            emitter_xyz,
            np.zeros(3),
            c1_xyz,
            c2_xyz,
            freq_hz,
            c1_vel_xyz,
            c2_vel_xyz,
        ),
    )

    # AOA from c3 toward emitter
    aoa_cosines = get_aoa_enu(c3_xyz, emitter_xyz)

    return {
        "collectors_xyz": collectors_xyz,
        "emitter_xyz": emitter_xyz,
        "c1_xyz": c1_xyz,
        "c2_xyz": c2_xyz,
        "c3_xyz": c3_xyz,
        "c1_vel_xyz": c1_vel_xyz,
        "c2_vel_xyz": c2_vel_xyz,
        "tdoa_s": tdoa_s,
        "fdoa_hz": fdoa_hz,
        "freq_hz": freq_hz,
        "aoa_cosines": aoa_cosines,
    }


def demo_mixed_no_sigma() -> None:
    """TDOA + FDOA + AOA on flat plane, no uncertainty."""
    print("--- Demo 3: TDOA + FDOA + AOA (No Sigma) ---")

    s = _make_mixed_scenario()
    origin_xyz = np.mean(s["collectors_xyz"], axis=0)
    sheet = _make_sheet(n_points=500)

    fig = FigureMap(
        style=MapStyle.BLANK,
        frame="enu",
        origin_xyz=origin_xyz,
        title="TDOA + FDOA + AOA Isochrones (No Sigma)",
    )

    # TDOA isochrone (c1-c2)
    tdoa_residual = get_tdoa_residual(s["c1_xyz"], s["c2_xyz"], s["tdoa_s"])
    tdoa_contours = intersect_sheet(sheet, tdoa_residual)
    fig.add_contours(tdoa_contours, label="TDOA")
    print(f"  TDOA: {len(tdoa_contours)} contour(s)")

    # FDOA isochrone (c1-c2 with velocities)
    fdoa_residual = get_fdoa_residual(
        s["c1_xyz"],
        s["c1_vel_xyz"],
        s["c2_xyz"],
        s["c2_vel_xyz"],
        s["fdoa_hz"],
        s["freq_hz"],
    )
    fdoa_contours = intersect_sheet(sheet, fdoa_residual)
    fig.add_contours(fdoa_contours, label="FDOA")
    print(f"  FDOA: {len(fdoa_contours)} contour(s)")

    # AOA: ray from collector to emitter, and intersection point marker
    # The AOA with no sigma is a bearing line. On a flat plane below an
    # elevated collector, the ray pierces the plane at a single point.
    # Show the ray as a thin line from collector down to the plane.
    from gri_utils.observables.aoa_enu import translate_aoa  # noqa: PLC0415

    aoa_far = translate_aoa(
        s["c3_xyz"],
        s["aoa_cosines"],
        80000.0,
        z_positive=False,
    )
    fig.add_line(
        s["c3_xyz"],
        aoa_far,
        color="#00ff00",
        width=1,
        label="AOA ray",
    )

    # Mark intersection point on the ground plane
    fig.add_markers(s["emitter_xyz"], labels=["Emitter"], color="#ff0000", size=12)

    fig.show()


# ---- Demo 4: TDOA + FDOA + AOA with sigma ----


def demo_mixed_with_sigma() -> None:
    """TDOA + FDOA + AOA with uncertainty bands and AOA ellipse."""
    print("--- Demo 4: TDOA + FDOA + AOA (With Sigma) ---")

    s = _make_mixed_scenario()
    origin_xyz = np.mean(s["collectors_xyz"], axis=0)
    sheet = _make_sheet(n_points=500)

    fig = FigureMap(
        style=MapStyle.BLANK,
        frame="enu",
        origin_xyz=origin_xyz,
        title="TDOA + FDOA + AOA with Uncertainty",
    )

    # TDOA with band
    tdoa_residual = get_tdoa_residual(s["c1_xyz"], s["c2_xyz"], s["tdoa_s"])
    tdoa_contours = intersect_sheet(sheet, tdoa_residual)
    fig.add_band(tdoa_contours, 1500.0, label="TDOA 1-sigma")
    fig.add_contours(tdoa_contours, label="TDOA")

    # FDOA with band
    fdoa_residual = get_fdoa_residual(
        s["c1_xyz"],
        s["c1_vel_xyz"],
        s["c2_xyz"],
        s["c2_vel_xyz"],
        s["fdoa_hz"],
        s["freq_hz"],
    )
    fdoa_contours = intersect_sheet(sheet, fdoa_residual)
    fig.add_band(fdoa_contours, 1500.0, label="FDOA 1-sigma")
    fig.add_contours(fdoa_contours, label="FDOA")

    # AOA ellipse: use small angular sigma to find the cone-plane
    # intersection as a closed contour. AOA residual returns angular
    # deviation in radians, so level = sigma_rad gives the 1-sigma cone.
    sigma_deg = 2.0
    sigma_rad = np.radians(sigma_deg)
    aoa_residual = get_aoa_residual(
        s["c3_xyz"],
        s["aoa_cosines"],
        z_positive=False,
    )
    aoa_contours = intersect_sheet(sheet, aoa_residual, level=sigma_rad)
    print(f"  AOA sigma contours: {len(aoa_contours)}")
    for contour in aoa_contours:
        fig.add_filled_contour(contour, label="AOA 1-sigma", alpha=0.3)

    # AOA center ray
    from gri_utils.observables.aoa_enu import translate_aoa  # noqa: PLC0415

    aoa_far = translate_aoa(
        s["c3_xyz"],
        s["aoa_cosines"],
        80000.0,
        z_positive=False,
    )
    fig.add_line(
        s["c3_xyz"],
        aoa_far,
        color="#00ff00",
        width=1,
        label="AOA ray",
    )

    fig.add_markers(s["emitter_xyz"], labels=["Emitter"], color="#ff0000", size=12)

    fig.show()


# ---- Demo 5: Terrain surface intersection ----


def demo_terrain_intersection() -> None:
    """Isochrones on a terrain surface with heatmap + flat plane.

    Uses Range (sphere) and FDOA observables which are altitude-sensitive,
    so terrain vs flat plane contours visibly diverge.
    """
    print("--- Demo 5: Terrain Surface Intersection ---")

    s = _make_mixed_scenario()
    origin_xyz = np.mean(s["collectors_xyz"], axis=0)

    # Create terrain sheet with realistic-looking elevation
    terrain_sheet = _make_terrain_sheet(n_points=300)

    # Also create a flat sheet at 1500m altitude for comparison
    flat_sheet = _make_sheet(n_points=300, altitude_m=1500.0)

    fig = FigureMap(
        style=MapStyle.BLANK,
        frame="enu",
        origin_xyz=origin_xyz,
        title="Isochrones on Terrain Surface vs. Flat Plane (1500m)",
    )

    # Show terrain as heatmap background (no colorbar label to avoid
    # covering legend)
    fig.add_heatmap(terrain_sheet, colorscale="earth", opacity=0.5)

    # Range sphere from c1 -- altitude-sensitive because the sphere
    # intersection with a surface shifts laterally as altitude changes.
    from gri_utils.observables import get_range_residual  # noqa: PLC0415

    range_m = float(np.linalg.norm(s["c1_xyz"] - s["emitter_xyz"]))
    range_residual = get_range_residual(s["c1_xyz"], range_m)
    terrain_range = intersect_sheet(terrain_sheet, range_residual)
    fig.add_contours(terrain_range, label="Range (terrain)", width=3)
    print(f"  Range on terrain: {len(terrain_range)} contour(s)")

    flat_range = intersect_sheet(flat_sheet, range_residual)
    fig.add_contours(flat_range, label="Range (1500m plane)", width=1)
    print(f"  Range on flat: {len(flat_range)} contour(s)")

    # FDOA on terrain
    fdoa_residual = get_fdoa_residual(
        s["c1_xyz"],
        s["c1_vel_xyz"],
        s["c2_xyz"],
        s["c2_vel_xyz"],
        s["fdoa_hz"],
        s["freq_hz"],
    )
    terrain_fdoa = intersect_sheet(terrain_sheet, fdoa_residual)
    fig.add_contours(terrain_fdoa, label="FDOA (terrain)", width=3)
    print(f"  FDOA on terrain: {len(terrain_fdoa)} contour(s)")

    flat_fdoa = intersect_sheet(flat_sheet, fdoa_residual)
    fig.add_contours(flat_fdoa, label="FDOA (1500m plane)", width=1)

    # AOA cone on terrain -- collector at az 45, el 45 from emitter,
    # range 15km (~10.6km altitude, above terrain max of ~8km).
    # 30-degree sigma gives a large footprint on flat and a smaller one
    # on terrain (closer to the collector = smaller cone cross-section).
    from gri_utils.conversion import translate_aer  # noqa: PLC0415

    aoa_collector_xyz = translate_aer(
        s["emitter_xyz"],
        np.array([45.0, 45.0, 15000.0]),
    )
    aoa_cosines = get_aoa_enu(aoa_collector_xyz, s["emitter_xyz"])
    aoa_residual = get_aoa_residual(
        aoa_collector_xyz,
        aoa_cosines,
        z_positive=False,
    )

    sigma_deg = 30.0
    sigma_rad = np.radians(sigma_deg)
    terrain_aoa = intersect_sheet(terrain_sheet, aoa_residual, level=sigma_rad)
    print(f"  AOA on terrain: {len(terrain_aoa)} contour(s)")
    for contour in terrain_aoa:
        fig.add_filled_contour(
            contour,
            color="#ff00ff",
            label="AOA (terrain)",
            alpha=0.4,
        )

    flat_aoa = intersect_sheet(flat_sheet, aoa_residual, level=sigma_rad)
    print(f"  AOA on flat: {len(flat_aoa)} contour(s)")
    for contour in flat_aoa:
        fig.add_filled_contour(
            contour,
            color="#ffff00",
            label="AOA (1500m plane)",
            alpha=0.25,
        )

    fig.show()


if __name__ == "__main__":
    demo_tdoa_satellite()
    demo_tdoa_bands()
    demo_mixed_no_sigma()
    demo_mixed_with_sigma()
    demo_terrain_intersection()
