"""
kalman_odom_filter.py — Extended Kalman Filter for odometry fusion.

Fuses wheel odometry and visual odometry into a single best-estimate
ground-truth trajectory.  Handles wheel-slip failover by down-weighting
the wheel odometry measurement covariance when slip is detected.

State vector:   x = [px, py, θ, v, ω]
  px, py  — position in map/odom frame (m)
  θ       — heading (rad), wrapped to [-π, π]
  v       — linear velocity (m/s, forward)
  ω       — angular velocity (rad/s)

Process model (unicycle, constant velocity over dt):
  px' = px + v·cos(θ)·dt
  py' = py + v·sin(θ)·dt
  θ'  = θ  + ω·dt
  v'  = v              (constant velocity prediction)
  ω'  = ω              (constant angular velocity prediction)

Measurement models:
  Wheel odom:    z_w = [v_wheel, ω_wheel]       (velocity measurements)
  Visual odom:   z_v = [v_visual, ω_visual]     (velocity measurements)
  RTAB-Map pose: z_r = [px_rtab, py_rtab, θ_rtab] (absolute pose correction)

Covariances:
  Process noise Q is fixed (low).
  Wheel measurement noise R_w increases 10× during slip.
  Visual measurement noise R_v increases 3× when VO returns invalid.

Units: SI throughout (m, rad, m/s, rad/s).
"""

from __future__ import annotations
import math
from typing import Tuple

import numpy as np


