import jax.numpy as jnp
import numpy as np

class KalmanFilter:
    def __init__(self, A, B, Q, C, R):
        self.params = dict()
        self._init_params(A, B, Q, C, R)

    def _init_params(self, A, B, Q, C, R):
        self.params["A"] = A # process matrix
        self.params["B"] = B # control matrix
        self.params["Q"] = Q # symetric postiive definite process noise covariance
        self.params["C"] = C # measurement matrix 
        self.params["R"] = R # symetric postiive definite measurement noise covariance

    def dynamic(self, x, u):
        return jnp.add(jnp.add(jnp.dot(self.params["A"], x), jnp.dot(self.params["B"], u)), jnp.dot(jnp.linalg.cholesky(self.params["Q"]),np.random.randn(len(x))))

    def observation(self, x):
        return jnp.add(jnp.dot(self.params["C"],x), jnp.dot(jnp.linalg.cholesky(self.params["R"]),np.random.randn(len(x))))

    def prediction(self, s, u):
        mu, sigma = s # unroll states
        mu_hat = jnp.add(jnp.dot(self.params["A"], mu), jnp.dot(self.params["B"], u))
        sigma_hat = jnp.add(jnp.dot(jnp.dot(self.params["A"], sigma), jnp.transpose(self.params["A"])), self.params["Q"])
        return (mu_hat, sigma_hat)
    
    def correction(self, s, y):
        mu, sigma = s # unroll states
        v = jnp.subtract(y, jnp.dot(self.params["C"], mu)) # innovation
        S = jnp.add(jnp.dot(jnp.dot(self.params["C"], sigma), jnp.transpose(self.params["C"])), self.params["R"]) # innovation covariance
        K = jnp.dot(jnp.dot(sigma, jnp.transpose(self.params["C"])), jnp.linalg.inv(S)) # Kalman gain
        new_mu = jnp.add(mu, jnp.dot(K, v)) # update state estimation
        new_sigma = jnp.dot(jnp.subtract(jnp.eye(len(mu)), jnp.dot(K, self.params["C"])), sigma) # update state covariance estimation
        return (new_mu, new_sigma)
    
    def update(self, s, u, y):
        mu, sigma = s
        s = self.prediction(s, u)
        s = self.correction(s, y)
        return s

    def mse_loss(self, s_hat, s_post):
        mu_hat, sigma_hat = s_hat
        mu_post, sigma_post = s_post
        # ex = jnp.linalg.norm(jnp.subtract(mu, jnp.add(jnp.dot(self.params["A"], mu), jnp.dot(self.params["B"], u))))
        # ey = jnp.linalg.norm(jnp.subtract(y, jnp.dot(self.params["C"], mu)))
        return jnp.linalg.norm(mu_hat - mu_post)
    

