"""
pose_fusion_ekf.py — Extended Kalman Filter for multi-sensor pose fusion (Issue #595).

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

Fuses three sensor streams
──────────────────────────
  1. IMU (200 Hz) — predict step (body-frame accel + gyro)
  2. UWB+IMU fused pose (10 Hz) — absolute position and heading update
  3. Visual odometry (30 Hz) — velocity-domain update (drift-free short-term motion)

State vector [x, y, θ, vx, vy, ω]ᵀ
────────────────────────────────────
  x   — position, forward    (m, map frame)
  y   — position, lateral    (m, map frame)
  θ   — heading, yaw         (rad, CCW positive)
  vx  — velocity x           (m/s, map frame)
  vy  — velocity y           (m/s, map frame)
  ω   — yaw rate             (rad/s)

Process model (IMU as control input):
  θ_{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

Measurement models
──────────────────
  UWB position:  z = [x, y]             H[2×6] = [[1,0,0,0,0,0],[0,1,0,0,0,0]]
  UWB heading:   z = [θ]                H[1×6] = [[0,0,1,0,0,0]]
  VO velocity:   z = [vx_w, vy_w, ω]   H[3×6] = diag of [vx,vy,ω] rows in state

  VO velocity measurement: VO gives vx_body (forward speed) and ω.
  We rotate to world frame: vx_w = vx_body*cos(θ), vy_w = vx_body*sin(θ)
  then update [vx, vy, ω] state directly (linear measurement step).

Sensor dropout handling
───────────────────────
  _uwb_dropout_s: seconds since last UWB update
  _vo_dropout_s:  seconds since last VO update
  Both increment during predict(), reset to 0 on update.
  When uwb_dropout_s > uwb_max_s, velocity damping activates and output
  is suppressed by the node.

Sensor weights (configurable via constructor)
─────────────────────────────────────────────
  sigma_uwb_pos_m    — UWB+IMU fused position σ (m)        default 0.10
  sigma_uwb_head_rad — UWB+IMU fused heading σ (rad)       default 0.15
  sigma_vo_vel_m_s   — VO linear velocity σ (m/s)          default 0.05
  sigma_vo_omega_r_s — VO yaw-rate σ (rad/s)               default 0.03
  sigma_imu_accel    — IMU accel noise σ (m/s²)            default 0.05
  sigma_imu_gyro     — IMU gyro noise σ (rad/s)            default 0.003
  sigma_vel_drift    — velocity process noise σ (m/s)       default 0.10
  dropout_vel_damp   — velocity damping/s during UWB drop   default 0.95
"""

from __future__ import annotations

import math
from typing import Tuple

import numpy as np


# ── State indices ──────────────────────────────────────────────────────────────
N = 6
IX, IY, IT, IVX, IVY, IW = range(N)


def _wrap(a: float) -> float:
    """Wrap angle to (−π, π]."""
    return float((a + math.pi) % (2.0 * math.pi) - math.pi)


