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, 0]).transpose()
		self.Xp = Xp
		self.X = X

		np.set_printoptions(precision=3)

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

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

	def process_truth(self):
		# ----------------------------
		# Process ground truth
		self.X = self.A * self.X

		return self.X

	def process_measurement(self, y):

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

		return Z

	def process(self, Z):

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

		# ----------------------------
		# 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

		return self.Xp

# Main
if __name__ == '__main__':
	dt =1

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

	_x1 = np.empty(0)
	_y1 = np.empty(0)
	_x2 = np.empty(0)
	_y2 = np.empty(0)

	N = int(100/dt)
	seqn = range(0, N)

	for n in seqn:
		X = k.process_truth()
		Z = k.process_measurement((X[0,0], 0))
		#Kalman.print("Z:", Z)
		Xp = k.process(Z)

		_x1 = np.append(_x1, Z[0])
		_x2 = np.append(_x2, Z[1])
		_y1 = np.append(_y1, Xp[0])
		_y2 = np.append(_y2, Xp[1])

	figure(1)
	subplot(2, 1, 1)
	plot(seqn, _x1, 'bx', seqn, _y1, '-r', linewidth=1)
	grid(True)
	subplot(2, 1, 2)
	plot(seqn, _x2, 'bx', seqn, _y2, '-r', linewidth=1)
	grid(True)
	show()

	print("End of program")

