# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_kalman_filter.ipynb.

# %% auto 0
__all__ = ['KalmanFilter']

# %% ../nbs/02_kalman_filter.ipynb 3
import numpy as np
import scipy.linalg

# %% ../nbs/02_kalman_filter.ipynb 4
class KalmanFilter:
    """
    A Kalman filter class designed for tracking bounding boxes in image space.

    Attributes:

    - `ndim` (int): The dimension of the state space.
    - `_motion_mat` (ndarray): The motion model matrix.
    - `_update_mat` (ndarray): The update matrix used for projecting state distribution to measurement space.
    - `_std_weight_position` (float): Standard deviation weight for the position.
    - `_std_weight_velocity` (float): Standard deviation weight for the velocity.
    """

    def __init__(self):
        """Initialize the Kalman filter with default parameters."""
        # Define the dimension of state space (x, y, dx, dy)
        self.ndim = 4
        dt = 1.0  # Time step

        # Define the motion model matrix
        self._motion_mat = np.eye(2 * self.ndim, 2 * self.ndim)
        self._motion_mat[:self.ndim, self.ndim:] = dt * np.eye(self.ndim)
        # Define the update matrix
        self._update_mat = np.eye(self.ndim, 2 * self.ndim)

        # Define weights for calculating uncertainties
        self._std_weight_position = 1. / 20
        self._std_weight_velocity = 1. / 160

    def _create_std(self, 
                    mean:np.ndarray # The mean values.
                   ) -> np.ndarray: # The computed standard deviations.
        """
        Compute standard deviations based on the mean. 
        """
        std_pos = self._std_weight_position * mean[..., 3]
        std_vel = self._std_weight_velocity * mean[..., 3]
        ones = np.ones_like(std_pos)

        return np.stack([std_pos, std_pos, 1e-2 * ones, std_pos, std_vel, std_vel, 1e-5 * ones, std_vel], axis=-1)

    def initiate(self, 
                 measurement: np.ndarray # The initial measurement for the track. 
                ) -> tuple: # The mean and covariance of the initiated track. 
        """
        Initialize a new track from an unassociated measurement. 
        """
        mean = np.hstack([measurement, np.zeros_like(measurement)])
        covariance = np.diag(self._create_std(mean)**2)
        return mean, covariance

    def predict(self, 
                mean:np.ndarray, # The current state mean. 
                covariance:np.ndarray # The current state covariance.
               ) -> tuple: # The predicted state mean and covariance.
        """
        Run the Kalman filter prediction step.
        """
        motion_cov = np.diag(self._create_std(mean)**2)

        mean = mean @ self._motion_mat.T
        covariance = self._motion_mat @ covariance @ self._motion_mat.T + motion_cov
        return mean, covariance

    def project(self, 
                mean:np.ndarray, # The current state mean.
                covariance:np.ndarray # The current state covariance.
               ) -> tuple: # The mean and covariance in the measurement space.
        """
        Project the state distribution to the measurement space.
        """
        innovation_cov = np.diag(self._create_std(mean)[:self.ndim]**2)

        mean = self._update_mat @ mean
        covariance = self._update_mat @ covariance @ self._update_mat.T + innovation_cov
        return mean, covariance

    def multi_predict(self, 
                      mean:np.ndarray, # The current state mean.
                      covariance:np.ndarray # The current state covariance.
                     ) -> tuple: # The predicted state mean and covariance for multiple measurements.
        """
        Run the Kalman filter prediction step for multiple measurements (Vectorized version).
        """
        motion_cov = np.array([np.diag(s) for s in self._create_std(mean)**2])

        mean = mean @ self._motion_mat.T
        covariance = self._motion_mat @ covariance @ self._motion_mat.T[np.newaxis, ...] + motion_cov
        return mean, covariance

    def update(self, 
               mean:np.ndarray, # The predicted state mean.
               covariance:np.ndarray, # The predicted state covariance.
               measurement:np.ndarray # The new measurement.
              ) -> tuple: # The updated state mean and covariance after correction.
        """
        Run the Kalman filter correction step.
        """
        projected_mean, projected_cov = self.project(mean, covariance)

        chol_factor, lower = scipy.linalg.cho_factor(projected_cov, lower=True, check_finite=False)
        kalman_gain = scipy.linalg.cho_solve((chol_factor, lower), self._update_mat @ covariance, check_finite=False).T

        new_mean = mean + (measurement - projected_mean) @ kalman_gain.T
        new_covariance = covariance - kalman_gain @ projected_cov @ kalman_gain.T
        return new_mean, new_covariance

    def gating_distance(self, 
                        mean:np.ndarray, # The state mean.
                        covariance:np.ndarray, # The state covariance.
                        measurements:np.ndarray, # The given measurements.
                        only_position:bool=False, # If True, consider only position in the gating distance. Defaults to False.
                        metric:str='maha' # The metric to use for distance calculation ('gaussian' or 'maha'). Defaults to 'maha'.
                       ) -> np.ndarray: # The gating distances.
        """
        Compute the gating distance between the state distribution and given measurements.

        Raises:
            ValueError: If an invalid distance metric is provided.
        """
        mean, covariance = self.project(mean, covariance)

        if only_position:
            mean, covariance = mean[:2], covariance[:2, :2]
            measurements = measurements[:, :2]

        diff = measurements - mean
        if metric == 'gaussian':
            return np.sum(diff**2, axis=1)
        elif metric == 'maha':
            chol_factor = np.linalg.cholesky(covariance)
            z = scipy.linalg.solve_triangular(chol_factor, diff.T, lower=True, check_finite=False)
            return np.sum(z**2, axis=0)
        else:
            raise ValueError('Invalid distance metric')
