import numpy as np
import pandas as pd
import matplotlib.pyplot as plt  
from gekko import GEKKO
import time

# -------------------------------------
# import or generate data
# -------------------------------------
filename = 'tclab_ss_data1.csv'
try:
    try:
        data = pd.read_csv(filename)
    except:
        url = 'https://apmonitor.com/do/uploads/Main/tclab_ss_data1.txt'
        data = pd.read_csv(url)
except:
    # generate training data if data file not available
    import tclab
    # Connect to Arduino
    a = tclab.TCLab()
    fid = open(filename,'w')
    fid.write('Heater,Temperature\n')
    # test takes 2 hours = 40 pts * 3 minutes each
    npts = 40
    Q = np.sin(np.linspace(0,np.pi,npts))*100
    for i in range(npts):
        # set heater value
        a.Q1(Q[i])
        print('Heater 1: ' + str(Q[i]) + ' %')
        # wait 3 minutes
        time.sleep(3*60)
        # record temperature and heater value
        print('Temperature 1: ' + str(a.T1) + ' degC')
        fid.write(str(Q[i])+','+str(a.T1)+'\n')
    # close file
    fid.close()
    # close connection to Arduino
    a.close()
    # read data file
    data = pd.read_csv(filename)

# -------------------------------------
# scale data
# -------------------------------------
x = data['Heater'].values
y = data['Temperature'].values
# minimum of x,y
x_min = min(x)
y_min = min(y)
# range of x,y
x_range = max(x)-min(x)
y_range = max(y)-min(y)
# scaled data
xs = (x - x_min)/x_range
ys = (y - y_min)/y_range

# -------------------------------------
# build neural network
# -------------------------------------
nin = 1  # inputs
n1 = 1   # hidden layer 1 (linear)
n2 = 1   # hidden layer 2 (nonlinear)
n3 = 1   # hidden layer 3 (linear)
nout = 1 # outputs

# Initialize gekko models
train = GEKKO() 
test  = GEKKO()
dyn   = GEKKO()
model = [train,test,dyn]

for m in model:
    # input(s)
    m.inpt = m.Param()

    # layer 1
    m.w1 = m.Array(m.FV, (nin,n1))
    m.l1 = [m.Intermediate(m.w1[0,i]*m.inpt) for i in range(n1)]

    # layer 2
    m.w2 = m.Array(m.FV, (n1,n2))
    m.l2 = [m.Intermediate(sum([m.tanh(m.w2[j,i]*m.l1[j]) \
            for j in range(n1)])) for i in range(n2)]

    # layer 3
    m.w3 = m.Array(m.FV, (n2,n3))
    m.l3 = [m.Intermediate(sum([m.w3[j,i]*m.l2[j] \
            for j in range(n2)])) for i in range(n3)]

    # output(s)
    m.outpt = m.CV()
    m.Equation(m.outpt==sum([m.l3[i] for i in range(n3)]))

    # flatten matrices
    m.w1 = m.w1.flatten()
    m.w2 = m.w2.flatten()
    m.w3 = m.w3.flatten()

# -------------------------------------
# fit parameter weights
# -------------------------------------
m = train
m.inpt.value=xs
m.outpt.value=ys
m.outpt.FSTATUS = 1
for i in range(len(m.w1)):
    m.w1[i].FSTATUS=1
    m.w1[i].STATUS=1
    m.w1[i].MEAS=1.0
for i in range(len(m.w2)):
    m.w2[i].STATUS=1
    m.w2[i].FSTATUS=1
    m.w2[i].MEAS=0.5
for i in range(len(m.w3)):
    m.w3[i].FSTATUS=1
    m.w3[i].STATUS=1
    m.w3[i].MEAS=1.0
m.options.IMODE = 2
m.options.SOLVER = 3
m.options.EV_TYPE = 2
m.solve(disp=False)

# -------------------------------------
# test sample points
# -------------------------------------
m = test
for i in range(len(m.w1)):
    m.w1[i].MEAS=train.w1[i].NEWVAL
    m.w1[i].FSTATUS = 1
    print('w1['+str(i)+']: '+str(m.w1[i].MEAS))
for i in range(len(m.w2)):
    m.w2[i].MEAS=train.w2[i].NEWVAL
    m.w2[i].FSTATUS = 1
    print('w2['+str(i)+']: '+str(m.w2[i].MEAS))
for i in range(len(m.w3)):
    m.w3[i].MEAS=train.w3[i].NEWVAL
    m.w3[i].FSTATUS = 1
    print('w3['+str(i)+']: '+str(m.w3[i].MEAS))
m.inpt.value=np.linspace(-0.1,1.5,100)
m.options.IMODE = 2
m.options.SOLVER = 3
m.solve(disp=False)

# -------------------------------------
# un-scale predictions
# -------------------------------------
xp = np.array(test.inpt.value) * x_range + x_min
yp = np.array(test.outpt.value) * y_range + y_min

# -------------------------------------
# plot results
# -------------------------------------
plt.figure()
plt.plot(x,y,'bo',label='data')
plt.plot(xp,yp,'r-',label='predict')
plt.legend(loc='best')
plt.ylabel('y')
plt.xlabel('x')
plt.savefig('tclab_ss_data1.png')

# -------------------------------------
# generate dynamic predictions
# -------------------------------------
m = dyn
m.time = np.linspace(0,600,601)
# load neural network parameters
for i in range(len(m.w1)):
    m.w1[i].MEAS=train.w1[i].NEWVAL
    m.w1[i].FSTATUS = 1
for i in range(len(m.w2)):
    m.w2[i].MEAS=train.w2[i].NEWVAL
    m.w2[i].FSTATUS = 1
for i in range(len(m.w3)):
    m.w3[i].MEAS=train.w3[i].NEWVAL
    m.w3[i].FSTATUS = 1
# doublet test
Qd = np.zeros(601)
Qd[10:200] = 80
Qd[200:400] = 20
Qd[400:] = 50
Q = m.Param()
Q.value = Qd
# scaled input
m.inpt.value = (Qd - x_min) / x_range

# define Temperature output
Q0 = 0   # initial heater
T0 = 23  # initial temperature
# scaled steady state ouput
T_ss = m.Var(value=T0)
m.Equation(T_ss == m.outpt*y_range + y_min)
# dynamic prediction
T = m.Var(value=T0)
# time constant
tau = m.Param(value=120) # determine in a later exercise
# additional model equation for dynamics
m.Equation(tau*T.dt()==-(T-T0)+(T_ss-T0))

# solve dynamic simulation
m.options.IMODE=4
m.solve()

plt.figure()
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,T.value,'r-',label='temperature')
#plt.plot(m.time,T_ss.value,'k--',label='target temperature')
plt.ylabel('Temperature (degC)')
plt.legend(loc='best')
plt.xlabel('Time (sec)')
plt.savefig('tclab_dyn_pred.png')
plt.show()