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

# initialize GEKKO model
m = GEKKO()

# model discretized time
n = 60*10+1  # Number of second time points (10min)
m.time = np.linspace(0,n-1,n) # Time vector

# Parameters
Qd = np.zeros(601)
Qd[10:200] = 80
Qd[200:400] = 20
Qd[400:] = 50
Q = m.Param(value=Qd) # Percent Heater (0-100%)

T0 = m.Param(value=23.0+273.15)     # Initial temperature
Ta = m.Param(value=23.0+273.15)     # K
U =  m.Param(value=10.0)            # W/m^2-K
mass = m.Param(value=4.0/1000.0)    # kg
Cp = m.Param(value=0.5*1000.0)      # J/kg-K    
A = m.Param(value=12.0/100.0**2)    # Area in m^2
alpha = m.Param(value=0.01)         # W / % heater
eps = m.Param(value=0.9)            # Emissivity
sigma = m.Const(5.67e-8)            # Stefan-Boltzman

T = m.Var(value=T0)         #Temperature state as GEKKO variable

m.Equation(T.dt() == (1.0/(mass*Cp))*(U*A*(Ta-T) \
                    + eps * sigma * A * (Ta**4 - T**4) \
                    + alpha*Q))

# simulation mode
m.options.IMODE = 4

# simulation model
m.solve()

# plot results
plt.figure(1)
plt.subplot(2,1,1)
plt.plot(m.time,Q.value,'b-',label='heater')
plt.ylabel('Heater (%)')
plt.legend(loc='best')
plt.subplot(2,1,2)
plt.plot(m.time,np.array(T.value)-273.15,'r-',label='temperature')
plt.ylabel('Temperature (degC)')
plt.legend(loc='best')
plt.xlabel('Time (sec)')
plt.savefig('tclab_eb_pred.png')
plt.show()