import numpy as np
import matplotlib.pyplot as plt

N = 500 
T = 50 
ts = np.linspace(0,T,N+1)
delt = T/N
gamma = .05 # damping

p_0   = np.array([ 10, -20])
v_0   = np.array([ 15,  -5])
p_des = np.array([100,  50])
v_des = np.array([  0,   0])

import cvxpy as cp

###########################
p_lb = np.array([  0, -35])
p_ub = np.array([115,  70])
u_ub = 1.0
###########################

p = cp.Variable((2,N+1))  # p_{0},...,p_{N}
v = cp.Variable((2,N+1))  # v_{0},...,v_{N}
u = cp.Variable((2,N))    # u_{0},...,u_{N-1}

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

constr  = [ p[:,-1] == p_des,  p[:,0] == p_0 ]
constr += [ v[:,-1] == v_des,  v[:,0] == v_0 ]
for t in range(N):
    constr += [ v[:,t+1] == (1-gamma*delt)*v[:,t] + delt*u[:,t] ]
    constr += [ p[:,t+1] == p[:,t] + delt*(1-0.5*gamma*delt)*v[:,t] 
                + 0.5*delt**2*u[:,t] ]
    ################################################
    constr += [ p_lb <= p[:,t+1], p[:,t+1] <= p_ub ]
    constr += [ cp.norm(u[:,t]) <= u_ub ]
    ################################################
    
cp.Problem(obj, constr).solve(verbose=True)

p_cp = np.array(p.value)
v_cp = np.array(p.value)
u_cp = np.array(u.value)

plt.figure(figsize=(14,6), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts,p_cp[0,:])
plt.ylabel(r'$p_x$')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts,p_cp[1,:])
plt.ylabel(r'$p_y$')
plt.grid()
plt.subplot(2,2,3)
plt.plot(ts,v_cp[0,:])
plt.xlabel('time')
plt.ylabel(r'$v_x$')
plt.grid()
plt.subplot(2,2,4)
plt.plot(ts,v_cp[1,:])
plt.xlabel('time')
plt.ylabel(r'$v_y$')
plt.grid()
plt.show()

plt.figure(figsize=(14,6), dpi=100)
plt.subplot(2,2,1)
plt.plot(ts[:-1],u_cp[0,:])
plt.ylabel(r'$u_x$')
plt.grid()
plt.subplot(2,2,2)
plt.plot(ts[:-1],u_cp[1,:])
plt.xlabel('time')
plt.ylabel(r'$u_y$')
plt.grid()
plt.show()
plt.figure(figsize=(14,4), dpi=100)
plt.plot(ts[:-1],np.linalg.norm(u_cp,axis=0))
plt.axhline(u_ub, color=plt.cm.tab20(2), alpha=0.4, label=r'$u_{ub}$')
plt.xlabel('time')
plt.ylabel(r'$||u||$')
plt.grid()
plt.show()
plt.show()

plt.figure(figsize=(14,9), dpi=100)
plt.broken_barh([(p_lb[0], p_ub[0])], (p_lb[1], p_ub[1]-p_lb[1]), \
                alpha = 0.1, label='Feasible region')
