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


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


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

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

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

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

	_seqn = range(0, 2*N)
	X = np.array([1, 0])
	for n in seqn:
		Z = k.process_measurement(X, 0.5)
		# 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])

	for n in seqn:
		X[0] = X[0] + 1
		Z = k.process_measurement(X, 0.5)
		# 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")