# ruff: noqa: T201
"""Single-model IMM demos and IMM model identification.

Run with:
    python -m gri_sandbox.kalman_tracker_demo.demo_imm

Shows:
1-3. Single-model trackers (CV, CA, CT) with 60s prediction arrows
4.   Smart single-model tracker with outlier gating
5-7. IMM model identification on CV, CA, CT data with mode probabilities
"""

from __future__ import annotations

from dataclasses import dataclass, field

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

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

_SIGMA = SIGMA

# IMM transition matrix: 0.98 diagonal gives enough inertia for model
# probabilities to accumulate rather than chasing single-step noise
_IMM_PI = np.full((3, 3), 0.01)
np.fill_diagonal(_IMM_PI, 0.98)

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


# ---------------------------------------------------------------------------
# Dataclasses
# ---------------------------------------------------------------------------


@dataclass
class _ScenarioResult:
    """Collected data for plotting a single-model scenario."""

    title: str
    truth: np.ndarray
    obs: list[np.ndarray]
    filtered: list[np.ndarray]
    predicted: list[np.ndarray]
    outlier_indices: list[int] = field(default_factory=list)
    discarded: list[int] = field(default_factory=list)


@dataclass
class _IMMScenarioResult:
    """Collected data for plotting an IMM scenario."""

    title: str
    truth: np.ndarray
    obs: list[np.ndarray]
    filtered: list[np.ndarray]
    mode_probs: list[np.ndarray]


# ---------------------------------------------------------------------------
# Runners
# ---------------------------------------------------------------------------


