import numpy as np
from numpy.linalg import inv


class Kalman:
	def __init__(self, dt, params):
		var_P = params['var_P']
		var_Q = params['var_Q']
		var_R = params['var_R']
		model = np.matrix([1, dt, 1/2*dt**2]).transpose()
		N = len(model)-1

		# Process Covariance Matrix
		self.P = var_P*np.eye(N)

		# Sensor Noise Covariance Matrix
		self.R = var_R*np.eye(N)

		self.H = np.eye(N)

		self.A = np.eye(N)
		for row in range(0, N):
			self.A[row, row:N] = model.transpose()[0, 0:N-row]

		G = np.matrix(model[N:0:-1])

		# Process Noise Covariance Matrix
		self.Q = G * G.transpose() * var_Q

		# State Matrix
		self.X = np.matrix([0, 0]).transpose()

		self.N = N

		np.set_printoptions(precision=3)

	@staticmethod
	def print(p, d):
		print(p)
		print(d)

	def initial(self, X):
		self.X = np.matrix([X[0], X[1]]).transpose()

	def process_measurement(self, y, var_Z):
		# ----------------------------
		# Take noisy measurement
		Y = self.H * np.matrix([y[0], y[1]]).transpose() + var_Z * np.random.randn(self.N, 1)

		return Y

	def process(self, Y):

		# ----------------------------
		# Predict State estimate
		X = self.A * self.X

		# ----------------------------
		# Predict State	covariance
		P = self.A * self.P * self.A.transpose() + self.Q

		# ----------------------------
		# Measurement prediction covariance
		S = self.H * P * self.H.transpose() + self.R

		# ----------------------------
		# Kalman gain
		K = P * self.H.transpose() * inv(S)

		# Update state estimate
		self.X = X + K*(Y - self.H * X)

		# ----------------------------
		# Updated state covariance
		I = np.eye(self.N)
		self.P = (I - K * self.H) * P
		return self.X