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

m = GEKKO(remote=True)

# Add 0.01 as first step
m.time = np.insert(np.linspace(0,50,101),1,0.01)

# change solver options
m.solver_options = ['minlp_gap_tol 0.001',\
                    'minlp_maximum_iterations 10000',\
                    'minlp_max_iter_with_int_sol 100',\
                    'minlp_branch_method 1',\
                    'minlp_integer_tol 0.001',\
                    'minlp_integer_leaves 0',\
                    'minlp_maximum_iterations 200']

SP = 3

last = m.Param(np.zeros(102))
last.value[-1] = 1

sigma=m.MV(value=1,lb=1,ub=2,integer=True)
sigma.STATUS = 1; sigma.DCOST=1e-5
x1 = m.Var(value=2)
x2 = m.CV(value=2); x2.STATUS=1
x2.SPLO = SP-0.2; x2.SPHI = SP+0.2

m.Equations([x1.dt() == sigma - m.sqrt(x1),\
             x2.dt() == m.sqrt(x1) - m.sqrt(x2)])

m.options.IMODE = 6
m.options.NODES = 2
m.options.SOLVER = 1
m.options.MV_TYPE = 0
m.solve()

plt.figure(1)
plt.step(m.time,sigma.value,'r-',label=r'$\sigma$ (1/2)')
plt.plot(m.time,x1.value,'k-',label=r'$x_1$')
plt.plot([0,m.time[-1]],[x2.SPHI,x2.SPHI],'b:',label=r'$SP_{HI}$')
plt.plot(m.time,x2.value,'g-',label=r'$x_2$')
plt.plot([0,m.time[-1]],[x2.SPLO,x2.SPLO],'b:',label=r'$SP_{LO}$')
plt.xlabel('Time'); plt.ylabel('Levels')
plt.grid(); plt.legend(); plt.tight_layout()
plt.savefig('double_tank.png',dpi=300)
plt.show()