# ruff: noqa: T201
"""SmartIMM outlier detection on single segments.

Run with:
    python -m gri_sandbox.kalman_tracker_demo.demo_smart_imm

Shows SmartIMM with CV+CA+CT models rejecting outliers on:
1. Constant velocity data + outliers
2. Constant acceleration data + outliers
3. Coordinated turn data + outliers
"""

from __future__ import annotations

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

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

_SIGMA = SIGMA

# Transition matrix matching the IMM demos
_IMM_PI = np.full((3, 3), 0.01)
np.fill_diagonal(_IMM_PI, 0.98)

_MODEL_NAMES = ["CV", "CA", "CT"]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_smart_imm() -> SmartIMM:
    """SmartIMM with CV+CA+CT for single-segment demos.

    Uses max_norm=5.0 (tighter than default 10.0) since these are
    single-segment scenarios with no changepoint concern.
    """
    return SmartIMM(
        [
            ConstantVelocity(process_noise_m_s2=0.01),
            ConstantAcceleration(process_noise_m_s3=0.001),
            CoordinatedTurn(
                process_noise_m_s2=0.1,
                process_noise_omega_rad_s2=1e-5,
            ),
        ],
        max_norm=5.0,
        transition_matrix=_IMM_PI,
        mixing=False,
        accumulate_likelihood=True,
    )


def _plot_smart_imm(  # noqa: PLR0913
    title: str,
    truth: np.ndarray,
    obs: list[np.ndarray],
    filtered: list[np.ndarray],
    outlier_indices: list[int],
    discarded: list[int],
    mode_probs: np.ndarray,
    model_names: list[str],
) -> None:
    """Plot SmartIMM result showing outliers and mode probabilities."""
    origin = truth[0]
    truth_rel = truth - origin
    obs_rel = np.array(obs) - origin
    filt_rel = np.array(filtered) - origin

    # Separate clean obs from injected outliers
    clean = [i for i in range(len(obs)) if i not in outlier_indices]

    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[clean, 0],
            y=obs_rel[clean, 1],
            mode="markers",
            name="Observations",
            marker={"color": "#636EFA", "symbol": "x", "size": 8},
        ),
    )

    if outlier_indices:
        fig.add_trace(
            go.Scatter(
                x=obs_rel[outlier_indices, 0],
                y=obs_rel[outlier_indices, 1],
                mode="markers",
                name="Outliers (injected)",
                marker={"color": "#EF553B", "symbol": "x", "size": 10},
            ),
        )

    # Filtered track colored by dominant model
    model_colors = ["#00CC96", "#AB63FA", "#FFA15A", "#636EFA"]
    hover = []
    for i in range(len(filtered)):
        parts = [f"Step {i}"]
        for j, name in enumerate(model_names):
            parts.append(f"{name}: {mode_probs[i, j]:.1%}")
        hover.append("<br>".join(parts))

    marker_colors = [
        model_colors[int(np.argmax(mode_probs[i])) % len(model_colors)]
        for i in range(len(filtered))
    ]

    fig.add_trace(
        go.Scatter(
            x=filt_rel[:, 0],
            y=filt_rel[:, 1],
            mode="lines+markers",
            name="Filtered",
            line={"color": "#888888", "width": 1},
            marker={"size": 8, "color": marker_colors},
            text=hover,
            hoverinfo="text",
        ),
    )

    # Mark discarded observations with circles
    if discarded:
        fig.add_trace(
            go.Scatter(
                x=obs_rel[discarded, 0],
                y=obs_rel[discarded, 1],
                mode="markers",
                name="Discarded (gated)",
                marker={
                    "color": "rgba(0,0,0,0)",
                    "line": {"color": "#EF553B", "width": 2},
                    "symbol": "circle",
                    "size": 14,
                },
            ),
        )

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

    # Console
    print(f"\n{title}")
    print(f"Injected outliers: {outlier_indices}")
    print(f"Discarded: {discarded}")
    detected = [i for i in outlier_indices if i in discarded]
    missed = [i for i in outlier_indices if i not in discarded]
    false_pos = [i for i in discarded if i not in outlier_indices]
    print(f"Detected: {detected}, Missed: {missed}, False+: {false_pos}")


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


def _run_smart_imm_scenario(
    rng: np.random.Generator,
    title: str,
    truth: np.ndarray,
    outlier_indices: list[int],
) -> None:
    """Run SmartIMM on a single-segment trajectory with outliers."""
    tracker = _make_smart_imm()

    obs_list: list[np.ndarray] = []
    times: list[float] = []
    for i in range(len(truth)):
        noise = rng.normal(0, _SIGMA, 3)
        if i in outlier_indices:
            noise += rng.normal(0, 15000, 3)
        obs_list.append(truth[i] + noise)
        times.append(i * DT)

    filtered = []
    probs = []
    for xyz, t in zip(obs_list, times, strict=True):
        tracker.update(make_ell(xyz, sigma_m=_SIGMA), timestamp_s=t)
        filtered.append(np.asarray(tracker.result.state.xyz))
        probs.append(tracker.result.mode_probabilities.copy())

    _plot_smart_imm(
        title,
        truth,
        obs_list,
        filtered,
        outlier_indices,
        tracker.result.discarded,
        np.array(probs),
        _MODEL_NAMES,
    )


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


def _generate_smart_imm_cv(rng: np.random.Generator) -> None:
    """SmartIMM on CV data with outliers."""
    vel = np.array([90.0, 45.0, 0.0])
    truth = np.array([X0 + vel * i * DT for i in range(N)])
    _run_smart_imm_scenario(
        rng,
        "SmartIMM: CV + Outliers",
        truth,
        outlier_indices=[8, 16],
    )


def _generate_smart_imm_ca(rng: np.random.Generator) -> None:
    """SmartIMM on lateral acceleration data with outliers."""
    v0 = np.array([100.0, 0.0, 0.0])
    acc = np.array([0.0, 1.0, 0.0])
    truth = np.array(
        [X0 + v0 * i * DT + 0.5 * acc * (i * DT) ** 2 for i in range(N)],
    )
    _run_smart_imm_scenario(
        rng,
        "SmartIMM: CA (lateral) + Outliers",
        truth,
        outlier_indices=[8, 16],
    )


def _generate_smart_imm_ct(rng: np.random.Generator) -> None:
    """SmartIMM on coordinated turn data with outliers."""
    speed = 100.0
    omega = 0.005
    truth = []
    for i in range(N):
        angle = omega * i * DT
        dx = speed / omega * np.sin(angle)
        dy = speed / omega * (1 - np.cos(angle))
        truth.append(X0 + np.array([dx, dy, 0.0]))
    _run_smart_imm_scenario(
        rng,
        "SmartIMM: CT + Outliers",
        np.array(truth),
        outlier_indices=[8, 16],
    )


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


def main() -> None:
    """Run SmartIMM outlier detection demos."""
    smart_imm_rng = np.random.default_rng(42)
    _generate_smart_imm_cv(smart_imm_rng)
    _generate_smart_imm_ca(smart_imm_rng)
    _generate_smart_imm_ct(smart_imm_rng)

    print("\nSmartIMM demos done.")


if __name__ == "__main__":
    main()
