import numpy as np
import matplotlib.pyplot as plt

n = 1000 # number of timesteps
T = 50 # 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))
B = np.zeros((4,2))
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

B[0,0] = delt**2/2
B[1,1] = delt**2/2
B[2,0] = delt
B[3,1] = 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([1,1])
R = np.diag([4,4])
Qh = sqrtm(Q)
Rh = sqrtm(R)

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

for t in range(n):
    x[:,t+1] = A.dot(x[:,t]) + B.dot(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,:], label='True')
plt.plot(ts[:-1],y[0,:], 'k.', alpha=0.1, label='Measurement')
plt.xlabel('time')
plt.ylabel('x position')
plt.legend()
plt.ylim([-80,10])
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x[1,:])
plt.plot(ts[:-1],y[1,:], 'k.', alpha=0.1)
plt.xlabel('time')
plt.ylabel('y position')
plt.ylim([-20,40])
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,6), dpi=100)
plt.subplot(2,2,1)
plt.hist(v_true[0,:], bins=10)
plt.xlabel(r'$v_x$')
plt.ylabel('frequency')
plt.title('Gaussian noise')
plt.grid()
plt.subplot(2,2,2)
plt.hist(v_true[1,:], bins=10)
plt.xlabel(r'$v_y$')
plt.ylabel('frequency')
plt.title('Gaussian noise')
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 with Gaussian measurements')
plt.legend()
plt.grid()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.xlim([-80,10])
plt.ylim([-20,40])
plt.show()

import cvxpy as cp

x = cp.Variable(shape=(4, n+1))
w = cp.Variable(shape=(2, n))
v = cp.Variable(shape=(2, n))

tau = 0.16
    
obj = cp.sum_squares(w) + tau*cp.sum_squares(v)
obj = cp.Minimize(obj)

constr = []
for t in range(n):
    constr += [ x[:,t+1] == A@x[:,t] + B@w[:,t] ,
                y[:,t]   == C@x[:,t] + v[:,t]   ]

cp.Problem(obj, constr).solve(verbose=True)

x_cvx = np.array(x.value)
w_cvx = np.array(w.value)
v_cvx = np.array(v.value)

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x_true[0,:], label='True')
plt.plot(ts,x_cvx[0,:], label='Kalman estimates')
plt.xlabel('time')
plt.ylabel('x position')
plt.legend()
plt.ylim([-80,10])
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x_true[1,:])
plt.plot(ts,x_cvx[1,:])
plt.xlabel('time')
plt.ylabel('y position')
plt.ylim([-20,40])
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x_true[2,:], label='True')
plt.plot(ts,x_cvx[2,:], label='Kalman estimates')
plt.xlabel('time')
plt.ylabel('x velocity')
plt.legend()
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x_true[3,:])
plt.plot(ts,x_cvx[3,:])
plt.xlabel('time')
plt.ylabel('y velocity')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x_true[0,:],x_true[1,:],'-',alpha=0.8, linewidth=5, label='True')
plt.plot(y[0,:],y[1,:],'ko',alpha=0.1, label='Measurement')
plt.plot(x_cvx[0,:],x_cvx[1,:],'-',alpha=0.8, linewidth=5, label='KF estimates')
plt.title('Trajectory estimates (KF) with Gaussian measurements')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.xlim([-80,10])
plt.ylim([-20,40])
plt.show()

np.random.seed(7030)

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

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

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

for t in range(n):
  if np.random.rand() < 0.25:
    v[:,t] *= 30

for t in range(n):
    x[:,t+1] = A.dot(x[:,t]) + B.dot(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_true[0,:], label='True')
plt.plot(ts[:-1],y[0,:], 'k.', alpha=0.1, label='Measurement')
plt.xlabel('time')
plt.ylabel('x position')
plt.legend()
plt.ylim([-80,10])
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x_true[1,:])
plt.plot(ts[:-1],y[1,:], 'k.', alpha=0.1)
plt.xlabel('time')
plt.ylabel('y position')
plt.ylim([-20,40])
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x_true[2,:])
plt.xlabel('time')
plt.ylabel('x velocity')
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x_true[3,:])
plt.xlabel('time')
plt.ylabel('y velocity')
plt.grid()
plt.show()

plt.figure(figsize=(14,6), dpi=100)
plt.subplot(2,2,1)
plt.hist(v_true[0,:], bins=10)
plt.xlabel(r'$v_x$')
plt.ylabel('frequency')
plt.title('Non-Gaussian noise')
plt.grid()
plt.subplot(2,2,2)
plt.hist(v_true[1,:], bins=10)
plt.xlabel(r'$v_y$')
plt.ylabel('frequency')
plt.title('Non-Gaussian noise')
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 with non-Gaussian measurements')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.xlim([-80,10])
plt.ylim([-20,40])
plt.show()

