"""
Kalman filter implementation
"""
import cv2
import numpy as np


class KalmanFilter:
    """Linear Kalman Filter - Constant velocity model
    """
    def __init__(self, meas_size, init_measurement, dT=None):
        """
        Args:
            meas_size (int): size of measurement space
            init_measurement (np.ndarray): initial measurement
            dT (float, optional): if elapsed time between updates is constant, if None -
                                  dT is measured automatically. Defaults to None.
        """
        # every measured state has its velocity, state size is twice in size
        self.meas_size = meas_size
        self.state_size = 2 * meas_size
        # kalman filter object
        self.kf = cv2.KalmanFilter(self.state_size, self.meas_size, 0)
        # transition, measurement  matrices
        self.kf.transitionMatrix = np.eye(self.kf.transitionMatrix.shape[0], dtype=np.float32)
        self.kf.measurementMatrix = np.zeros((self.meas_size, self.state_size), dtype=np.float32)
        self.kf.measurementMatrix[np.arange(self.meas_size), np.arange(self.meas_size)] = 1.0
        # process and measurement noise covariances
        self.kf.processNoiseCov[np.arange(self.meas_size), np.arange(self.meas_size)] = 5.0
        self.kf.processNoiseCov[np.arange(self.meas_size, self.state_size), np.arange(self.meas_size, self.state_size)] = 25.0
        self.kf.measurementNoiseCov = np.eye(self.kf.measurementNoiseCov.shape[0], dtype=np.float32) * 2.5
        # error covariance prior to correction
        self.kf.errorCovPre = 30.0 * (np.eye(self.kf.errorCovPre.shape[0], dtype=np.float32) +
                                      np.eye(self.kf.errorCovPre.shape[0], k=meas_size, dtype=np.float32) +
                                      np.eye(self.kf.errorCovPre.shape[0], k=-meas_size, dtype=np.float32))
        # initialize state from initial measurement
        state = np.zeros_like((self.kf.statePre), dtype=np.float32)
        state[:init_measurement.shape[0], :] = init_measurement[:, np.newaxis]
        self.kf.statePost = state
        # set constant dT if passed, otherwise setup timer
        self.dt = dT
        if self.dt is None:
            self.ticks = 0.0
            self.tick_freq = cv2.getTickFrequency()

    def update(self, measurement):
        """Update Kalman fileter with measurement

        Args:
            measurement (np.array): measured vector
        """
        measurement = np.array(measurement, dtype=np.float32).reshape(-1, 1)
        self.kf.correct(measurement)

    def predict(self):
        """Predict measurement

        Returns:
            np.array: predicted measurement vector
        """
        # if dT is not constant, calculate it from timer
        if self.dt is None:
            prev_ticks = self.ticks
            self.ticks = cv2.getTickCount()
            dT = (self.ticks - prev_ticks) / self.tick_freq
        else:
            dT = self.dt
        # set dT to tensition matrix
        self.kf.transitionMatrix[np.arange(self.meas_size), np.arange(self.meas_size, self.state_size)] = dT
        # predict output state
        state = self.kf.predict()[:self.meas_size, 0]
        return state

    def project(self):
        """Project current state and covariance to measurement space

        Returns:
            np.array, np.ndarray: _description_
        """
        # mean is measurement state
        mean = self.kf.predict()[:self.meas_size, 0]
        # project covariance to measurement space HPH' and add measurement covariance R
        covariance = np.linalg.multi_dot((self.kf.measurementMatrix, self.kf.errorCovPre, self.kf.measurementMatrix.T)) + self.kf.measurementNoiseCov
        return mean, covariance
