import numpy as np
import matplotlib.pyplot as plt

n = 200 # number of timesteps
T = 100 # time will vary from 0 to T with step delt
ts = np.linspace(0,T,n+1)
delt = T/n
gamma = 0.05 # damping, 0 is no damping

A = np.zeros((4,4))
C = np.zeros((2,4))

A[0,0] = 1
A[1,1] = 1
A[0,2] = (1-gamma*delt/2)*delt
A[1,3] = (1-gamma*delt/2)*delt
A[2,2] = 1 - gamma*delt
A[3,3] = 1 - gamma*delt

C[0,0] = 1
C[1,1] = 1

from scipy.linalg import sqrtm 

np.random.seed(7030)

x = np.zeros((4,n+1))
x[:,0] = [0,0,0,0]
y = np.zeros((2,n))

Q = np.diag([0.01,0.01,0.1,0.1])
R = np.diag([1,1])
Qh = sqrtm(Q)
Rh = sqrtm(R)

w = Qh@np.random.randn(4,n)
v = Rh@np.random.randn(2,n)

for t in range(n):
    x[:,t+1] = A.dot(x[:,t]) + w[:,t]
    y[:,t]   = C.dot(x[:,t]) + v[:,t]
    
x_true = x.copy()
w_true = w.copy()
v_true = v.copy()

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x[0,:])
plt.plot(ts[:-1],y[0,:], alpha=0.5)
plt.xlabel('time')
plt.ylabel('x position')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x[1,:])
plt.plot(ts[:-1],y[1,:], alpha=0.5)
plt.xlabel('time')
plt.ylabel('y position')
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x[2,:])
plt.xlabel('time')
plt.ylabel('x velocity')
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x[3,:])
plt.xlabel('time')
plt.ylabel('y velocity')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x[0,:],x[1,:],'-',alpha=0.8, linewidth=5, label='True')
plt.plot(y[0,:],y[1,:],'ko',alpha=0.1, label='Measurement')
plt.title('Trajectory')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.show()

import scipy.sparse as ssp
import scipy.sparse.linalg as sla

tau = 1

G = np.zeros( (2*n,4) )
for i in range(n):
  G[2*i:2*(i+1),:] = C@np.linalg.matrix_power(A,i)
  
H = np.zeros( (2*n,4*n) )
H_first = np.zeros( (2*n,4) )
for i in range(1,n):
  H_first[2*i:2*(i+1),:] = C@np.linalg.matrix_power(A,i-1)
for i in range(n):
  H[2*(i+1):,4*i:4*(i+1)] = H_first[2:2*(n-i),:]

A_1 = ssp.hstack(( ssp.csr_matrix((4*n,4)), ssp.eye(4*n),              ssp.csr_matrix((4*n,2*n)) ))
A_2 = ssp.hstack(( ssp.csr_matrix((2*n,4)), ssp.csr_matrix((2*n,4*n)), np.sqrt(tau)*ssp.eye(2*n) ))
A_tilde = ssp.vstack((A_1, A_2))
b_tilde = np.zeros(A_tilde.shape[0])
C_tilde = ssp.hstack(( G, H, ssp.eye(2*n) ))
d_tilde = y.T.flatten()

S_1 = ssp.hstack(( A_tilde.T@A_tilde, C_tilde.T                 ))
S_2 = ssp.hstack(( C_tilde,           ssp.csr_matrix((2*n,2*n)) ))
S = ssp.vstack(( S_1, S_2 ))
r = np.hstack(( A_tilde.T.dot(b_tilde), d_tilde ))
sol = sla.spsolve(S,r)

x_hat = sol[:A_tilde.shape[1]]

x0_hat = x_hat[:4]

w_kf_vec = x_hat[4:4*n+4]
v_kf_vec = x_hat[4*n+4:]

w_kf = w_kf_vec.reshape(n,4).T
v_kf = v_kf_vec.reshape(n,2).T

x_kf = np.zeros((4,n+1))
x_kf[:,0] = x0_hat

for t in range(n):
    x_kf[:,t+1] = A.dot(x_kf[:,t]) + w_kf[:,t]

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x[0,:])
plt.plot(ts,x_kf[0,:])
plt.xlabel('time')
plt.ylabel('x position')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x[1,:])
plt.plot(ts,x_kf[1,:])
plt.xlabel('time')
plt.ylabel('y position')
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x[2,:])
plt.plot(ts,x_kf[2,:])
plt.xlabel('time')
plt.ylabel('x velocity')
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x[3,:])
plt.plot(ts,x_kf[3,:])
plt.xlabel('time')
plt.ylabel('y velocity')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x[0,:],x[1,:],'-',alpha=0.8, linewidth=5, label='True')
plt.plot(y[0,:],y[1,:],'ko',alpha=0.1, label='Measurement')
plt.plot(x_kf[0,:],x_kf[1,:],'-',alpha=0.8, linewidth=5, label='KF estimates')
plt.title('Trajectory')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.show()

Sigma_w = Q
Sigma_v = R
P = np.eye(4)

x_rk = np.zeros((4,n+1))
w_rk = np.zeros((4,n))
v_rk = np.zeros((2,n))
x_rk[:,0] = np.array([0,0,0,0])

for t in range(n):
  K = A@P@C.T@np.linalg.inv(C@P@C.T+Sigma_v)
  x_rk[:,t+1] = A@x_rk[:,t] + K@(y[:,t] - C@x_rk[:,t])
  P = A@P@A.T - K@C@P@A.T + Sigma_w
  w_rk[:,t] = x_rk[:,t+1] - A@x_rk[:,t]
  v_rk[:,t] = y[:,t] - C@x_rk[:,t]

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x[0,:])
plt.plot(ts,x_kf[0,:])
plt.plot(ts,x_rk[0,:])
plt.xlabel('time')
plt.ylabel('x position')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x[1,:])
plt.plot(ts,x_kf[1,:])
plt.plot(ts,x_rk[1,:])
plt.xlabel('time')
plt.ylabel('y position')
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x[2,:])
plt.plot(ts,x_kf[2,:])
plt.plot(ts,x_rk[2,:])
plt.xlabel('time')
plt.ylabel('x velocity')
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x[3,:])
plt.plot(ts,x_kf[3,:])
plt.plot(ts,x_rk[3,:])
plt.xlabel('time')
plt.ylabel('y velocity')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x[0,:],x[1,:],'-',alpha=0.8, linewidth=5, label='True')
plt.plot(y[0,:],y[1,:],'ko',alpha=0.1, label='Measurement')
plt.plot(x_kf[0,:],x_kf[1,:],'-',alpha=0.8, linewidth=5, label='KF estimates')
plt.plot(x_rk[0,:],x_rk[1,:],'-',alpha=0.6, linewidth=5, label='KF (recursive)')
plt.title('Trajectory')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.show()