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

m = GEKKO(remote=False)

nt = 501
m.time = np.linspace(0,1,nt)

# Variables
x1 = m.Var(value=1)
x2 = m.Var(value=0)
x3 = m.Var(value=0)
T = m.MV(value=398,lb=298,ub=398)
T.STATUS = 1
T.DCOST = 1e-6
k1 = m.Intermediate(4000*m.exp(-2500/T))
k2 = m.Intermediate(6.2e5*m.exp(-5000/T))

p = np.zeros(nt)
p[-1] = 1.0
final = m.Param(value=p)

# Intermediates
r1 = m.Intermediate(k1*x1**2)
r2 = m.Intermediate(k2*x2)

# Equations
m.Equation(x1.dt()==-r1)
m.Equation(x2.dt()== r1 - r2)
m.Equation(x3.dt()== r2)

# Objective Function
# maximize final x2
m.Maximize(x2*final)

m.options.IMODE = 6
m.options.NODES = 3
m.solve()

print(f'Optimal x2(tf): {x2.value[-1]:0.4f}')

plt.figure(figsize=(6,4))
plt.subplot(2,1,1)
plt.plot(m.time,x1.value,'k:',\
         lw=2,label=r'$x_1$')
plt.plot(m.time,x2.value,'b-',\
         lw=2,label=r'$x_2$')
plt.plot(m.time,x3.value,'r--',\
         lw=2,label=r'$x_3$')
plt.plot(m.time[-1],x2.value[-1],'o',\
         color='orange',markersize=5,\
         label=r'$x_2\left(t_f\right)$')
plt.legend(loc='best'); plt.grid()
plt.ylabel('Mole Fraction')
plt.subplot(2,1,2)
plt.plot(m.time,T.value,'r-',\
         lw=2,label='Temperature')
plt.ylabel(r'$T/;(^oC)$'); plt.legend(loc='best')
plt.xlabel('Time'); plt.ylabel('Temperature (K)')
plt.grid(); plt.tight_layout()
plt.savefig('batch_results.png',dpi=300)
plt.show()