"""
Adaptive 2D Kalman Filter — 6-State Constant Acceleration Model
================================================================
State vector: [x, y, vx, vy, ax, ay]
Measurement:  [x, y]

Features:
    - NIS (Normalized Innovation Squared) based adaptive process noise
    - Automatic maneuver detection and Q-matrix scaling
    - Confidence scoring based on covariance trace
    - Covariance ellipse parameters for visualization
    - Pure NumPy — no external dependencies beyond numpy

Author: ByIbos
License: MIT
"""

import numpy as np
from collections import deque


class KalmanFilter2D:
    """
    Constant-acceleration 2D Kalman Filter with adaptive process noise.

    The filter uses a 6-dimensional state vector [x, y, vx, vy, ax, ay]
    and a 2-dimensional measurement vector [x, y]. When the target performs
    sudden maneuvers, the NIS (Normalized Innovation Squared) metric detects
    the model mismatch and dynamically scales the process noise covariance
    matrix Q, allowing the filter to adapt in real-time.

    Args:
        dt: Time step between frames in seconds (default: 1/30 for 30 FPS).
        process_noise: Base process noise level. Higher = more responsive
            to measurements but noisier estimates.
        measurement_noise: Measurement noise level. Higher = smoother
            estimates but slower response to real changes.

    Example:
        >>> kf = KalmanFilter2D(dt=1/30, process_noise=1.0, measurement_noise=5.0)
        >>> kf.init_state(320, 240)
        >>> px, py = kf.predict()
        >>> fx, fy = kf.update(325, 238)
    """

    def __init__(self, dt=1/30, process_noise=1.0, measurement_noise=5.0):
        self.dt = dt
        self.base_process_noise = process_noise
        self.measurement_noise = measurement_noise

        # State vector: [x, y, vx, vy, ax, ay]
        self.x = np.zeros(6)

        # State covariance matrix (6x6)
        self.P = np.eye(6) * 500.0

        # State transition matrix (updated each predict)
        self.F = np.eye(6)
        self._update_F(dt)

        # Measurement matrix: we only observe [x, y]
        self.H = np.zeros((2, 6))
        self.H[0, 0] = 1.0  # x
        self.H[1, 1] = 1.0  # y

        # Measurement noise covariance
        self.R = np.eye(2) * measurement_noise ** 2

        # Process noise covariance
        self.Q = self._build_Q(dt, process_noise)

        # NIS (Normalized Innovation Squared) history for adaptive filtering
        self.nis_history = deque(maxlen=20)
        self.nis_threshold = 5.99  # Chi-squared 95% for 2 DOF
        self.adaptive_scale = 1.0

        # Position history for trajectory prediction
        self.position_history = deque(maxlen=120)

        # Track state
        self.initialized = False
        self.predict_count = 0  # Consecutive predictions without update
        self.update_count = 0   # Total updates received

    def _update_F(self, dt):
        """Update the state transition matrix for current dt."""
        self.F = np.eye(6)
        # Position += velocity * dt + 0.5 * acceleration * dt^2
        self.F[0, 2] = dt       # x += vx * dt
        self.F[1, 3] = dt       # y += vy * dt
        self.F[0, 4] = 0.5 * dt ** 2  # x += 0.5 * ax * dt^2
        self.F[1, 5] = 0.5 * dt ** 2  # y += 0.5 * ay * dt^2
        # Velocity += acceleration * dt
        self.F[2, 4] = dt       # vx += ax * dt
        self.F[3, 5] = dt       # vy += ay * dt

    def _build_Q(self, dt, sigma):
        """
        Build the process noise covariance matrix using
        the discrete white noise acceleration model.

        Args:
            dt: Time step (seconds).
            sigma: Process noise standard deviation.

        Returns:
            6x6 process noise covariance matrix.
        """
        dt2 = dt ** 2
        dt3 = dt ** 3
        dt4 = dt ** 4

        # Single-axis block (constant acceleration model)
        q_block = np.array([
            [dt4 / 4, dt3 / 2, dt2 / 2],
            [dt3 / 2, dt2,     dt],
            [dt2 / 2, dt,      1.0]
        ]) * sigma ** 2

        Q = np.zeros((6, 6))
        # X-axis block: indices 0, 2, 4 (x, vx, ax)
        for i, ii in enumerate([0, 2, 4]):
            for j, jj in enumerate([0, 2, 4]):
                Q[ii, jj] = q_block[i, j]
        # Y-axis block: indices 1, 3, 5 (y, vy, ay)
        for i, ii in enumerate([1, 3, 5]):
            for j, jj in enumerate([1, 3, 5]):
                Q[ii, jj] = q_block[i, j]

        return Q

    def init_state(self, x, y):
        """
        Initialize the filter state with a first measurement.

        Args:
            x: Initial x position (pixels).
            y: Initial y position (pixels).
        """
        self.x = np.array([x, y, 0.0, 0.0, 0.0, 0.0], dtype=float)
        self.P = np.eye(6) * 500.0
        self.P[0, 0] = self.measurement_noise ** 2
        self.P[1, 1] = self.measurement_noise ** 2
        self.initialized = True
        self.predict_count = 0
        self.update_count = 0
        self.position_history.clear()
        self.position_history.append((x, y))

    def predict(self):
        """
        Predict the state one step forward using the motion model.

        Returns:
            tuple: (predicted_x, predicted_y) position.
        """
        if not self.initialized:
            return 0.0, 0.0

        self._update_F(self.dt)

        # State prediction: x_k|k-1 = F * x_k-1|k-1
        self.x = self.F @ self.x

        # Covariance prediction: P_k|k-1 = F * P * F^T + Q * adaptive_scale
        self.P = self.F @ self.P @ self.F.T + self.Q * self.adaptive_scale

        self.predict_count += 1
        return self.x[0], self.x[1]

    def update(self, z_x, z_y):
        """
        Update the state with a new measurement (adaptive Q included).

        The NIS metric is computed to detect sudden maneuvers. When the NIS
        exceeds the chi-squared threshold (5.99 for 2 DOF at 95%), the
        process noise is scaled up to allow the filter to track aggressive
        movements.

        Args:
            z_x: Measured x position (pixels).
            z_y: Measured y position (pixels).

        Returns:
            tuple: (filtered_x, filtered_y) — the corrected position.
        """
        if not self.initialized:
            self.init_state(z_x, z_y)
            return z_x, z_y

        z = np.array([z_x, z_y])

        # Innovation (measurement residual)
        y = z - self.H @ self.x

        # Innovation covariance
        S = self.H @ self.P @ self.H.T + self.R

        # NIS (Normalized Innovation Squared)
        try:
            S_inv = np.linalg.inv(S)
            nis = float(y.T @ S_inv @ y)
        except np.linalg.LinAlgError:
            nis = 0.0

        self.nis_history.append(nis)

        # Adaptive Q scaling based on NIS
        if nis > self.nis_threshold:
            # Maneuver detected → loosen the filter
            self.adaptive_scale = min(50.0, self.adaptive_scale * 1.5 + 1.0)
        else:
            # Normal operation → tighten the filter
            self.adaptive_scale = max(1.0, self.adaptive_scale * 0.85)

        # Kalman gain
        K = self.P @ self.H.T @ S_inv

        # State update
        self.x = self.x + K @ y

        # Covariance update (Joseph form for numerical stability)
        I_KH = np.eye(6) - K @ self.H
        self.P = I_KH @ self.P @ I_KH.T + K @ self.R @ K.T

        self.predict_count = 0
        self.update_count += 1

        # Store position for trajectory prediction
        self.position_history.append((self.x[0], self.x[1]))

        return self.x[0], self.x[1]

    @property
    def position(self):
        """Current estimated position (x, y)."""
        return self.x[0], self.x[1]

    @property
    def velocity(self):
        """Current estimated velocity (vx, vy) in pixels/frame."""
        return self.x[2], self.x[3]

    @property
    def acceleration(self):
        """Current estimated acceleration (ax, ay) in pixels/frame²."""
        return self.x[4], self.x[5]

    @property
    def speed(self):
        """Total speed magnitude in pixels/frame."""
        vx, vy = self.velocity
        return float(np.sqrt(vx ** 2 + vy ** 2))

    @property
    def confidence(self):
        """
        Prediction confidence score in [0.0, 1.0].

        Derived from the covariance matrix trace — lower covariance
        means higher confidence.

        Returns:
            float: Confidence score (1.0 = very confident, 0.0 = uncertain).
        """
        if not self.initialized:
            return 0.0

        trace = self.P[0, 0] + self.P[1, 1]
        conf = max(0.0, 1.0 - trace / 2000.0)

        # Reduce confidence during prolonged prediction (no measurements)
        if self.predict_count > 5:
            conf *= max(0.2, 1.0 - self.predict_count * 0.05)

        return min(1.0, conf)

    def covariance_ellipse(self):
        """
        Compute the 95% confidence ellipse parameters for visualization.

        Returns:
            tuple: (width, height, angle_degrees) of the ellipse.
        """
        if not self.initialized:
            return 0, 0, 0

        # Extract position covariance (2x2 sub-matrix)
        P_pos = self.P[:2, :2]

        try:
            eigenvalues, eigenvectors = np.linalg.eigh(P_pos)
            eigenvalues = np.maximum(eigenvalues, 0)

            # 95% confidence interval → chi-squared scale factor
            scale = 2.4477  # sqrt(5.99)
            width = 2 * scale * np.sqrt(eigenvalues[1])
            height = 2 * scale * np.sqrt(eigenvalues[0])

            angle = np.degrees(np.arctan2(eigenvectors[1, 1], eigenvectors[0, 1]))

            return int(width), int(height), angle
        except np.linalg.LinAlgError:
            return 50, 50, 0

    def reset(self):
        """Reset the filter to its uninitialized state."""
        self.x = np.zeros(6)
        self.P = np.eye(6) * 500.0
        self.initialized = False
        self.predict_count = 0
        self.update_count = 0
        self.adaptive_scale = 1.0
        self.nis_history.clear()
        self.position_history.clear()
