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

m1 = GEKKO(remote=False)
m2 = GEKKO(remote=False)

m = m1
# known parameters
# number of biweeks in a year
nb = 26
ny = 3 # number of years

biweeks = np.zeros((nb,ny*nb+1))
biweeks[0][0] = 1
for i in range(nb):
   for j in range(ny):
      biweeks[i][j*nb+i+1] = 1

# write csv data file
tm = np.linspace(0,78,79)
# case data
cases = np.array([180,180,271,423,465,523,649,624,556,420,\
                  423,488,441,268,260,163,83,60,41,48,65,82,\
                  145,122,194,237,318,450,671,1387,1617,2058,\
                  3099,3340,2965,1873,1641,1122,884,591,427,282,\
                  174,127,84,97,68,88,79,58,85,75,121,174,209,458,\
                  742,929,1027,1411,1885,2110,1764,2001,2154,1843,\
                  1427,970,726,416,218,160,160,188,224,298,436,482,468])
data = np.vstack((tm,cases))
data = data.T
np.savetxt('measles_biweek_2.csv',data,delimiter=',',header='time,cases')

#load data from csv
m.time, cases_meas = np.loadtxt('measles_biweek_2.csv', \
                delimiter=',',skiprows=1,unpack=True)

m.Vr = m.Param(value = 0)
# Variables
m.N = m.FV(value = 3.2e6)
m.mu = m.FV(value = 7.8e-4)

m.rep_frac = m.FV(value = 0.45)
# beta values (unknown parameters in the model)
m.beta = [m.FV(value=1, lb=0.1, ub=5) for i in range(nb)]
# predicted values
m.S = m.SV(value = 0.06*m.N.value, lb=0,ub=m.N)
m.I = m.SV(value = 0.001*m.N.value, lb=0,ub=m.N)

m.V = m.Var(value = 2e5)

# measured values
m.cases = m.CV(value = cases_meas, lb=0)
# turn on feedback status for CASES
m.cases.FSTATUS = 1
# weight on prior model predictions
m.cases.WMODEL = 0
# meas_gap = deadband that represents level of
#   accuracy / measurement noise
db = 100
m.cases.MEAS_GAP = db

for i in range(nb):
   m.beta[i].STATUS=1

m.gamma = m.FV(value=0.07)
m.gamma.STATUS = 1
m.gamma.LOWER = 0.05
m.gamma.UPPER = 0.5

m.biweek=[None]*nb
for i in range(nb):
   m.biweek[i] = m.Param(value=biweeks[i])

# Intermediate
m.Rs = m.Intermediate(m.S*m.I/m.N)

# Equations
sum_biweek = sum([m.biweek[i]*m.beta[i]*m.Rs for i in range(nb)])
m.Equation(m.S.dt()== -sum_biweek + m.mu*m.N - m.Vr)
m.Equation(m.I.dt()== sum_biweek - m.gamma*m.I)
m.Equation(m.cases == m.rep_frac*sum_biweek)
m.Equation(m.V.dt()==-m.Vr)

# options
m.options.SOLVER = 1
m.options.NODES=3
# imode = 5, dynamic estimation
m.options.IMODE = 5
# ev_type = 1 (L1-norm) or 2 (squared error)
m.options.EV_TYPE = 1

# solve model and print solver output
m.solve()

[print('beta['+str(i+1)+'] = '+str(m.beta[i][0])) \
 for i in range(nb)]
print('gamma = '+str(m.gamma.value[0]))

# export data
# stack time and avg as column vectors
my_data = np.vstack((m.time,np.asarray(m.beta),m.gamma))
# transpose data
my_data = my_data.T
# save text file with comma delimiter
beta_str = ''
for i in range(nb):
   beta_str = beta_str + ',beta[' + str(i+1) + ']'
header_name = 'time,gamma' + beta_str
##np.savetxt('solution_data.csv',my_data,delimiter=',',\
##           header = header_name, comments='')
np.savetxt('solution_data_EVTYPE_'+str(m.options.EV_TYPE)+\
           '_gamma'+str(m.gamma.STATUS)+'.csv',\
           my_data,delimiter=',',header = header_name)

plt.figure(num=1, figsize=(16,8))
plt.suptitle('Estimation')
plt.subplot(2,2,1)
plt.plot(m.time,m.cases, label='Cases (model)')
plt.plot(m.time,cases_meas, label='Cases (measured)')
if m.options.EV_TYPE==1:
   plt.plot(m.time,cases_meas+db/2, 'k-.',\
            lw=0.5, label=r'$Cases_{db-hi}$')
   plt.plot(m.time,cases_meas-db/2, 'k-.',\
            lw=0.5, label=r'$Cases_{db-lo}$')
   plt.fill_between(m.time,cases_meas-db/2,\
                    cases_meas+db/2,color='gold',alpha=.5)
plt.legend(loc='best')
plt.ylabel('Cases')
plt.subplot(2,2,2)
plt.plot(m.time,m.S,'r--')
plt.ylabel('S')
plt.subplot(2,2,3)
[plt.plot(m.time,m.beta[i], label='_nolegend_')\
 for i in range(nb)]
plt.plot(m.time,m.gamma,'c--', label=r'$\gamma$')
plt.legend(loc='best')
plt.ylabel(r'$\beta, \gamma$')
plt.xlabel('Time')
plt.subplot(2,2,4)
plt.plot(m.time,m.I,'g--')
plt.xlabel('Time')
plt.ylabel('I')
plt.subplots_adjust(hspace=0.2,wspace=0.4)