plt.plot(p_0[0], p_0[1], 'o', markersize=7, label='Initial position')
plt.plot(p_des[0], p_des[1], '*', markersize=10, label='Target position')
plt.plot(p_cp[0,:],p_cp[1,:], label='Optimal trajectory (constrained)')
plt.title('Trajectory')
plt.legend()
for i in range(0,N-1,10):
  plt.arrow(p_cp[0,i], p_cp[1,i], 10*u_cp[0,i], 10*u_cp[1,i], \
            head_width=1.0, 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()

p_cp_hist = []
v_cp_hist = []
u_cp_hist = []        

p = cp.Variable((2,N+1))  # p_{0},...,p_{N}
v = cp.Variable((2,N+1))  # v_{0},...,v_{N}
u = cp.Variable((2,N))    # u_{0},...,u_{N-1}

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

constr  = [ p[:,-1] == p_des,  p[:,0] == p_0 ]
constr += [ v[:,-1] == v_des,  v[:,0] == v_0 ]
for t in range(N):
    constr += [ v[:,t+1] == (1-gamma*delt)*v[:,t] + delt*u[:,t] ]
    constr += [ p[:,t+1] == p[:,t] + delt*(1-0.5*gamma*delt)*v[:,t] 
                + 0.5*delt**2*u[:,t] ]
    constr += [ p_lb <= p[:,t+1], p[:,t+1] <= p_ub ]
    constr += [ cp.norm(u[:,t]) <= u_ub ]

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

p_cp = np.array(p.value)
v_cp = np.array(v.value)
u_cp = np.array(u.value)

p_cp_hist.append(p_cp)
v_cp_hist.append(v_cp)
u_cp_hist.append(u_cp)

max_iters = 10

##########################
p_obs = np.array([120,20])
d_sft = 20
u_lb = 0.1
##########################

for iter in range(max_iters):

  p0 = p_cp
  v0 = v_cp
  u0 = u_cp

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

  constr  = [ p[:,-1] == p_des,  p[:,0] == p_0 ]
  constr += [ v[:,-1] == v_des,  v[:,0] == v_0 ]
  for t in range(N):
    constr += [ v[:,t+1] == (1-gamma*delt)*v[:,t] + delt*u[:,t] ]
    constr += [ p[:,t+1] == p[:,t] + delt*(1-0.5*gamma*delt)*v[:,t] 
                + 0.5*delt**2*u[:,t] ]
    constr += [ p_lb <= p[:,t+1], p[:,t+1] <= p_ub ]
    constr += [ cp.norm(u[:,t]) <= u_ub ]
    ##################################################################
    constr += [ 2*(p0[:,t] - p_obs).T@p[:,t] >= d_sft**2 
               + cp.sum_squares(p0[:,t]) - p_obs.dot(p_obs) ]
    constr += [ 2*u0[:,t].T@u[:,t] >= u_lb**2 + u0[:,t].dot(u0[:,t]) ]
    ##################################################################

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

  p_cp = p.value
  v_cp = v.value
  u_cp = u.value
       
  p_cp_hist.append(p_cp)
  v_cp_hist.append(v_cp)
  u_cp_hist.append(u_cp)

  print(iter)
  if (np.linalg.norm(p_cp_hist[-1]-p_cp_hist[-2],'fro')<1.0):
    n_iters = len(p_cp_hist)
    print(f'terminated after {iter+1} iterations')
    break


plt.figure(figsize=(14,6), dpi=100)
plt.subplot(2,2,1)
for id in range(n_iters):
  plt.plot(ts,p_cp_hist[id][0,:], color=plt.cm.tab20(0), \
           linewidth=0.5+(id+1)/n_iters*2, alpha = (id+1+2)/(n_iters+2))
plt.xlabel('time')
plt.ylabel(r'$p_x$')
plt.grid()
plt.subplot(2,2,2)
for id in range(n_iters):
  plt.plot(ts,p_cp_hist[id][1,:], color=plt.cm.tab20(0), \
           linewidth=0.5+(id+1)/n_iters*2, alpha = (id+1+2)/(n_iters+2))
plt.xlabel('time')
plt.ylabel(r'$p_y$')
plt.grid()
plt.subplot(2,2,3)
for id in range(n_iters):
  plt.plot(ts,v_cp_hist[id][0,:], color=plt.cm.tab20(0), \
           linewidth=0.5+(id+1)/n_iters*2, alpha = (id+1+2)/(n_iters+2))
plt.xlabel('time')
plt.ylabel(r'$v_x$')
plt.grid()
plt.subplot(2,2,4)
for id in range(n_iters):
  plt.plot(ts,v_cp_hist[id][1,:], color=plt.cm.tab20(0), \
           linewidth=0.5+(id+1)/n_iters*2, alpha = (id+1+2)/(n_iters+2))
plt.xlabel('time')
plt.ylabel(r'$v_y$')
plt.grid()
plt.show()

plt.figure(figsize=(14,6), dpi=100)
plt.subplot(2,2,1)
for id in range(n_iters):
  plt.plot(ts[:-1],u_cp_hist[id][0,:], color=plt.cm.tab20(0), \
           linewidth=0.5+(id+1)/n_iters*2, alpha = (id+1+2)/(n_iters+2))
plt.xlabel('time')
plt.ylabel(r'$u_x$')
plt.grid()
plt.subplot(2,2,2)
for id in range(n_iters):
  plt.plot(ts[:-1],u_cp_hist[id][1,:], color=plt.cm.tab20(0), \
           linewidth=0.5+(id+1)/n_iters*2, alpha = (id+1+2)/(n_iters+2))
plt.xlabel('time')
plt.ylabel(r'$u_y$')
plt.grid()
plt.show()
plt.figure(figsize=(14,4), dpi=100)
for id in range(n_iters):
  plt.plot(ts[:-1],np.linalg.norm(u_cp_hist[id],axis=0), \
           color=plt.cm.tab20(0), linewidth=0.5+(id+1)/n_iters*2, \
           alpha = (id+1+2)/(n_iters+2), label=f'SCP control, iter={id}')
plt.xlabel('time')
plt.ylabel(r'$||u||$')
plt.axhline(u_ub, color=plt.cm.tab20(2), alpha=0.5, label=r'$u_{ub}$')
plt.axhline(u_lb, color=plt.cm.tab20(4), alpha=0.5, label=r'$u_{lb}$')
plt.legend()
plt.grid()
plt.show()

axes = plt.figure(num=1, figsize=(14,9), dpi=100)
plt.broken_barh([(p_lb[0], p_ub[0])], (p_lb[1], p_ub[1]-p_lb[1]), \
                alpha = 0.1, label='Feasible region')
plt.plot(p_0[0], p_0[1], 'o', markersize=7, label='Initial position')
plt.plot(p_des[0], p_des[1], '*', markersize=10, label='Target position')
plt.gca().add_artist(plt.Circle((p_obs[0],p_obs[1]),d_sft, label='Obstable', \
                                color=plt.cm.tab20(8), alpha=0.5))
for id in range(n_iters):
  plt.plot(p_cp_hist[id][0,:], p_cp_hist[id][1,:], \
           color=plt.cm.tab20(4), linewidth=0.5+(id+1)/n_iters*2, \
           alpha = (id+1+2)/(n_iters+2), label=f'SCP trajectory, iter={id}')

plt.legend()
plt.title('Trajectories')

for t in range(0,N,10):
  plt.arrow(p_cp_hist[id][0,t], p_cp_hist[id][1,t], \
            10*u_cp_hist[id][0,t], 10*u_cp_hist[id][1,t], \
            head_width=1.1, width=0.2, color='r', ec='none')
plt.axis('equal')
plt.xlabel(r'$x$ position')
plt.ylabel(r'$y$ position')
plt.grid()
plt.show()