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

# initialize GEKKO model
m = GEKKO()

# time
m.time = np.linspace(0,20,41)

# constants
mass = 500

# Parameters
b = m.Param(value=50)
K = m.Param(value=0.8)
# Manipulated variable
p = m.MV(value=0, lb=-100, ub=100)

# Reference trajectory
sine = 10*np.sin(m.time/20*4*np.pi)
traj = m.Param(value=sine)

# Controlled Variable
v = m.SV(value=0,name='v')

# Error
e = m.CV(value=0,name='e')

# Equations
m.Equation(mass*v.dt() == -v*b + K*b*p)
m.Equation(e==v-traj)

m.options.IMODE = 6 # control

# MV tuning
p.STATUS = 1 #allow optimizer to change
p.DCOST = 0.1 #smooth out MV
p.DMAX = 50 #slow down change of MV

# CV tuning
e.STATUS = 1 #add the CV to the objective
m.options.CV_TYPE = 1 #Dead-band
db = 2
e.SPHI = db #set point
e.SPLO = -db #set point
e.TR_INIT = 0 #setpoint trajectory
e.TAU = 5 #time constant of setpoint trajectory

# Solve
m.solve()

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

# Plot solution
plt.figure(figsize=(6,4))
plt.subplot(3,1,1)
plt.plot(m.time,p.value,'b-',lw=2,label='MV')
plt.legend(loc='best')
plt.ylabel('MV')
plt.subplot(3,1,2)
plt.plot(m.time,sine+db,'k-',label='SPHI')
plt.plot(m.time,sine-db,'k-',label='SPLO')
plt.plot(m.time,v.value,'r--',lw=2,label='CV')
plt.legend(loc='best')
plt.ylabel('CV')
plt.subplot(3,1,3)
plt.plot(m.time,results['e.tr_hi'],'k-',label='SPHI')
plt.plot(m.time,results['e.tr_lo'],'k-',label='SPLO')
plt.plot(m.time,e.value,'r--',lw=2,label='Error')
plt.legend(loc='best')
plt.ylabel('Error')
plt.xlabel('time')
plt.tight_layout()
plt.show()