name = 'cases_EVTYPE_'+ str(m.options.EV_TYPE) +\
       '_gamma' + str(m.gamma.STATUS) +  '.png'
plt.savefig(name)

##-----------------------------------------------##
## Control
##-----------------------------------------------##

m = m2

m.time = m1.time

# Variables
N = m.FV(value = 3.2e6)
mu = m.FV(value = 7.8e-4)

rep_frac = m.FV(value = 0.45)
# beta values (unknown parameters in the model)
beta = [m.FV(value = m1.beta[i].NEWVAL) for i in range(nb)]
gamma = m.FV(value = m1.gamma.NEWVAL)

cases = m.CV(value = cases_meas[0],lb=0)

# predicted values
S = m.SV(value=0.06*N, lb=0,ub=N)
I = m.SV(value = 0.001*N, lb=0,ub=N)

V = m.CV(value = 2e5)
Vr = m.MV(value = 0)

cases.STATUS = 1
cases.FSTATUS = 0
cases.TR_INIT = 0
cases.SPHI = 50
cases.SPLO = 0
cases_SPHI = np.full(len(m.time),cases.SPHI)
cases_SPLO = np.full(len(m.time),cases.SPLO)

Vr.STATUS = 1
Vr.UPPER = 1e4
Vr.LOWER = 0
Vr.COST = 1e-5

V.SPHI = 2e5
V.SPLO = 0
V.STATUS = 0
V.TR_INIT = 0
V_SPHI = np.full(len(m.time),V.SPHI)
V_SPLO = np.full(len(m.time),V.SPLO)

biweek=[None]*nb
for i in range(nb):
   biweek[i] = m.Param(value=biweeks[i])

# Intermediates
Rs = m.Intermediate(S*I/N)

#Equations
sum_biweek = sum([biweek[i]*beta[i]*Rs for i in range(nb)])
m.Equation(S.dt()== -sum_biweek + mu*N - Vr)
m.Equation(I.dt()== sum_biweek - gamma*I)
m.Equation(cases == rep_frac*sum_biweek)
m.Equation(V.dt() == -Vr)

# options
m.options.SOLVER = 1
# imode = 6, dynamic control
m.options.IMODE = 6
# ctrl_units = 5, time units are in years
m.options.CTRL_UNITS = 5
m.options.CV_TYPE = 1

# solve model and print solver output
m.solve()

[print('beta['+str(i+1)+'] = '+str(beta[i][0]))\
 for i in range(nb)]
print('gamma = '+str(gamma.value[0]))

# export data
# stack time and avg as column vectors
my_data = np.vstack((m.time,np.asarray(beta),V,Vr,gamma))
# transpose data
my_data = my_data.T
# save text file with comma delimiter
beta_str = ''
for i in range(nb):
   beta_str = beta_str + ',beta[' + str(i+1) + ']'
header_name = 'time,gamma' + beta_str
##np.savetxt('solution_data.csv',my_data,delimiter=',',\
##           header = header_name, comments='')
np.savetxt('solution_control_EVTYPE_'+str(m.options.EV_TYPE)+\
           '_gamma'+str(gamma.STATUS)+'.csv',\
           my_data,delimiter=',',header = header_name)

plt.figure(num=2, figsize=(16,8))
plt.suptitle('Control')
plt.subplot2grid((6,2),(0,0), rowspan=2)
plt.plot(m.time,cases, label='Cases (model)')
plt.plot(m.time,cases_SPHI, 'k-.', lw=0.5,\
         label=r'$Cases _{SP-HI}$')
plt.plot(m.time,cases_SPLO, 'k-.', lw=0.5,\
         label=r'$Cases _{SP-LO}$')
plt.fill_between(m.time,cases_SPLO,cases_SPHI,\
                 color='gold',alpha=.25)
plt.legend(loc='best')
plt.ylabel('Cases')

plt.subplot2grid((6,2),(0,1), rowspan=3)
plt.plot(m.time,S,'r--')
plt.ylabel('S')

plt.subplot2grid((6,2),(2,0), rowspan=2)
[plt.plot(m.time,V, label='_nolegend_') for i in range(nb)]
plt.plot(m.time,V_SPHI, 'k-.', lw=0.5,\
         label=r'$V _{SP-HI}$')
plt.plot(m.time,V_SPLO, 'k-.', lw=0.5,\
         label=r'$V _{SP-LO}$')
plt.fill_between(m.time,V_SPLO,V_SPHI,color='gold',alpha=.25)
plt.legend(loc='best')
plt.ylabel(r'$V$')
plt.xlabel('Time')

plt.subplot2grid((6,2),(3,1),rowspan=3)
plt.plot(m.time,I,'g--')
plt.xlabel('Time')
plt.ylabel('I')

plt.subplot2grid((6,2),(4,0),rowspan=2)
plt.plot(m.time,Vr, 'b--')
plt.ylabel(r'$V_{r}$')
plt.xlabel('Time')
plt.tight_layout()
plt.subplots_adjust(top=0.95,wspace=0.2)
name = 'cases_EVTYPE_'+ str(m.options.EV_TYPE) + '_gamma'+\
       str(gamma.STATUS) +  '.png'
plt.savefig(name)

plt.show()