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

m = GEKKO(remote=False)
m.time = np.linspace(0, 1, 100)

# variables
Z1 = m.Var(value=0, ub=330, lb=0)
Z2 = m.Var(value=0, ub=33, lb=0)
m.fix_final(Z2, 0)
m.fix_final(Z1, 300)

# adjust final time
tf = m.FV(value=500, lb=0.1)
tf.STATUS = 1

# discrete decision - manipulated variable
u = m.MV(value=1, lb=-2, ub=1, integer=True)
u.STATUS = 1

m.Equation(Z1.dt() / tf == Z2)
m.Equation(Z2.dt() / tf == u)

m.Minimize(tf)

# options
m.options.IMODE = 6
m.options.SOLVER = 1

m.solve(disp=False)

# display solution
print("Total time: " + str(tf.value[0]))
tm = m.time * tf.value[0]
plt.figure(figsize=(6,4))
plt.subplot(211)
plt.plot(tm, Z1, label=r'$Z_1$')
plt.plot(tm, Z2, label=r'$Z_2$')
plt.ylabel('Z'); plt.legend(); plt.grid()
plt.subplot(212)
plt.plot(tm, u, label=r'$u$')
plt.ylabel('u'); plt.xlabel('Time')
plt.legend(); plt.grid(); plt.tight_layout()
plt.savefig('optimal_time.png',dpi=300); plt.show()