import tclab
import time
import numpy as np
import matplotlib.pyplot as plt
import pickle
from keras.models import load_model

# PID: 1 sec, MPC: 2 sec 
cycle_time = 2
tclab_hardware = False # switch to True if hardware available

h5_file,s_x,s_y,window = pickle.load(open('lstm_control.pkl','rb'))
model = load_model(h5_file)

if tclab_hardware:
    mlab = tclab.TCLab  # Physical hardware
else:
    mlab = tclab.setup(connected=False,speedup=10) # Emulator

# LSTM Controller
def lstm(T1_m,Tsp1_m,T2_m,Tsp2_m):
    # Calculate error (necessary feature for LSTM input)
    err1 = Tsp1_m - T1_m
    err2 = Tsp2_m - T2_m

    # Format data for LSTM input
    X = np.vstack((Tsp1_m,err1,Tsp2_m,err2)).T
    Xs = s_x.transform(X)
    Xs = np.reshape(Xs, (1, Xs.shape[0], Xs.shape[1]))

    # Predict Q for controller and unscale
    Qc_s = model.predict(Xs)
    Qc = s_y.inverse_transform(Qc_s)[0]

    # Ensure Qc is between 0 and 100
    Qc = np.clip(Qc,0.0,100.0)
    return Qc

tf = 300                          # final time (sec)
n = int(tf/cycle_time)            # cycles
tm=[]; T1=[]; T2=[]; Q1=[]; Q2=[] # storage

lab = mlab()
i = 0
for t in tclab.clock(tf, cycle_time):
    # record time
    tm.append(t)
    # read temperatures
    T1.append(lab.T1)
    T2.append(lab.T2)
    # LSTM control
    if i>=window:
        T1_m = T1[i-window:i]
        T2_m = T2[i-window:i]
    else:
        insert1 = np.ones(window-i)*T1[0]
        insert2 = np.ones(window-i)*T2[0]
        T1_m = np.concatenate((insert1,T1[0:-1]))
        T2_m = np.concatenate((insert2,T2[0:-1]))
    Tsp1_m = 50*np.ones(window)
    Tsp2_m = 40*np.ones(window)
    Q = lstm(T1_m,Tsp1_m,T2_m,Tsp2_m)
    Q1.append(Q[0]); Q2.append(Q[1])
    lab.Q1(Q1[-1]); lab.Q2(Q2[-1])
    if i%50==0:            
        print('  Time_____Q1___Tsp1_____T1_____Q2___Tsp2_____T2')
    if i%5==0:
        s='{:6.1f} {:6.2f} {:6.2f} {:6.2f} {:6.2f} {:6.2f} {:6.2f}'
        print((s).format(tm[-1],Q1[-1],Tsp1_m[-1],T1[-1],\
                                Q2[-1],Tsp2_m[-1],T2[-1]))
    i+=1
lab.close()

# Create Figure
tmin = np.array(tm)/60.0
plt.figure(figsize=(10,6))
plt.subplot(3,1,1)
plt.plot(tmin,np.ones_like(tmin)*50,'k-',label='$T_{SP,1} (^oC)$')
plt.plot(tmin,T1,'b-',label='$T_1 (^oC)$')
plt.legend(loc=1)
plt.subplot(3,1,2)
plt.plot(tmin,np.ones_like(tmin)*40,'k--',label='$T_{SP,2} (^oC)$')
plt.plot(tmin,T2,'r-',label='$T_2 (^oC)$')
plt.legend(loc=1)
plt.subplot(3,1,3)
plt.plot(tmin,Q1,'b-',label='$Q1_{LSTM}$ (%)')
plt.plot(tmin,Q2,'r--',label='$Q2_{LSTM}$ (%)')
plt.legend(loc=1)
plt.ylim([0,100])
plt.xlabel('Time (min)')
plt.savefig('LSTM_Control.png',dpi=300)
plt.show()