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 = .05 # damping, 0 is no damping

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

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


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

C = np.array([[1,0,0,0],[0,1,0,0]])

x_0 = np.array([10, -20, 30, -10])

K = 3
tk = np.array([ 300, 600, 1000])-1
wp = np.array([[100,  50,  100],
               [  0,  50,   50]])


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

S = np.zeros( (2*K,2*n))
for k in range(K):
  S[2*k:2*k+2,2*tk[k]:2*tk[k]+2] = np.eye(2)
  
u_hat = sla.lsqr(S@H,wp.T.flatten() - S@G@x_0)[0]

u_vec = u_hat

u_opt = u_vec.reshape(1000,2).T

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

for t in range(n):
    x[:,t+1] = A.dot(x[:,t]) + B.dot(u_opt[:,t])

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x[0,:])
plt.xlabel('time')
plt.ylabel('x position')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x[1,:])
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.subplot(2,2,1)
plt.plot(ts[:-1],u_opt[0,:])
plt.xlabel('time')
plt.ylabel(r'$u_1$')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts[:-1],u_opt[1,:])
plt.xlabel('time')
plt.ylabel(r'$u_2$')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x_0[0], x_0[1], 'o', markersize=10, label='Initial position')
for k in range(K):
  plt.plot(wp[0,k], wp[1,k], '^', markersize=10, label=f'Waypoint #{k+1}')
plt.title('Trajectory')
plt.plot(x[0,:],x[1,:], label='Optimal trajectory')
plt.legend()
for i in range(0,n-1,10):
  plt.arrow(x[0,i], x[1,i], 2*u_opt[0,i], 2*u_opt[1,i], head_width=1, width=0.2, fc='tab:red', ec='none')
plt.axis('equal')
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.grid()
plt.show()

W = 1*ssp.diags([1, 1, 1, 1, 10, 10])
A_tilde = ssp.vstack((ssp.eye(2*n), W@S@H))
y_tilde = np.hstack((np.zeros(2*n), W@(wp.T.flatten() - S@G@x_0))) 

u_hat_rlx = sla.lsqr(A_tilde, y_tilde)[0]

u_vec_rlx = u_hat_rlx

u_opt_rlx = u_vec_rlx.reshape(1000,2).T

x_rlx = np.zeros((4,n+1))
x_rlx[:,0] = x_0

for t in range(n):
    x_rlx[:,t+1] = A.dot(x_rlx[:,t]) + B.dot(u_opt_rlx[:,t])

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x_0[0], x_0[1], 'o', markersize=10, label='Initial position')
for k in range(K):
  plt.plot(wp[0,k], wp[1,k], '^', markersize=10, label=f'Waypoint #{k+1}')
plt.plot(x[0,:],x[1,:], label='Optimal trajectory')
plt.plot(x_rlx[0,:],x_rlx[1,:], label='Optimal trajectory (relaxed)')
plt.title('Trajectory')
plt.legend()
plt.axis('equal')
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.grid()
plt.show()