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

m = GEKKO()
m.time = np.linspace(0,10,101)

# Dynamic control options
m.options.IMODE = 6
m.options.CV_TYPE = 1
m.options.MV_TYPE = 0
m.options.SOLVER = 3
m.options.MV_STEP_HOR = 1
m.options.NODES = 3

#Define Manipulated Variables
u = m.MV(name='u')

#Define Controled Variables
y = m.CV(1,name='y')
z = m.CV(1,name='z')
s = m.CV(1,name='s')

# Environmental Constraint
#setup CV
# tau is the speed of the CV response, 0=step, 1 = 63.2# of the way
#   to the new setpoint in 1 sec, only if tr_init is 1 or 2.
# with tr_init=0, it is just a pure dead-band
# specifying the speed to get to the set point
# get to 63.2# of sp withing tau seconds
y.TAU = 5
y.STATUS = 1
y.TR_INIT = 2
y.SPHI = 5
y.SPLO = 4
y.FSTATUS = 0
y.WSPHI = 100
y.WSPLO = 100

# Operational Constraint
z.TAU = 4
z.STATUS = 1
z.TR_INIT = 2
z.SPHI = 7
z.SPLO = 6
z.FSTATUS = 0
z.WSPHI = 50
z.WSPLO = 50

# Safety Constraint
s.TAU = 10
s.STATUS = 1
s.TR_INIT = 2
s.TR_OPEN = 3
s.SPHI = 11
s.SPLO = 10
s.FSTATUS = 0
s.WSPHI = 200
s.WSPLO = 200

#setup MV (u)
u.STATUS = 1
u.DCOST = 0
u.LOWER = 0
u.UPPER = 1000
u.COST = 0

# process model
tau = 1
K = 3
m.Equation(tau*y.dt()+y==u)
m.Equation(z==y)
m.Equation(s==y)

# solve problem
m.solve(disp=True)

# get additional solution information
import json
with open(m.path+'//results.json') as f:
    results = json.load(f)

# create plot
p, ax = plt.subplots(nrows=2, ncols=1, \
                     gridspec_kw={'height_ratios':[3,1]})

ax[0].plot(m.time,results['s.tr_hi'],'r-.',lw=2)
ax[0].plot(m.time,results['y.tr_hi'],'b:',lw=2)
ax[0].plot(m.time,results['z.tr_hi'],'--',color='orange',lw=2)
ax[0].plot(m.time,results['z'],'k-',lw=3)
ax[0].legend(['Priority 1: Safety Constraint',\
            'Priority 2: Environmental Constraint',\
            'Priority 3: Economic Constraint','Response'],loc=4)
ax[0].plot(m.time,results['z.tr_lo'],'--',color='orange',lw=2)
ax[0].plot(m.time,results['y.tr_lo'],'b:',lw=2)
ax[0].plot(m.time,results['s.tr_lo'],'r-.',lw=2)
ax[0].set_ylabel('Pressure (bar)')

ax[1].plot(m.time,u.value,'b-',lw=2)
ax[1].legend(['Manipulated Variable'])
ax[1].set_ylabel('MV')
ax[1].set_xlabel('Time (min)')

plt.show()