import numpy as np
from numpy.linalg import inv
from matplotlib.pyplot import plot, figure, subplot, title, xlabel, ylabel, grid, show

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

		P = var_P*np.eye(N)
		R = var_R*np.eye(N)
		H = np.eye(N)

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

		G = np.matrix(model[N:0:-1])
		Q = G * G.transpose() * var_Q

		self.P = P
		self.Q = Q
		self.R = R
		self.N = N
		self.A = A
		self.H = H
		self.var_Z = var_Z

		X = np.matrix([0, 1.0]).transpose()
		Xp = np.matrix([0, -1.0]).transpose()
		self.Xp = Xp
		self.X = X

		np.set_printoptions(precision=3)

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

	def process(self):

		_x1 = np.empty(0)
		_y1 = np.empty(0)
		_x2 = np.empty(0)
		_y2 = np.empty(0)
		for n in range(0, 100):

			# ----------------------------
			# State estimate
			self.Xp = self.A * self.Xp

			# ----------------------------
			# Process ground truth
			self.X = self.A * self.X

			# ----------------------------
			# Take noisy measurement
			Z = self.H * self.X + self.var_Z * np.random.randn(self.N, 1)

			# ----------------------------
			# Measurement prediction
			Zp = self.H * self.Xp

			# ----------------------------
			# Measurement residual
			V = Z - Zp

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

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

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

			# ----------------------------
			# Update state estimate
			self.Xp = self.Xp + K * V

			# ----------------------------
			# Updated state covariance
			self.P = self.P - K * S * K

			# ----------------------------
			# Plot vars
			_x1 = np.append(_x1, Z[0])
			_x2 = np.append(_x2, Z[1])
			_y1 = np.append(_y1, self.Xp[0])
			_y2 = np.append(_y2, self.Xp[1])

		n = range(0,100)
		figure(1)
		subplot(2, 1, 1)
		plot(n, _x1, 'bx', n, _y1, '-r', linewidth=1)
		grid(True)
		subplot(2, 1, 2)
		plot(n, _x2, 'bx', n, _y2, '-r', linewidth=1)
		grid(True)
		show()
# Main
if __name__ == '__main__':

	params = {
		'dt'	: 1,
		'var_P'	: 1,
		'var_Q'	: .001,
		'var_R'	: 1,
		'var_Z'	: 1
	}
	k = Kalman(params)

	k.process()

	print("End of program")