"""
ekf_math.py — Extended Kalman Filter for UWB + IMU fusion (Issue #573)

No ROS2 dependencies — pure Python + numpy, fully unit-testable.

State vector (6-DOF ground robot):
───────────────────────────────────
    x  = [x, y, θ, vx, vy, ω]ᵀ    (world frame)
    x  — forward position (m)
    y  — lateral position (m)
    θ  — heading (rad, CCW positive)
    vx — longitudinal velocity (m/s, world frame)
    vy — lateral velocity (m/s, world frame)
    ω  — yaw rate (rad/s)

Process model — IMU as control input u = [ax_body, ay_body, ω_imu]
───────────────────────────────────────────────────────────────────
    θ_k+1  = θ_k  + ω_imu * dt
    vx_k+1 = vx_k + (ax_body * cos(θ) - ay_body * sin(θ)) * dt
    vy_k+1 = vy_k + (ax_body * sin(θ) + ay_body * cos(θ)) * dt
    x_k+1  = x_k  + vx_k * dt
    y_k+1  = y_k  + vy_k * dt
    ω_k+1  = ω_imu  (direct observation, no integration)

UWB measurement model (position observation):
──────────────────────────────────────────────
    z = [x, y]ᵀ    H = [[1,0,0,0,0,0],
                         [0,1,0,0,0,0]]

IMU bias handling:
──────────────────
    Simple first-order high-pass on accel to reduce drift.
    For longer deployments extend state to include accel bias.
"""

from __future__ import annotations

import math
import numpy as np
from typing import Tuple

# State dimension
N = 6   # [x, y, θ, vx, vy, ω]

# Indices
IX, IY, IT, IVX, IVY, IW = range(N)


