from gekko import GEKKO
import numpy as np
import matplotlib.pyplot as plt

m = GEKKO(remote=False)

n = 101; Lp = 4; b = 3; a = 1

m.time = np.linspace(0,1,n)
x1,x2,x3 = m.Array(m.Var,3,lb=0,ub=10)
x1.value = a
u = m.MV(value=0, lb=-10, ub=20) # integer=True

z = np.zeros(n); z[-1] = 1
final = m.Param(value=z)

m.Equation(x1.dt() == u)
m.Equation(x2.dt() == x1 * (1+u**2)**0.5)
m.Equation(x3.dt() == (1+u**2)**0.5)

m.Minimize(1000*(x1-b)**2*final)
m.Minimize(1000*(x3-Lp)**2*final)

m.options.SOLVER = 3
m.options.NODES = 3
m.options.IMODE = 6
m.Minimize(x2*final)

# initialize
u.STATUS = 0
m.solve()

# optimize
m.options.TIME_SHIFT = 0
u.STATUS = 1; u.DCOST = 1e-3
m.solve()

plt.figure(figsize=(6,4))
plt.subplot(2,1,1)
plt.plot(m.time,x1.value,'r--',label=r'$x_1$')
plt.plot(m.time,x2.value,'g-',label=r'$x_2$')
plt.plot(m.time,x3.value,'k:',label=r'$x_3$')
plt.ylabel('x'); plt.xlim([0,1])
plt.legend(); plt.grid()
plt.subplot(2,1,2)
plt.plot(2,1,2)
plt.plot(m.time,u.value,'k-',label=r'$u$')
plt.grid(); plt.xlim([0,1]); plt.legend()
plt.xlabel('Distance'); plt.ylabel('u')
plt.tight_layout()
plt.savefig('chain.png',dpi=300); plt.show()