import numpy as np
import numpy.linalg as la
import scipy.stats as stats

from .bayesian_single_target_tracker import BayesianSingleTargetTracker
from utils import maharanobis_distance


class NNKF(BayesianSingleTargetTracker):
    """Nearest neighbor kalman filter."""

    def __init__(self, model, covar_w, covar_v, x, covar_x, z, PG):

        self.PG = PG  # gate threshold probability
        self.gate_thresh = None

        # linear state space model of the plant
        # x_{t+1} = F_t * x_t + G_t * u_t + L * w_t
        # y_{t} = H_t * x_t + v_t
        self.model = model

        # covariance
        self.Q = covar_w
        self.R = covar_v

        # expectaion to estimate.
        self.x_pred = x
        self.x_est = None

        self.z_pred = z

        # covariance to estimate.
        self.P_pred = covar_x
        self.P_est = None

        # kalman gain
        self.Gain = None

        # innvation: z - z_pred
        self.innovation = None
        self.valid_innovations = None
        self.covar_innovation = None

    def extract_valid_measurments(self, z, z_hat, inv_covar):
        """Extract measurments in the validation region."""

        if z.shape[0] == 0:
            # shape: (0, nz)
            return np.empty_like(z), np.empty_like(z)

        distances = maharanobis_distance(z, z_hat, inv_covar)

        if not self.gate_thresh:
            self.gate_thresh = stats.chi2.ppf(q=self.PG, df=z_hat.shape[0])

        mask = distances <= self.gate_thresh
        return z[mask], distances[mask]

    def associate(self, t, z):
        """Associate the measurements with the predicted sate variable."""

        P = self.P_pred
        H = self.model.H
        R = self.R

        self.covar_innovation = H @ P @ H.T + R
        inv_covar = la.inv(self.covar_innovation)
        z, d = self.extract_valid_measurments(z, self.z_pred, inv_covar)
        print(f"step {t} : the number of valid measurments is {z.shape[0]}")

        if z.shape[0] == 0:
            self.valid_innovations = None
            self.innovation = None
        else:
            self.valid_innovations = z - self.z_pred
            self.innovation = self.valid_innovations[np.argmin(d)]

    def filtering(self):
        """Estimate the state variable."""

        # extract variables.
        x, P = self.x_pred, self.P_pred
        S = self.covar_innovation
        H = self.model.H

        # update the gain.
        W = P @ H.T @ np.linalg.inv(S)
        self.Gain = W  # store

        if self.valid_innovations is None:
            self.x_est = x
            self.P_est = P
            return

        # x(t|t), P(t|t)
        self.x_est = x + W @ self.innovation
        self.P_est = (np.eye(W.shape[0]) - W @ H) @ P

    def predict(self, t, u):
        """Predict the sate variable."""

        x, P = self.x_est, self.P_est
        F = self.model.F
        L = self.model.L
        Q = self.Q

        # x(t+1|t), P(t+1|t)
        self.x_pred = self.model.state_equation(t, x, u)
        self.P_pred = F @ P @ F.T + L @ Q @ L.T

        # z(t+1|t)
        self.z_pred = self.model.observation_equation(t+1, self.x_pred)