class UwbImuEKF:
    """
    EKF fusing UWB position (10 Hz) with IMU accel+gyro (200 Hz).

    Parameters
    ----------
    sigma_uwb_pos_m : UWB position std-dev (m)           default 0.12
    sigma_imu_accel : IMU accelerometer noise (m/s²)     default 0.05
    sigma_imu_gyro  : IMU gyroscope noise (rad/s)        default 0.003
    sigma_vel_drift : velocity drift process noise (m/s) default 0.1
    dropout_vel_damping : velocity damping per second during UWB dropout
    """

    def __init__(
        self,
        sigma_uwb_pos_m: float    = 0.12,
        sigma_imu_accel: float    = 0.05,
        sigma_imu_gyro: float     = 0.003,
        sigma_vel_drift: float    = 0.10,
        dropout_vel_damping: float = 0.95,
    ) -> None:
        self._R_uwb  = sigma_uwb_pos_m ** 2   # UWB position variance
        self._q_a    = sigma_imu_accel ** 2    # accel noise variance
        self._q_g    = sigma_imu_gyro  ** 2    # gyro noise variance
        self._q_v    = sigma_vel_drift ** 2    # velocity drift variance
        self._damp   = dropout_vel_damping

        # State and covariance
        self._x = np.zeros(N, dtype=float)
        self._P = np.eye(N, dtype=float) * 9.0   # large initial uncertainty

        # Accel bias (simple running mean for high-pass)
        self._accel_bias = np.zeros(2, dtype=float)
        self._bias_alpha = 0.005   # low-pass coefficient for bias estimation

        self._initialised = False
        self._uwb_dropout_s = 0.0   # time since last UWB update

    # ── Initialise ───────────────────────────────────────────────────────────

    def initialise(self, x: float, y: float, heading_rad: float = 0.0) -> None:
        """Seed the filter with a known position."""
        self._x[:] = 0.0
        self._x[IX] = x
        self._x[IY] = y
        self._x[IT] = heading_rad
        self._P = np.diag([0.25, 0.25, 0.1, 1.0, 1.0, 0.1])
        self._initialised = True

    # ── IMU predict ──────────────────────────────────────────────────────────

    def predict(self, ax_body: float, ay_body: float, omega: float, dt: float) -> None:
        """
        EKF predict step driven by IMU measurement.

        Parameters
        ----------
        ax_body : forward acceleration in body frame (m/s²)
        ay_body : lateral acceleration in body frame (m/s², left positive)
        omega   : yaw rate (rad/s, CCW positive)
        dt      : time step (s)
        """
        if not self._initialised:
            return

        # Subtract estimated accel bias
        ax_c = ax_body - self._accel_bias[0]
        ay_c = ay_body - self._accel_bias[1]

        x  = self._x
        th = x[IT]
        ct = math.cos(th)
        st = math.sin(th)

        # World-frame acceleration
        ax_w = ax_c * ct - ay_c * st
        ay_w = ax_c * st + ay_c * ct

        # State prediction
        xp = x.copy()
        xp[IX]  = x[IX]  + x[IVX] * dt
        xp[IY]  = x[IY]  + x[IVY] * dt
        xp[IT]  = _wrap_angle(x[IT] + omega * dt)
        xp[IVX] = x[IVX] + ax_w * dt
        xp[IVY] = x[IVY] + ay_w * dt
        xp[IW]  = omega   # direct observation from gyro

        # Jacobian F = ∂f/∂x
        F = np.eye(N, dtype=float)
        # ∂x / ∂vx, ∂x / ∂vy
        F[IX, IVX] = dt
        F[IY, IVY] = dt
        # ∂vx / ∂θ  = (-ax_c·sin(θ) - ay_c·cos(θ)) * dt
        F[IVX, IT] = (-ax_c * st - ay_c * ct) * dt
        # ∂vy / ∂θ  = ( ax_c·cos(θ) - ay_c·sin(θ)) * dt
        F[IVY, IT] = ( ax_c * ct - ay_c * st) * dt

        # Process noise Q
        dt2 = dt * dt
        Q = np.diag([
            0.5 * self._q_a * dt2 * dt2,   # x: from accel double-integrated
            0.5 * self._q_a * dt2 * dt2,   # y
            self._q_g * dt2,               # θ: from gyro integrated
            self._q_v * dt + self._q_a * dt2,  # vx
            self._q_v * dt + self._q_a * dt2,  # vy
            self._q_g,                     # ω: direct gyro noise
        ])

        self._x = xp
        self._P = F @ self._P @ F.T + Q
        self._uwb_dropout_s += dt

        # Velocity damping during extended UWB dropout (dead-reckoning coasting)
        if self._uwb_dropout_s > 2.0:
            damping = self._damp ** dt
            self._x[IVX] *= damping
            self._x[IVY] *= damping

        # Update accel bias estimate
        self._accel_bias[0] += self._bias_alpha * (ax_body - self._accel_bias[0])
        self._accel_bias[1] += self._bias_alpha * (ay_body - self._accel_bias[1])

    # ── UWB update ───────────────────────────────────────────────────────────

    def update_uwb(self, x_meas: float, y_meas: float,
                   sigma_m: float | None = None) -> None:
        """
        EKF update step with a UWB position measurement.

        Parameters
        ----------
        x_meas, y_meas : measured position (m, world frame)
        sigma_m        : override measurement std-dev (m); uses node default if None
        """
        if not self._initialised:
            self.initialise(x_meas, y_meas)
            return

        R_val = (sigma_m ** 2) if sigma_m is not None else self._R_uwb
        R = np.eye(2) * R_val

        # H: position rows only
        H = np.zeros((2, N))
        H[0, IX] = 1.0
        H[1, IY] = 1.0

        innov = np.array([x_meas - self._x[IX], y_meas - self._x[IY]])
        S     = H @ self._P @ H.T + R
        K     = self._P @ H.T @ np.linalg.inv(S)

        self._x = self._x + K @ innov
        self._x[IT] = _wrap_angle(self._x[IT])
        self._P = (np.eye(N) - K @ H) @ self._P

        self._uwb_dropout_s = 0.0

    # ── Update heading from magnetometer / other source ───────────────────────

    def update_heading(self, heading_rad: float, sigma_rad: float = 0.1) -> None:
        """Optional: update θ directly from compass or visual odometry."""
        if not self._initialised:
            return
        H = np.zeros((1, N))
        H[0, IT] = 1.0
        innov = np.array([_wrap_angle(heading_rad - self._x[IT])])
        S     = H @ self._P @ H.T + np.array([[sigma_rad ** 2]])
        K     = self._P @ H.T @ np.linalg.inv(S)
        self._x = self._x + K.flatten() * innov[0]
        self._x[IT] = _wrap_angle(self._x[IT])
        self._P = (np.eye(N) - K @ H) @ self._P

    # ── Accessors ────────────────────────────────────────────────────────────

    @property
    def position(self) -> Tuple[float, float]:
        return float(self._x[IX]), float(self._x[IY])

    @property
    def heading(self) -> float:
        return float(self._x[IT])

    @property
    def velocity(self) -> Tuple[float, float]:
        return float(self._x[IVX]), float(self._x[IVY])

    @property
    def yaw_rate(self) -> float:
        return float(self._x[IW])

    @property
    def covariance_position(self) -> np.ndarray:
        """2×2 position covariance sub-matrix."""
        return self._P[:2, :2].copy()

    @property
    def covariance_full(self) -> np.ndarray:
        """Full 6×6 state covariance."""
        return self._P.copy()

    @property
    def is_initialised(self) -> bool:
        return self._initialised

    @property
    def uwb_dropout_s(self) -> float:
        """Seconds since last UWB update."""
        return self._uwb_dropout_s

    def position_uncertainty_m(self) -> float:
        """Approximate 1-σ position uncertainty (m)."""
        return float(math.sqrt((self._P[IX, IX] + self._P[IY, IY]) / 2.0))


# ── Helpers ───────────────────────────────────────────────────────────────────

def _wrap_angle(a: float) -> float:
    """Wrap angle to [-π, π]."""
    return float((a + math.pi) % (2 * math.pi) - math.pi)