class KalmanOdomFilter:
    """
    5-state EKF for ground-robot odometry fusion.

    Usage:
        kf = KalmanOdomFilter()
        kf.predict(dt)
        kf.update_wheel(v, omega)         # wheel odom measurement
        kf.update_visual(v, omega)        # visual odom measurement
        kf.update_rtabmap(px, py, theta)  # absolute RTAB-Map correction
        px, py, theta = kf.pose
    """

    _N = 5   # state dimension

    def __init__(
        self,
        # Process noise (standard deviations)
        q_pos:   float = 0.02,   # m per step
        q_theta: float = 0.01,   # rad per step
        q_vel:   float = 0.10,   # m/s per step
        q_omega: float = 0.05,   # rad/s per step
        # Wheel odom noise (normal)
        r_wheel_v:     float = 0.05,   # m/s
        r_wheel_omega: float = 0.05,   # rad/s
        # Visual odom noise
        r_visual_v:    float = 0.08,
        r_visual_omega:float = 0.08,
        # RTAB-Map pose noise
        r_rtab_pos:    float = 0.10,   # m
        r_rtab_theta:  float = 0.05,   # rad
        # Slip multiplier for wheel noise
        slip_noise_multiplier: float = 10.0,
    ):
        self._slip_mult = slip_noise_multiplier

        # State estimate
        self._x = np.zeros(self._N, dtype=np.float64)   # [px, py, theta, v, omega]

        # Covariance
        self._P = np.eye(self._N, dtype=np.float64) * 0.01

        # Process noise covariance Q
        self._Q = np.diag([q_pos**2, q_pos**2, q_theta**2, q_vel**2, q_omega**2])

        # Measurement noise matrices
        self._R_wheel_normal = np.diag([r_wheel_v**2, r_wheel_omega**2])
        self._R_wheel_slip   = np.diag([
            (r_wheel_v   * slip_noise_multiplier)**2,
            (r_wheel_omega * slip_noise_multiplier)**2,
        ])
        self._R_visual_valid   = np.diag([r_visual_v**2, r_visual_omega**2])
        self._R_visual_invalid = np.diag([
            (r_visual_v   * 3)**2,
            (r_visual_omega * 3)**2,
        ])
        self._R_rtabmap = np.diag([r_rtab_pos**2, r_rtab_pos**2, r_rtab_theta**2])

    # ── Public interface ──────────────────────────────────────────────────────

    @property
    def pose(self) -> Tuple[float, float, float]:
        """Return (x, y, theta) current best estimate."""
        return float(self._x[0]), float(self._x[1]), float(self._x[2])

    @property
    def velocity(self) -> Tuple[float, float]:
        """Return (v_linear, omega_angular) current estimate."""
        return float(self._x[3]), float(self._x[4])

    @property
    def covariance_3x3(self) -> np.ndarray:
        """Return 3×3 position+heading covariance (px, py, theta sub-block)."""
        return self._P[:3, :3].copy()

    def predict(self, dt: float) -> None:
        """Advance the state estimate by dt seconds (unicycle motion model)."""
        px, py, th, v, w = self._x
        c, s = math.cos(th), math.sin(th)

        # State transition
        self._x[0] = px + v * c * dt
        self._x[1] = py + v * s * dt
        self._x[2] = self._wrap(th + w * dt)
        # v and omega unchanged (constant velocity prediction)

        # Jacobian F of state transition
        F = np.eye(self._N, dtype=np.float64)
        F[0, 2] = -v * s * dt
        F[0, 3] =  c * dt
        F[1, 2] =  v * c * dt
        F[1, 3] =  s * dt
        F[2, 4] =  dt

        # Covariance prediction
        self._P = F @ self._P @ F.T + self._Q

    def update_wheel(self, v: float, omega: float, slipping: bool = False) -> None:
        """
        Fuse wheel-odometry velocity measurement.

        Args:
            v:        linear velocity from encoders (m/s)
            omega:    angular velocity from encoders (rad/s)
            slipping: if True, use high-noise covariance for wheel measurement
        """
        R = self._R_wheel_slip if slipping else self._R_wheel_normal
        z = np.array([v, omega], dtype=np.float64)
        H = np.zeros((2, self._N), dtype=np.float64)
        H[0, 3] = 1.0   # measure v
        H[1, 4] = 1.0   # measure omega
        self._measurement_update(z, H, R)

    def update_visual(self, v: float, omega: float, valid: bool = True) -> None:
        """
        Fuse visual-odometry velocity measurement.

        Args:
            v, omega: velocity from VO
            valid:    False if VO returned invalid (low feature count etc.)
        """
        R = self._R_visual_valid if valid else self._R_visual_invalid
        z = np.array([v, omega], dtype=np.float64)
        H = np.zeros((2, self._N), dtype=np.float64)
        H[0, 3] = 1.0
        H[1, 4] = 1.0
        self._measurement_update(z, H, R)

    def update_rtabmap(self, px: float, py: float, theta: float) -> None:
        """
        Fuse an absolute pose from RTAB-Map (loop closure correction).

        Uses the full pose measurement [px, py, theta] against the current
        position state.
        """
        z = np.array([px, py, theta], dtype=np.float64)
        H = np.zeros((3, self._N), dtype=np.float64)
        H[0, 0] = 1.0   # measure px
        H[1, 1] = 1.0   # measure py
        H[2, 2] = 1.0   # measure theta

        # Innovation with angle wrapping
        innov = z - H @ self._x
        innov[2] = self._wrap(innov[2])

        self._measurement_update_with_innov(innov, H, self._R_rtabmap)

    # ── Internal EKF machinery ────────────────────────────────────────────────

    def _measurement_update(
        self, z: np.ndarray, H: np.ndarray, R: np.ndarray
    ) -> None:
        innov = z - H @ self._x
        self._measurement_update_with_innov(innov, H, R)

    def _measurement_update_with_innov(
        self, innov: np.ndarray, H: np.ndarray, R: np.ndarray
    ) -> None:
        S   = H @ self._P @ H.T + R
        K   = self._P @ H.T @ np.linalg.inv(S)
        self._x = self._x + K @ innov
        self._x[2] = self._wrap(self._x[2])
        self._P = (np.eye(self._N) - K @ H) @ self._P
        # Joseph form for numerical stability
        I_KH = np.eye(self._N) - K @ H
        self._P = I_KH @ self._P @ I_KH.T + K @ R @ K.T

    @staticmethod
    def _wrap(angle: float) -> float:
        return (angle + math.pi) % (2 * math.pi) - math.pi