import cvxpy as cp

x = cp.Variable(shape=(4, n+1))
w = cp.Variable(shape=(2, n))
v = cp.Variable(shape=(2, n))

tau = 0.16
    
obj = cp.sum_squares(w) + tau*cp.sum_squares(v)
obj = cp.Minimize(obj)

constr = []
for t in range(n):
    constr += [ x[:,t+1] == A@x[:,t] + B@w[:,t] ,
                y[:,t]   == C@x[:,t] + v[:,t]   ]

cp.Problem(obj, constr).solve(verbose=True)

x_kf = np.array(x.value)
w_kf = np.array(w.value)
v_kf = np.array(v.value)

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x_true[0,:], label='True')
plt.plot(ts,x_kf[0,:], label='Kalman estimates')
plt.plot(ts[:-1],y[0,:], 'k.', alpha=0.1, label='Measurement')
plt.xlabel('time')
plt.ylabel('x position')
plt.legend()
plt.ylim([-80,10])
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x_true[1,:])
plt.plot(ts,x_kf[1,:])
plt.plot(ts[:-1],y[1,:], 'k.', alpha=0.1)
plt.xlabel('time')
plt.ylabel('y position')
plt.ylim([-20,40])
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x_true[2,:], label='True')
plt.plot(ts,x_kf[2,:], label='Kalman estimates')
plt.xlabel('time')
plt.ylabel('x velocity')
plt.legend()
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x_true[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_true[0,:],x_true[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 estimates (KF) with non-Gaussian measurements')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.xlim([-80,10])
plt.ylim([-20,40])
plt.show()

from scipy.special import huber
x_test = np.linspace(-5,5,500)
plt.figure(figsize=(12,6), dpi=100)
plt.plot(x_test,x_test**2, label='Quadratic')
plt.plot(x_test,2*huber(1,x_test), linewidth=4, alpha=0.5, label='Huber')
plt.axis('equal')
plt.legend()
plt.ylim(0,4)
plt.grid()
plt.show()

import cvxpy as cp

x = cp.Variable(shape=(4, n+1)) 
w = cp.Variable(shape=(2, n))
v = cp.Variable(shape=(2, n))

tau = 0.16
rho = 2.0
    
obj = cp.sum_squares(w)
obj += sum(tau*cp.huber(cp.norm(v[:,t]),rho) for t in range(n))
obj = cp.Minimize(obj)

constr = []
for t in range(n):
    constr += [ x[:,t+1] == A@x[:,t] + B@w[:,t] ,
                y[:,t]   == C@x[:,t] + v[:,t]   ]

cp.Problem(obj, constr).solve(verbose=True)

x_rkf = np.array(x.value)
w_rkf = np.array(w.value)
v_rkf = np.array(v.value)

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x_true[0,:], label='True')
plt.plot(ts,x_kf[0,:], label='Kalman estimates')
plt.plot(ts,x_rkf[0,:], label='Robust Kalman estimates')
plt.plot(ts[:-1],y[0,:], 'k.', alpha=0.1, label='Measurement')
plt.xlabel('time')
plt.ylabel('x position')
plt.legend()
plt.ylim([-80,10])
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x_true[1,:])
plt.plot(ts,x_kf[1,:])
plt.plot(ts,x_rkf[1,:])
plt.plot(ts[:-1],y[1,:], 'k.', alpha=0.1)
plt.xlabel('time')
plt.ylabel('y position')
plt.ylim([-20,40])
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x_true[2,:], label='True') 
plt.plot(ts,x_kf[2,:], label='Kalman estimates')
plt.plot(ts,x_rkf[2,:], label='Robust Kalman estimates')
plt.xlabel('time')
plt.ylabel('x velocity')
plt.legend()
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x_true[3,:])
plt.plot(ts,x_kf[3,:])
plt.plot(ts,x_rkf[3,:])
plt.xlabel('time')
plt.ylabel('y velocity')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x_true[0,:],x_true[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.5, linewidth=5, label='Kalman estimates')
plt.plot(x_rkf[0,:],x_rkf[1,:],'-',alpha=0.8, linewidth=5, label='Robust Kalman estimates')
plt.title('Trajectory estimates (KF and Huber) with non-Gaussian measurements')
plt.grid()
plt.legend()
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.axis('equal')
plt.xlim([-80,10])
plt.ylim([-20,40])
plt.show()