# ruff: noqa: T201
"""SegmentedIMM demos: sailboat tacking and racetrack pattern.

Run with:
    python -m gri_sandbox.kalman_tracker_demo.demo_segmented_imm

Shows SegmentedIMM (no outlier rejection) on:
1. Sailboat tacking: CV-only, alternating headings
2. Racetrack: CV+CT, straight-turn-straight-turn
"""

from __future__ import annotations

import numpy as np
import plotly.graph_objects as go
from gri_convolve.kalman import (
    ConstantVelocity,
    CoordinatedTurn,
    SegmentedIMM,
)

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

# ---------------------------------------------------------------------------
# Plot function
# ---------------------------------------------------------------------------


def _plot_segmented(  # noqa: PLR0913
    title: str,
    truth: np.ndarray,
    obs: list[np.ndarray],
    segments: list,
    model_names: list[str],
    true_boundaries: list[int],
) -> None:
    """Plot a segmented track with color-coded legs."""
    origin = truth[0]
    truth_rel = truth - origin
    obs_rel = np.array(obs) - origin

    seg_colors = [
        "#00CC96",
        "#AB63FA",
        "#FFA15A",
        "#EF553B",
        "#636EFA",
        "#FF6692",
        "#B6E880",
    ]

    fig = go.Figure()

    fig.add_trace(
        go.Scatter(
            x=truth_rel[:, 0],
            y=truth_rel[:, 1],
            mode="lines",
            name="Truth",
            line={"color": "white", "width": 2},
        ),
    )

    fig.add_trace(
        go.Scatter(
            x=obs_rel[:, 0],
            y=obs_rel[:, 1],
            mode="markers",
            name="Observations",
            marker={"color": "#636EFA", "symbol": "x", "size": 6, "opacity": 0.4},
        ),
    )

    # Plot each segment as a colored line
    for i, seg in enumerate(segments):
        color = seg_colors[i % len(seg_colors)]
        seg_obs = obs_rel[seg.start : seg.end]
        model_name = model_names[seg.model_index]
        fig.add_trace(
            go.Scatter(
                x=seg_obs[:, 0],
                y=seg_obs[:, 1],
                mode="lines+markers",
                name=f"Seg {i}: {model_name} [{seg.start}-{seg.end})",
                line={"color": color, "width": 3},
                marker={"size": 5},
            ),
        )

    # Mark true boundaries with diamond markers
    for b in true_boundaries:
        if b < len(truth_rel):
            fig.add_trace(
                go.Scatter(
                    x=[truth_rel[b, 0]],
                    y=[truth_rel[b, 1]],
                    mode="markers",
                    name=f"True boundary @ {b}",
                    marker={
                        "color": "white",
                        "symbol": "diamond",
                        "size": 12,
                        "line": {"width": 2, "color": "white"},
                    },
                ),
            )

    fig.update_layout(
        template="plotly_dark",
        title=title,
        xaxis_title="X offset (m)",
        yaxis_title="Y offset (m)",
        yaxis_scaleanchor="x",
    )
    fig.show()

    # Console printout
    print(f"\n{title}")
    print(f"True boundaries: {true_boundaries}")
    print(f"Found {len(segments)} segments:")
    for i, seg in enumerate(segments):
        print(
            f"  Seg {i}: [{seg.start:3d}-{seg.end:3d}) "
            f"model={model_names[seg.model_index]}  len={seg.length}",
        )


# ---------------------------------------------------------------------------
# Scenario generators
# ---------------------------------------------------------------------------


def _generate_sailboat(rng: np.random.Generator) -> None:
    """Sailboat tacking: alternating CV headings."""
    speed = 100.0
    headings_deg = [45.0, -45.0, 45.0, -45.0]
    leg_lengths = [30, 20, 30, 20]
    true_boundaries = [30, 50, 80]

    tracker = SegmentedIMM(
        [ConstantVelocity(process_noise_m_s2=0.01)],
        window=10,
    )

    truth_list: list[np.ndarray] = []
    obs_list: list[np.ndarray] = []
    pos = X0.copy()
    obs_idx = 0

    for heading_deg, n_obs in zip(headings_deg, leg_lengths, strict=True):
        heading = np.radians(heading_deg)
        vel = speed * np.array([np.cos(heading), np.sin(heading), 0.0])
        for i in range(n_obs):
            t = obs_idx * DT
            pos_true = pos + vel * i * DT
            truth_list.append(pos_true)
            noisy = pos_true + rng.normal(0, SIGMA, 3)
            obs_list.append(noisy)
            tracker.update(
                make_ell(noisy, sigma_m=SIGMA),
                timestamp_s=t,
            )
            obs_idx += 1
        pos = pos + vel * n_obs * DT

    _plot_segmented(
        "Segmented: Sailboat Tacking (CV-only, 4 legs)",
        np.array(truth_list),
        obs_list,
        tracker.segments,
        ["CV"],
        true_boundaries,
    )


def _generate_racetrack(rng: np.random.Generator) -> None:
    """Racetrack: straight-turn-straight-turn."""
    speed = 100.0
    omega = 0.01  # rad/s for turns

    tracker = SegmentedIMM(
        [
            ConstantVelocity(process_noise_m_s2=0.01),
            CoordinatedTurn(
                process_noise_m_s2=0.1,
                process_noise_omega_rad_s2=1e-5,
            ),
        ],
        window=10,
    )

    legs = [
        ("straight", 30),
        ("turn", 15),
        ("straight", 30),
        ("turn", 15),
    ]
    true_boundaries = [30, 45, 75]

    truth_list: list[np.ndarray] = []
    obs_list: list[np.ndarray] = []
    pos = X0.copy()
    heading = 0.0
    obs_idx = 0

    for leg_type, n_obs in legs:
        vel = speed * np.array([np.cos(heading), np.sin(heading), 0.0])
        for i in range(n_obs):
            t = obs_idx * DT
            if leg_type == "straight":
                pos_true = pos + vel * i * DT
            else:
                angle = omega * i * DT
                dx = speed / omega * np.sin(angle)
                dy = speed / omega * (1 - np.cos(angle))
                c, s = np.cos(heading), np.sin(heading)
                pos_true = pos + np.array(
                    [c * dx - s * dy, s * dx + c * dy, 0.0],
                )
            truth_list.append(pos_true)
            noisy = pos_true + rng.normal(0, SIGMA, 3)
            obs_list.append(noisy)
            tracker.update(
                make_ell(noisy, sigma_m=SIGMA),
                timestamp_s=t,
            )
            obs_idx += 1

        if leg_type == "straight":
            pos = pos + vel * n_obs * DT
        else:
            total_turn = omega * n_obs * DT
            dx = speed / omega * np.sin(total_turn)
            dy = speed / omega * (1 - np.cos(total_turn))
            c, s = np.cos(heading), np.sin(heading)
            pos = pos + np.array(
                [c * dx - s * dy, s * dx + c * dy, 0.0],
            )
            heading += total_turn

    _plot_segmented(
        "Segmented: Racetrack (CV+CT, 4 legs)",
        np.array(truth_list),
        obs_list,
        tracker.segments,
        ["CV", "CT"],
        true_boundaries,
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    """Run SegmentedIMM demos."""
    _generate_sailboat(np.random.default_rng(42))
    _generate_racetrack(np.random.default_rng(42))

    print("\nSegmentedIMM demos done.")


if __name__ == "__main__":
    main()
