"""
_landmark_kalman.py — Vectorised constant-velocity Kalman filter for 2-D facial landmarks.

No ROS2 dependencies — pure numpy.

Each face track has one LandmarkKalmanFilter instance that maintains a
(N, 4) state array where each row is [x, y, vx, vy] for one landmark.
All N points are predicted / updated simultaneously using broadcast numpy
operations, so there is no per-point Python loop.

Usage
-----
    kf = LandmarkKalmanFilter(n_points=68)
    smoothed = kf.update(obs, dt)   # obs: (68, 2) ndarray, dt: float (s)
    smoothed = kf.update(obs, dt)   # subsequent calls
"""

from __future__ import annotations

import numpy as np
from typing import Optional


class LandmarkKalmanFilter:
    """
    Constant-velocity Kalman filter over N 2-D landmark points.

    State per point: [x, y, vx, vy]
    Observation per point: [x, y]

    Parameters
    ----------
    n_points          : number of landmarks (default 68)
    process_noise_std : std-dev of velocity random-walk noise  (normalised coords/s)
    obs_noise_std     : std-dev of observation noise           (normalised coords)
    """

    def __init__(
        self,
        n_points: int = 68,
        process_noise_std: float = 0.01,
        obs_noise_std: float = 0.005,
    ) -> None:
        self._n = n_points
        self._q_std = process_noise_std
        self._r_std = obs_noise_std

        # State: (n, 4)  [x, y, vx, vy]
        self._x: Optional[np.ndarray] = None   # initialised on first update

        # Covariance: (n, 4, 4)
        self._P: Optional[np.ndarray] = None

        # Fixed matrices (shared across all points via broadcasting)
        # Observation matrix H: (2, 4)  — selects x and y
        self._H = np.array([[1., 0., 0., 0.],
                             [0., 1., 0., 0.]])          # (2, 4)

        # Observation noise R: (2, 2)
        self._R = (obs_noise_std ** 2) * np.eye(2)       # (2, 2)

    # ── Public API ────────────────────────────────────────────────────────────

    def update(self, obs: np.ndarray, dt: float) -> np.ndarray:
        """
        Predict then update with new observations.

        Parameters
        ----------
        obs : (n_points, 2) array of [x, y] normalised landmark coordinates
        dt  : elapsed time since last call (seconds); unused on the first call

        Returns
        -------
        (n_points, 2) array of smoothed [x, y] coordinates
        """
        obs = np.asarray(obs, dtype=np.float64)
        if obs.shape != (self._n, 2):
            raise ValueError(f'Expected obs shape ({self._n}, 2), got {obs.shape}')

        if self._x is None:
            self._initialise(obs)
        else:
            self._predict(max(dt, 1e-6))
            self._update(obs)

        return self._x[:, :2].astype(np.float32)

    @property
    def is_initialised(self) -> bool:
        return self._x is not None

    # ── Internal ──────────────────────────────────────────────────────────────

    def _initialise(self, obs: np.ndarray) -> None:
        """Initialise state from first observation (zero velocity)."""
        self._x = np.zeros((self._n, 4), dtype=np.float64)
        self._x[:, :2] = obs

        # Start with moderate position uncertainty, zero velocity
        init_pos_var = (self._r_std * 10) ** 2
        self._P = np.zeros((self._n, 4, 4), dtype=np.float64)
        self._P[:, 0, 0] = init_pos_var
        self._P[:, 1, 1] = init_pos_var
        self._P[:, 2, 2] = (self._q_std * 5) ** 2
        self._P[:, 3, 3] = (self._q_std * 5) ** 2

    def _predict(self, dt: float) -> None:
        """Constant-velocity predict step: x' = F x,  P' = F P Fᵀ + Q."""
        # Transition matrix F: (4, 4)
        F = np.array([[1., 0., dt, 0.],
                      [0., 1., 0., dt],
                      [0., 0., 1., 0.],
                      [0., 0., 0., 1.]])

        # Process noise Q: (4, 4)  — velocity driven (discrete Wiener model)
        q = self._q_std ** 2
        dt2 = dt * dt
        dt3 = dt2 * dt
        dt4 = dt3 * dt
        Q = q * np.array([[dt4/4, 0.,    dt3/2, 0.   ],
                           [0.,    dt4/4, 0.,    dt3/2],
                           [dt3/2, 0.,    dt2,   0.   ],
                           [0.,    dt3/2, 0.,    dt2  ]])

        # Vectorised over n points: (n,4) @ (4,4).T  → (n,4)
        self._x = self._x @ F.T                              # (n, 4)

        # P' = F P Fᵀ + Q  for each of n points
        # (n,4,4) → matmul with (4,4) on both sides
        FP  = np.einsum('ij,njk->nik', F, self._P)          # (n, 4, 4)
        FPFt = np.einsum('nij,kj->nik', FP, F)              # (n, 4, 4)
        self._P = FPFt + Q[np.newaxis]                       # (n, 4, 4)

    def _update(self, obs: np.ndarray) -> None:
        """Kalman update step with measurement obs: (n, 2)."""
        H  = self._H                              # (2, 4)
        R  = self._R                              # (2, 2)

        # Innovation: y = obs − H x  →  (n, 2)
        Hx = np.einsum('ij,nj->ni', H, self._x)  # (n, 2)
        y  = obs - Hx                             # (n, 2)

        # Innovation covariance S = H P Hᵀ + R  →  (n, 2, 2)
        HP  = np.einsum('ij,njk->nik', H, self._P)           # (n, 2, 4)
        S   = np.einsum('nij,kj->nik', HP, H) + R[np.newaxis]  # (n, 2, 2)

        # Kalman gain K = P Hᵀ S⁻¹  →  (n, 4, 2)
        PHt = np.einsum('nij,kj->nik', self._P, H)           # (n, 4, 2)
        S_inv = np.linalg.inv(S)                              # (n, 2, 2)
        K   = np.einsum('nij,njk->nik', PHt, S_inv)          # (n, 4, 2)

        # State update: x += K y
        Ky = np.einsum('nij,nj->ni', K, y)                   # (n, 4)
        self._x = self._x + Ky

        # Covariance update: P = (I - K H) P
        n = self._n
        I = np.eye(4)[np.newaxis].repeat(n, axis=0)          # (n, 4, 4)
        KH = np.einsum('nij,jk->nik', K, H)                  # (n, 4, 4)
        self._P = np.einsum('nij,njk->nik', I - KH, self._P) # (n, 4, 4)
