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

#data from problem statement
time = np.array([0, 5, 10, 20, 30, 40, 50, 60, 90, 120,
                 150, 180, 240, 300, 360, 480, 600, 720])
Cc_meas = np.array([0,0.57, 0.78, 0.92, 1.04, 1.19, 1.29,
                    1.36, 1.59, 1.68, 1.84, 1.96, 2.01, 2.13,
                    2.21, 5.32, 2.38, 2.44])

m = GEKKO(remote=True)

m.time = time

Kf = m.FV(value =0.0025, lb = 0, ub = 0.01)
Kb = m.FV(value =0.0025, lb = 0, ub = 0.01)
Kf.STATUS = 1
Kb.STATUS = 1

Ca = m.Var(value = 4.84)
Cb = m.Var(value = 9.67)
Cd = m.Var(value = 0)

Cc = m.CV(Cc_meas)
Cc.FSTATUS = 1

#Reaction equations
m.Equation(Cc.dt() == Kf*Ca*(Cb**2) - Kb*Cc*Cd)
m.Equation(Ca.dt() == -Cc.dt())
m.Equation(Cb.dt() == -2*Cc.dt())
m.Equation(Cd.dt() == Cc.dt())

#turn GEKKO mode to estimation
m.options.IMODE = 5
m.options.NODES = 4
m.options.EV_TYPE = 1 #2 = Sum of squared error, 1 = Sum of absoluted error

m.solve(disp=False)

#solve and plot
plt.figure()
plt.subplot(2,1,1)
plt.plot(m.time, Ca.value, label = 'Ca')
plt.plot(m.time, Cb.value, label = 'Cb')
plt.ylabel('Concentration (mol/L)')
plt.legend(); plt.grid()
plt.subplot(2,1,2)
plt.plot(m.time, Cc.value, label = 'Cc')
plt.plot(m.time, Cc_meas, 'ro', label = 'Cc Data')
plt.ylabel('Concentration (mol/L)')
plt.xlabel('Time (min)')
plt.legend(); plt.grid()
plt.ylabel('Concentration (mol/L)')
plt.tight_layout()
plt.savefig('regression.png',dpi=300)

print("Kf = " + str(Kf.value[0]))
print("Kb = " + str(Kb.value[0]))