class PoseFusionEKF:
    """
    EKF fusing UWB+IMU absolute pose (10 Hz), visual-odometry velocity (30 Hz),
    and raw IMU (200 Hz) into a single authoritative map-frame pose.

    Thread model: single-threaded — caller must serialise access.
    """

    def __init__(
        self,
        sigma_uwb_pos_m:    float = 0.10,
        sigma_uwb_head_rad: float = 0.15,
        sigma_vo_vel_m_s:   float = 0.05,
        sigma_vo_omega_r_s: float = 0.03,
        sigma_imu_accel:    float = 0.05,
        sigma_imu_gyro:     float = 0.003,
        sigma_vel_drift:    float = 0.10,
        dropout_vel_damp:   float = 0.95,
    ) -> None:
        # Measurement noise variances
        self._R_uwb_pos  = sigma_uwb_pos_m    ** 2
        self._R_uwb_head = sigma_uwb_head_rad ** 2
        self._R_vo_vel   = sigma_vo_vel_m_s   ** 2
        self._R_vo_omega = sigma_vo_omega_r_s ** 2

        # Process noise parameters
        self._q_a    = sigma_imu_accel  ** 2
        self._q_g    = sigma_imu_gyro   ** 2
        self._q_v    = sigma_vel_drift  ** 2
        self._damp   = dropout_vel_damp

        # State & covariance
        self._x    = np.zeros(N, dtype=np.float64)
        self._P    = np.eye(N, dtype=np.float64) * 9.0

        # Simple accel bias estimator
        self._accel_bias  = np.zeros(2, dtype=np.float64)
        self._bias_alpha  = 0.005

        self._init           = False
        self._uwb_dropout_s  = 0.0
        self._vo_dropout_s   = 0.0

    # ── Initialisation ─────────────────────────────────────────────────────────

    def initialise(self, x: float, y: float, heading: float = 0.0) -> None:
        """Seed filter with known position (called on first UWB measurement)."""
        self._x[:] = 0.0
        self._x[IX] = x
        self._x[IY] = y
        self._x[IT] = heading
        self._P     = np.diag([0.25, 0.25, 0.1, 1.0, 1.0, 0.1])
        self._init  = True

    # ── IMU predict step (200 Hz) ──────────────────────────────────────────────

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

        Parameters
        ----------
        ax_body : forward accel (m/s², body frame)
        ay_body : lateral accel (m/s², body frame, left positive)
        omega   : yaw rate (rad/s, CCW positive)
        dt      : timestep (s)
        """
        if not self._init:
            return

        # Bias-compensated acceleration
        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

        # Propagate state
        xp       = x.copy()
        xp[IX]   = x[IX]  + x[IVX] * dt
        xp[IY]   = x[IY]  + x[IVY] * dt
        xp[IT]   = _wrap(x[IT] + omega * dt)
        xp[IVX]  = x[IVX] + ax_w * dt
        xp[IVY]  = x[IVY] + ay_w * dt
        xp[IW]   = omega

        # Jacobian F = ∂f/∂x
        F        = np.eye(N, dtype=np.float64)
        F[IX, IVX] = dt
        F[IY, IVY] = dt
        F[IVX, IT] = (-ax_c * st - ay_c * ct) * 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,
            0.5 * self._q_a * dt2 * dt2,
            self._q_g * dt2,
            self._q_v * dt + self._q_a * dt2,
            self._q_v * dt + self._q_a * dt2,
            self._q_g,
        ])

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

        # Advance dropout clocks
        self._uwb_dropout_s += dt
        self._vo_dropout_s  += dt

        # Velocity damping during extended UWB dropout
        if self._uwb_dropout_s > 2.0:
            d = self._damp ** dt
            self._x[IVX] *= d
            self._x[IVY] *= d

        # Running 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+IMU position update (10 Hz) ───────────────────────────────────────

    def update_uwb_position(
        self,
        x_m: float,
        y_m: float,
        sigma_m: float | None = None,
    ) -> None:
        """
        Update from UWB+IMU fused position measurement [x, y].

        Parameters
        ----------
        x_m, y_m : measured position (m, map frame)
        sigma_m  : std-dev override; uses constructor default if None
        """
        if not self._init:
            self.initialise(x_m, y_m)
            return

        R_val = (sigma_m ** 2) if sigma_m is not None else self._R_uwb_pos
        R     = np.eye(2, dtype=np.float64) * R_val

        H     = np.zeros((2, N), dtype=np.float64)
        H[0, IX] = 1.0
        H[1, IY] = 1.0

        innov = np.array([x_m - self._x[IX], y_m - self._x[IY]])
        _ekf_update(self._x, self._P, H, R, innov)
        self._x[IT] = _wrap(self._x[IT])

        self._uwb_dropout_s = 0.0

    # ── UWB+IMU heading update (10 Hz) ────────────────────────────────────────

    def update_uwb_heading(
        self,
        heading_rad: float,
        sigma_rad:   float | None = None,
    ) -> None:
        """Update heading from UWB+IMU fused estimate."""
        if not self._init:
            return

        R_val = (sigma_rad ** 2) if sigma_rad is not None else self._R_uwb_head
        R     = np.array([[R_val]])

        H     = np.zeros((1, N), dtype=np.float64)
        H[0, IT] = 1.0

        innov = np.array([_wrap(heading_rad - self._x[IT])])
        _ekf_update(self._x, self._P, H, R, innov)
        self._x[IT] = _wrap(self._x[IT])

    # ── Visual-odometry velocity update (30 Hz) ───────────────────────────────

    def update_vo_velocity(
        self,
        vx_body: float,
        omega:   float,
        sigma_vel:   float | None = None,
        sigma_omega: float | None = None,
    ) -> None:
        """
        Update from visual-odometry velocity measurement.

        VO gives forward speed (vx_body, m/s) in the robot body frame, and
        yaw rate (omega, rad/s).  We rotate vx_body to world frame using the
        current heading estimate, then update [vx, vy, ω] state.

        Parameters
        ----------
        vx_body   : VO forward linear speed (m/s)
        omega     : VO yaw rate (rad/s)
        sigma_vel : std-dev for linear velocity component (m/s)
        sigma_omega: std-dev for yaw rate (rad/s)
        """
        if not self._init:
            return

        R_vel   = (sigma_vel   ** 2) if sigma_vel   is not None else self._R_vo_vel
        R_omega = (sigma_omega ** 2) if sigma_omega is not None else self._R_vo_omega

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

        # Rotate body-frame vx to world-frame [vx_w, vy_w]
        vx_w = vx_body * ct
        vy_w = vx_body * st

        # Measurement: [vx_w, vy_w, ω]
        z = np.array([vx_w, vy_w, omega])

        H = np.zeros((3, N), dtype=np.float64)
        H[0, IVX] = 1.0
        H[1, IVY] = 1.0
        H[2, IW]  = 1.0

        R = np.diag([R_vel, R_vel, R_omega])

        innov = z - H @ self._x
        _ekf_update(self._x, self._P, H, R, innov)
        self._x[IT] = _wrap(self._x[IT])

        self._vo_dropout_s = 0.0

    # ── 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 sub-covariance [x, y]."""
        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._init

    @property
    def uwb_dropout_s(self) -> float:
        return self._uwb_dropout_s

    @property
    def vo_dropout_s(self) -> float:
        return self._vo_dropout_s

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


# ── EKF update helper ──────────────────────────────────────────────────────────

def _ekf_update(
    x: np.ndarray,
    P: np.ndarray,
    H: np.ndarray,
    R: np.ndarray,
    innov: np.ndarray,
) -> None:
    """Standard linear EKF update step — modifies x and P in place."""
    S = H @ P @ H.T + R
    K = P @ H.T @ np.linalg.inv(S)
    x += K @ innov
    I_KH = np.eye(N, dtype=np.float64) - K @ H
    # Joseph form for numerical stability
    P[:] = I_KH @ P @ I_KH.T + K @ R @ K.T
