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

x_0 = np.array([10, -20, 15, -5])
x_des = np.array([100, 50, 0, 0])

u_max = 1.0

import cvxpy as cp

x = cp.Variable((4,N+1))  # x_{0},...,x_{N}
u = cp.Variable((2,N))    # u_{0},...,u_{N-1}

obj = cp.Minimize(cp.sum_squares(u))

constr = [ x[:,-1] == x_des, 
           x[:,0]  == x_0    ]
for t in range(N):
    constr += [ x[:,t+1] == A@x[:,t] + B@u[:,t] ]
    constr += [ cp.norm(u[:,t]) <= u_max ]
    
prob = cp.Problem(obj, constr)
prob.solve(verbose=True)

x_cp = np.array(x.value)
u_cp = np.array(u.value)


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

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts[:-1],np.linalg.norm(u_cp,2,axis=0))
plt.xlabel('time')
plt.ylabel(r'$||u||$')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x_cp[0,:],x_cp[1,:], label='Minimum energy trajectory')
plt.plot(x_0[0], x_0[1], 'o', markersize=7, label='Initial position')
plt.plot(x_des[0], x_des[1], '*', markersize=10, label='Target position')
plt.title('Trajectory')
plt.legend()
for i in range(0,N-1,10):
  plt.arrow(x_cp[0,i], x_cp[1,i], 10*u_cp[0,i], 10*u_cp[1,i], head_width=0.2, 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()

import cvxpy as cp

p_upper = N
p_lower = 0
epsilon = 1

while (p_upper - p_lower > epsilon):
  
  s = int(0.5*(p_upper+p_lower))
  print (f'UB:{p_upper}    LB:{p_lower}')
  
  tm = np.arange(s+1)*delt

  x = cp.Variable((4,s+1))  # x_{0},...,x_{s}
  u = cp.Variable((2,s))    # u_{0},...,u_{s-1}

  obj = cp.Minimize(cp.sum_squares(u))

  constr = [ x[:,-1] == x_des, 
             x[:,0]  == x_0    ]
  for t in range(s):
      constr += [ x[:,t+1] == A@x[:,t] + B@u[:,t] ]
      constr += [ cp.norm(u[:,t]) <= u_max ]
    
  prob = cp.Problem(obj, constr)
  prob.solve(solver=cp.ECOS)
  
  if prob.status == 'optimal':
    p_upper = s
  else:
    p_lower = s

s = p_upper
tm = np.arange(s+1)*delt

x = cp.Variable((4,s+1))  # x_{0},...,x_{s}
u = cp.Variable((2,s))    # u_{0},...,u_{s-1}

obj = cp.Minimize(cp.sum_squares(u))

constr = [ x[:,-1] == x_des, 
           x[:,0]  == x_0    ]
for t in range(s):
    constr += [ x[:,t+1] == A@x[:,t] + B@u[:,t] ]
    constr += [ cp.norm(u[:,t]) <= u_max ]

    
prob = cp.Problem(obj, constr)
prob.solve()
    
x_mt = np.array(x.value)
u_mt = np.array(u.value)

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,x_cp[0,:], alpha=0.5, label='Minimum energy control')
plt.plot(tm,x_mt[0,:], label='Minimum time control')
plt.xlabel('time')
plt.ylabel('x position')
plt.legend()
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,x_cp[1,:], alpha=0.5)
plt.plot(tm,x_mt[1,:])
plt.xlabel('time')
plt.ylabel('y position')
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,x_cp[2,:], alpha=0.5)
plt.plot(tm,x_mt[2,:])
plt.xlabel('time')
plt.ylabel('x velocity')
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,x_cp[3,:], alpha=0.5)
plt.plot(tm,x_mt[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_cp[0,:], alpha=0.5, label='Minimum energy control')
plt.plot(tm[:-1],u_mt[0,:], label='Minimum time control')
plt.xlabel('time')
plt.ylabel(r'$u_1$')
plt.legend()
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts[:-1],u_cp[1,:], alpha=0.5)
plt.plot(tm[:-1],u_mt[1,:])
plt.xlabel('time')
plt.ylabel(r'$u_2$')
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts[:-1],np.linalg.norm(u_cp,2,axis=0), alpha=0.5, label='Minimum energy control')
plt.plot(tm[:-1],np.linalg.norm(u_mt,2,axis=0), label='Minimum time control')
plt.xlabel('time')
plt.ylabel(r'$||u||$')
plt.legend()
plt.grid()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.plot(x_cp[0,:],x_cp[1,:], alpha=0.5, label='Minimum energy trajectory')
plt.plot(x_mt[0,:],x_mt[1,:], label='Minimum time trajectory')
plt.plot(x_0[0], x_0[1], 'o', markersize=7, label='Initial position')
plt.plot(x_des[0], x_des[1], '*', markersize=10, label='Target position')
plt.title('Trajectory')
plt.legend()
for i in range(0,N-1,10):
  plt.arrow(x_cp[0,i], x_cp[1,i], 10*u_cp[0,i], 10*u_cp[1,i], head_width=0.6, width=0.2, alpha=0.5, fc='tab:green', ec='none')
for i in range(0,s-1,10):
  plt.arrow(x_mt[0,i], x_mt[1,i], 10*u_mt[0,i], 10*u_mt[1,i], head_width=0.6, width=0.2, alpha=1.0, fc='tab:red', ec='none')
plt.axis('equal')
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.grid()
plt.show()