def _run_tracker(
    tracker: IMM | SmartIMM,
    obs: list[np.ndarray],
    times: list[float],
    predict_dt: float = 60.0,
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Run tracker, collect filtered positions and predictions."""
    filtered = []
    predicted = []
    for xyz, t in zip(obs, times, strict=True):
        tracker.update(make_ell(xyz), timestamp_s=t)
        filtered.append(np.asarray(tracker.result.state.xyz))
        predicted.append(np.asarray(tracker.predict(t + predict_dt).xyz))
    return filtered, predicted


def _make_imm() -> IMM:
    """Standard 3-model IMM for demos.

    Tight CV process noise commits it to straight lines. Low CA jerk
    noise commits it to smooth acceleration. Low CT omega noise lets
    the turn rate accumulate from velocity trend.

    Uses mixing=False (parallel filter bank) so models maintain
    independent state, and accumulate_likelihood=True so mode
    probabilities reflect total fit over all observations rather
    than reacting to single-step noise.
    """
    return IMM(
        [
            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,
            ),
        ],
        transition_matrix=_IMM_PI,
        mixing=False,
        accumulate_likelihood=True,
    )


def _run_imm(
    obs: list[np.ndarray],
    times: list[float],
) -> tuple[list[np.ndarray], list[np.ndarray]]:
    """Run IMM tracker, collect filtered positions and mode probs."""
    imm = _make_imm()
    filtered = []
    mode_probs = []
    for xyz, t in zip(obs, times, strict=True):
        imm.update(make_ell(xyz, sigma_m=_SIGMA), timestamp_s=t)
        filtered.append(np.asarray(imm.result.state.xyz))
        mode_probs.append(imm.result.mode_probabilities.copy())
    return filtered, mode_probs


# ---------------------------------------------------------------------------
# Plot functions
# ---------------------------------------------------------------------------


def _plot(result: _ScenarioResult) -> None:
    """Show one plotly figure for a single-model scenario result."""
    origin = result.truth[0]
    truth = result.truth - origin
    obs = np.array(result.obs) - origin
    filt = np.array(result.filtered) - origin
    pred = np.array(result.predicted) - origin

    fig = go.Figure()

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

    # Observations
    if result.outlier_indices:
        clean = [i for i in range(len(result.obs)) if i not in result.outlier_indices]
        fig.add_trace(
            go.Scatter(
                x=obs[clean, 0],
                y=obs[clean, 1],
                mode="markers",
                name="Observations",
                marker={"color": "#636EFA", "symbol": "x", "size": 8},
            ),
        )
        fig.add_trace(
            go.Scatter(
                x=obs[result.outlier_indices, 0],
                y=obs[result.outlier_indices, 1],
                mode="markers",
                name="Outliers (injected)",
                marker={"color": "#EF553B", "symbol": "x", "size": 10},
            ),
        )
    else:
        fig.add_trace(
            go.Scatter(
                x=obs[:, 0],
                y=obs[:, 1],
                mode="markers",
                name="Observations",
                marker={"color": "#636EFA", "symbol": "x", "size": 8},
            ),
        )

    # Filtered
    fig.add_trace(
        go.Scatter(
            x=filt[:, 0],
            y=filt[:, 1],
            mode="lines+markers",
            name="Filtered",
            line={"color": "#00CC96", "width": 2},
            marker={"size": 6},
        ),
    )

    # Prediction arrows (lines from filtered to predicted)
    arrow_x: list[float | None] = []
    arrow_y: list[float | None] = []
    for i in range(len(filt)):
        arrow_x.extend([filt[i, 0], pred[i, 0], None])
        arrow_y.extend([filt[i, 1], pred[i, 1], None])

    fig.add_trace(
        go.Scatter(
            x=arrow_x,
            y=arrow_y,
            mode="lines",
            name="Predict (+60s)",
            line={"color": "#FFA15A", "width": 1, "dash": "dot"},
        ),
    )

    # Predicted endpoints
    fig.add_trace(
        go.Scatter(
            x=pred[:, 0].tolist(),
            y=pred[:, 1].tolist(),
            mode="markers",
            name="Predicted pos",
            marker={
                "color": "#FFA15A",
                "symbol": "diamond-open",
                "size": 6,
            },
        ),
    )

    # Mark discarded observations
    if result.discarded:
        fig.add_trace(
            go.Scatter(
                x=obs[result.discarded, 0],
                y=obs[result.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=result.title,
        xaxis_title="X offset (m)",
        yaxis_title="Y offset (m)",
        yaxis_scaleanchor="x",
    )
    fig.show()


def _plot_imm(result: _IMMScenarioResult) -> None:
    """Show one plotly figure for an IMM scenario with mode probabilities."""
    origin = result.truth[0]
    truth = result.truth - origin
    obs = np.array(result.obs) - origin
    filt = np.array(result.filtered) - origin
    probs = np.array(result.mode_probs)

    # Build hover text showing mode probabilities
    hover = []
    for i in range(len(result.filtered)):
        parts = [f"Step {i}"]
        for j, name in enumerate(_IMM_MODEL_NAMES):
            parts.append(f"{name}: {probs[i, j]:.1%}")
        best = _IMM_MODEL_NAMES[int(np.argmax(probs[i]))]
        parts.append(f"Best: {best}")
        hover.append("<br>".join(parts))

    fig = go.Figure()

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

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

    # Filtered track colored by dominant model
    # CV=#00CC96 (green), CA=#AB63FA (purple), CT=#FFA15A (orange)
    model_colors = ["#00CC96", "#AB63FA", "#FFA15A"]
    marker_colors = [
        model_colors[int(np.argmax(probs[i]))] for i in range(len(result.filtered))
    ]

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

    # Dummy traces for legend
    for j, name in enumerate(_IMM_MODEL_NAMES):
        fig.add_trace(
            go.Scatter(
                x=[None],
                y=[None],
                mode="markers",
                name=f"{name} dominant",
                marker={"color": model_colors[j], "size": 10},
            ),
        )

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

    # Console printout
    print(f"\n{result.title}")
    print(f"{'Step':>4}  {'CV':>6}  {'CA':>6}  {'CT':>6}  {'Best':>4}")
    for i in range(len(result.filtered)):
        best = _IMM_MODEL_NAMES[int(np.argmax(probs[i]))]
        print(
            f"{i:4d}  {probs[i, 0]:6.1%}  {probs[i, 1]:6.1%}  "
            f"{probs[i, 2]:6.1%}  {best:>4}",
        )


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


def _generate_cv(rng: np.random.Generator) -> _ScenarioResult:
    """Constant velocity: ~100 m/s, ~60km over 10 minutes."""
    vel = np.array([90.0, 45.0, 0.0])  # ~100 m/s total
    truth = np.array([X0 + vel * i * DT for i in range(N)])
    obs = [truth[i] + rng.normal(0, _SIGMA, 3) for i in range(N)]
    times = [i * DT for i in range(N)]

    tracker = IMM([ConstantVelocity(process_noise_m_s2=0.1)])
    filtered, predicted = _run_tracker(tracker, obs, times)
    return _ScenarioResult(
        title="Constant Velocity (~100 m/s, 10 min)",
        truth=truth,
        obs=obs,
        filtered=filtered,
        predicted=predicted,
    )


def _generate_ca(rng: np.random.Generator) -> _ScenarioResult:
    """Constant acceleration: 100 m/s +X with 1 m/s^2 lateral +Y."""
    v0 = np.array([100.0, 0.0, 0.0])
    acc = np.array([0.0, 1.0, 0.0])  # lateral acceleration
    truth = np.array(
        [X0 + v0 * i * DT + 0.5 * acc * (i * DT) ** 2 for i in range(N)],
    )
    obs = [truth[i] + rng.normal(0, _SIGMA, 3) for i in range(N)]
    times = [i * DT for i in range(N)]

    tracker = IMM([ConstantAcceleration(process_noise_m_s3=0.001)])
    filtered, predicted = _run_tracker(tracker, obs, times)
    return _ScenarioResult(
        title="Constant Acceleration (100 m/s +X, 1 m/s^2 +Y)",
        truth=truth,
        obs=obs,
        filtered=filtered,
        predicted=predicted,
    )


def _generate_ct(rng: np.random.Generator) -> _ScenarioResult:
    """Coordinated turn: CT model learning turn rate over time.

    Very low omega process noise tells the filter that the turn rate is
    nearly constant, letting it accumulate evidence gradually rather than
    chasing noise. The filter converges on omega after ~10 observations.
    """
    speed = 100.0
    omega = 0.003  # rad/s, gentle turn
    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]))
    truth_arr = np.array(truth)
    obs = [truth_arr[i] + rng.normal(0, _SIGMA, 3) for i in range(N)]
    times = [i * DT for i in range(N)]

    tracker = IMM(
        [
            CoordinatedTurn(
                process_noise_m_s2=0.5,
                process_noise_omega_rad_s2=1e-5,
            ),
        ],
    )
    filtered, predicted = _run_tracker(tracker, obs, times)
    return _ScenarioResult(
        title="Coordinated Turn (CT, omega learned from trend)",
        truth=truth_arr,
        obs=obs,
        filtered=filtered,
        predicted=predicted,
    )


def _generate_smart(rng: np.random.Generator) -> _ScenarioResult:
    """CV with injected outliers and smart gating."""
    vel = np.array([90.0, 45.0, 0.0])
    outlier_indices = [5, 12, 16]
    truth = np.array([X0 + vel * i * DT for i in range(N)])
    obs = []
    for i in range(N):
        noise = rng.normal(0, _SIGMA, 3)
        if i in outlier_indices:
            noise += rng.normal(0, 5000, 3)  # ~5km offset
        obs.append(truth[i] + noise)
    times = [i * DT for i in range(N)]

    tracker = SmartIMM(
        [ConstantVelocity(process_noise_m_s2=0.1)],
        max_norm=3.0,
    )
    filtered, predicted = _run_tracker(tracker, obs, times)
    return _ScenarioResult(
        title="Smart Tracker (CV + 3 outliers, gated)",
        truth=truth,
        obs=obs,
        filtered=filtered,
        predicted=predicted,
        outlier_indices=outlier_indices,
        discarded=tracker.result.discarded,
    )


def _generate_imm_cv(rng: np.random.Generator) -> _IMMScenarioResult:
    """IMM on constant velocity: straight line at ~100 m/s."""
    vel = np.array([90.0, 45.0, 0.0])
    truth = np.array([X0 + vel * i * DT for i in range(N)])
    obs = [truth[i] + rng.normal(0, _SIGMA, 3) for i in range(N)]
    times = [i * DT for i in range(N)]

    filtered, mode_probs = _run_imm(obs, times)
    return _IMMScenarioResult(
        title="IMM: Constant Velocity (straight line)",
        truth=truth,
        obs=obs,
        filtered=filtered,
        mode_probs=mode_probs,
    )


def _generate_imm_ca(rng: np.random.Generator) -> _IMMScenarioResult:
    """IMM on lateral acceleration: 100 m/s +X with 1 m/s^2 +Y.

    The lateral acceleration produces a visible curve that CV cannot
    absorb since it creates systematic cross-track residuals.
    """
    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)],
    )
    obs = [truth[i] + rng.normal(0, _SIGMA, 3) for i in range(N)]
    times = [i * DT for i in range(N)]

    filtered, mode_probs = _run_imm(obs, times)
    return _IMMScenarioResult(
        title="IMM: Lateral Acceleration (100 m/s +X, 1 m/s^2 +Y)",
        truth=truth,
        obs=obs,
        filtered=filtered,
        mode_probs=mode_probs,
    )


def _generate_imm_ct(rng: np.random.Generator) -> _IMMScenarioResult:
    """IMM on coordinated turn: 100 m/s with omega=0.005 rad/s."""
    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]))
    truth_arr = np.array(truth)
    obs = [truth_arr[i] + rng.normal(0, _SIGMA, 3) for i in range(N)]
    times = [i * DT for i in range(N)]

    filtered, mode_probs = _run_imm(obs, times)
    return _IMMScenarioResult(
        title="IMM: Coordinated Turn (omega=0.005 rad/s)",
        truth=truth_arr,
        obs=obs,
        filtered=filtered,
        mode_probs=mode_probs,
    )


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


def main() -> None:
    """Run single-model and IMM model identification demos."""
    rng = np.random.default_rng(42)

    # Single-model tracker demos (CV, CA, CT, Smart)
    for gen in [_generate_cv, _generate_ca, _generate_ct, _generate_smart]:
        _plot(gen(rng))

    # IMM model identification demos
    imm_rng = np.random.default_rng(58)
    for gen_imm in [_generate_imm_cv, _generate_imm_ca, _generate_imm_ct]:
        _plot_imm(gen_imm(imm_rng))

    print("\nIMM demos done.")


if __name__ == "__main__":
    main()
