# -*- coding: utf-8 -*-
"""Knowledge_Distillation.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1OAkgOFYmfsD7Um1WiHnvFm_WfJa9XdIW
"""

#Importing necessary Libraries

import tensorflow as tf
import numpy as np
import keras
import random
from keras import backend as K
from keras.models import Sequential
from keras.models import Model
from keras.layers.core import Dense, Dropout, Activation
from keras.layers.embeddings import Embedding
from keras.layers.recurrent import LSTM, GRU
from keras.layers import Input, Embedding, LSTM, Dense, Conv2D, Conv1D, AveragePooling1D, MaxPool2D,GlobalAveragePooling2D,MaxPool1D,GlobalAveragePooling1D,BatchNormalization,Flatten,Bidirectional,Conv3D,MaxPool3D,GlobalAveragePooling3D
from keras.optimizers import Adam
from keras.layers import concatenate
from keras.regularizers import l2
from keras.callbacks import ModelCheckpoint
from keras import optimizers, callbacks
from tensorflow.keras.utils import to_categorical 
import random as rn

#Setting up the SEED
import os
os.environ['PYTHONHASHSEED']='0'
np.random.seed(13817)
rn.seed(12513)
tf.random.set_seed(8130)

scale_factor=4 # defining the scaling factor

def line1(x):
  # received shape: (256*8)

  conv_1=Conv1D(int(128),3,activation='relu',padding='same')(x)
  
  pool_1=MaxPool1D(2)(conv_1)
  print(pool_1.shape)
  reshaped_1=tf.reshape(pool_1,[tf.shape(x)[0],int(128),int(128),1])





  #2D net
  block1_conv1 = Conv2D(int(32/scale_factor),(7,7),activation='relu',padding='same')(reshaped_1)
  block1_pool  = MaxPool2D(2,2)(block1_conv1)

  block2_conv1 = Conv2D(int(64/scale_factor),(5,5),activation='relu',padding='same')(block1_pool)
  block2_pool  = MaxPool2D(2,2)(block2_conv1)
  

  #block3_conv1 = Conv2D(128,(5,5),activation='relu',padding='same')(block1_pool)
  #block3_pool  = MaxPool2D(2,2)(block3_conv1)

  global_pool= GlobalAveragePooling2D()(block2_pool)

  fully_connected_1= Dense(int(512/scale_factor),activation='relu')(global_pool)

  return fully_connected_1

def line2(x):
  # received shape: (64*16)

  conv_1=Conv1D(int(64),14,activation='relu',padding='same')(x)
  
  reshaped_1=tf.reshape(conv_1,[tf.shape(x)[0],int(64),int(64),1])
  print(reshaped_1.shape)
  '''block1_pool=MaxPool2D(4)(conv_1)
  print(block1_pool.shape)'''
  #2D net
  block1_conv1 = Conv2D(int(32/scale_factor),(7,7),activation='relu',padding='same')(reshaped_1)
  block1_pool  = MaxPool2D(2,2)(block1_conv1)

  block2_conv1 = Conv2D(int(64/scale_factor),(5,5),activation='relu',padding='same')(block1_pool)
  block2_pool  = MaxPool2D(2,2)(block2_conv1)

  # block3_conv1 = Conv2D(128,(5,5),activation='relu',padding='same')(block1_pool)
  # block3_pool  = MaxPool2D(2,2)(block3_conv1)

  global_pool= GlobalAveragePooling2D()(block2_pool)

  fully_connected_1= Dense(int(512/scale_factor),activation='relu')(global_pool)

  return fully_connected_1

def line3(x):
  # received shape: (16*32)

  conv_1=Conv1D(int(16),14,activation='relu',padding='same')(x)
  reshaped_1=tf.reshape(conv_1,[tf.shape(x)[0],int(16),int(16),1])
  block1_conv1 = Conv2D(int(16),(3,3),activation='relu',padding='same')(reshaped_1)
  #print(block1_conv1.shape)
  
  reshaped_2=tf.reshape(block1_conv1,[tf.shape(x)[0],int(16),int(16),int(16),1])
  conv_1=Conv3D(int(32/scale_factor),(2,2,2),activation='relu',padding='same')(reshaped_2)
  conv_1=MaxPool3D(2)(conv_1)
  conv_1=Conv3D(int(64/scale_factor),(2,2,2),activation='relu',padding='same')(conv_1)
  conv_1=MaxPool3D(2)(conv_1)
  conv_1=Conv3D(int(128/scale_factor),(2,2,2),activation='relu',padding='same')(conv_1)
  conv_1=MaxPool3D(2)(conv_1)
  conv_1=GlobalAveragePooling3D()(conv_1)
  out=Dense(int(512/scale_factor),activation='relu')(conv_1)
  return out

def line4(x):
  #received shape: (16*256)
  conv_1=Conv1D(int(512/scale_factor),14,activation='relu',padding='same')(x)
  conv_1=MaxPool1D(4)(conv_1)
  conv_1=BatchNormalization()(conv_1)
  lstm_1=LSTM(int(256/scale_factor),return_sequences=False)(conv_1)
  flat=Flatten()(lstm_1)
  out=Dense(int(512/scale_factor),activation='relu')(flat)
  #global_pool=GlobalAveragePooling1D()(conv_1)

  return out

def line5(x):
  x=BatchNormalization()(x)
  lstm1=(LSTM(int(128/scale_factor),return_sequences=False))(x)
  out=Dense(int(512/scale_factor),activation='relu')(lstm1)

  return out

def getModel():
  x=Input(shape=(1024,12),name='x')
  f_lstm=line5(x)

  # converting to (256*64):

  conv_1=Conv1D(int(64/scale_factor),21,activation='relu',padding='same',strides=1)(x)
  conv_1=MaxPool1D(4)(conv_1) # (256*64)
  print("YO1",conv_1.shape)

  f1=line1(conv_1)

  # converting to (64*128):

  conv_1=Conv1D(int(4),21,activation='relu',padding='same',strides=1)(conv_1)
  conv_1=MaxPool1D(4)(conv_1) # (32,64)


  f2=line2(conv_1)

  # converting to (16*256):

  conv_1=Conv1D(int(8),21,activation='relu',padding='same',strides=1)(conv_1)
  conv_1=MaxPool1D(4)(conv_1) # (8*128)
 
  f3=line3(conv_1)


  # converting to (16*256): # same input for line4 function 

  f4=line4(conv_1)
  print(f4.shape)

  







  # concate the 4 dense layer channelwise
 


  f1=reshaped_1=tf.reshape(f1,[tf.shape(x)[0],int(128),1])
  f2=reshaped_1=tf.reshape(f2,[tf.shape(x)[0],int(128),1])
  f3=reshaped_1=tf.reshape(f3,[tf.shape(x)[0],int(128),1])
  f4=reshaped_1=tf.reshape(f4,[tf.shape(x)[0],int(128),1])
  f5=reshaped_1=tf.reshape(f_lstm,[tf.shape(x)[0],int(128),1])


  concat= tf.keras.layers.Concatenate()([f1,f2,f4,f5,f3])
  concat=tf.reduce_mean(concat,axis=2)


  dense_1=Dense(32,activation='relu')(concat)
  out=Dense(6,activation='softmax')(dense_1)
  

  #print(out.shape)
  model=Model(inputs=x,outputs=out)








  return model

model=getModel()
model.summary()

#Constructing Distiller() class
class Distiller(keras.Model):
    def __init__(self, student, teacher):
        super(Distiller, self).__init__()
        self.teacher = teacher
        self.student = student

    def compile(
        self,
        optimizer,
        metrics,
        student_loss_fn,
        distillation_loss_fn,
        alpha=0.1,
        temperature=3,
    ):
        """ Configure the distiller.

        Args:
            optimizer: Keras optimizer for the student weights
            metrics: Keras metrics for evaluation
            student_loss_fn: Loss function of difference between student
                predictions and ground-truth
            distillation_loss_fn: Loss function of difference between soft
                student predictions and soft teacher predictions
            alpha: weight to student_loss_fn and 1-alpha to distillation_loss_fn
            temperature: Temperature for softening probability distributions.
                Larger temperature gives softer distributions.
        """
        super(Distiller, self).compile(optimizer=optimizer, metrics=metrics)
        self.student_loss_fn = student_loss_fn
        self.distillation_loss_fn = distillation_loss_fn
        self.alpha = alpha
        self.temperature = temperature

    def train_step(self, data):
        # Unpack data
        x, y = data
       
        # Forward pass of teacher
        #teacher_predictions = self.teacher(x, training=False)
        teacher_layer_1 = K.function([self.teacher.input], [self.teacher.layers[-1].output,self.teacher.layers[-2].output,self.teacher.layers[-3].output,self.teacher.layers[-4].output,self.teacher.layers[-5].output,self.teacher.layers[-6].output,self.teacher.layers[-7].output,self.teacher.layers[-8].output,self.teacher.layers[-9].output])
        #teacher_layer_1 = K.function([self.teacher.input], [self.teacher.layers[-1].output])
       
        teacher_1 = teacher_layer_1([x])[0]
        teacher_2 = teacher_layer_1([x])[1]
        teacher_3 = teacher_layer_1([x])[2]
        teacher_4 = teacher_layer_1([x])[3]
        teacher_5 = teacher_layer_1([x])[4]
        teacher_6 = teacher_layer_1([x])[5]
        teacher_7 = teacher_layer_1([x])[6]
        teacher_8 = teacher_layer_1([x])[7]
        teacher_9 = teacher_layer_1([x])[8]
        


        with tf.GradientTape() as tape:
            # Forward pass of student
            #student_predictions = self.student(x, training=True)
            student_layer_1 = K.function([self.student.input], [self.student.layers[-1].output,self.student.layers[-2].output,self.student.layers[-3].output,self.student.layers[-4].output,self.student.layers[-5].output,self.student.layers[-6].output,self.student.layers[-7].output,self.student.layers[-8].output,self.student.layers[-9].output])
            #student_layer_1 = K.function([self.student.input], [self.student.layers[-1].output])


            ''',self.student.layers[-2].output,self.student.layers[-3].output,self.student.layers[-4].output]'''
            stu_1 = student_layer_1([x])[0]
            stu_2 = student_layer_1([x])[1]
            stu_3 = student_layer_1([x])[2]
            stu_4 = student_layer_1([x])[3]
            stu_5 = student_layer_1([x])[4]
            stu_6 = student_layer_1([x])[5]
            stu_7 = student_layer_1([x])[6]
            stu_8 = student_layer_1([x])[7]
            stu_9 = student_layer_1([x])[8]
          
            
            # Compute losses
            student_loss = self.student_loss_fn(y, student_predictions)

            '''distillation_loss = self.distillation_loss_fn(
                tf.nn.softmax(teacher_1 / self.temperature, axis=1),
                tf.nn.softmax(stu_1 / self.temperature, axis=1),
            )'''
            dis_loss_1=self.distillation_loss_fn(teacher_1,stu_1)
            dis_loss_2=self.distillation_loss_fn(teacher_2,stu_2)
            dis_loss_3=self.distillation_loss_fn(teacher_3,stu_3)
            dis_loss_4=self.distillation_loss_fn(teacher_4,stu_4)
            dis_loss_5=self.distillation_loss_fn(teacher_5,stu_5)
            dis_loss_6=self.distillation_loss_fn(teacher_6,stu_6)
            dis_loss_7=self.distillation_loss_fn(teacher_7,stu_7)
            dis_loss_8=self.distillation_loss_fn(teacher_8,stu_8)
            dis_loss_9=self.distillation_loss_fn(teacher_9,stu_9)
            distillation_loss=dis_loss_1+dis_loss_2+dis_loss_3+dis_loss_4+dis_loss_5+dis_loss_6+dis_loss_7+dis_loss_8+dis_loss_9
            #loss = self.alpha * student_loss + (1 - self.alpha) *dis_loss_1
            loss = self.alpha * student_loss + (1 - self.alpha) * distillation_loss
        # Compute gradients
        trainable_vars = self.student.trainable_variables
        gradients = tape.gradient(loss, trainable_vars)

        # Update weights
        self.optimizer.apply_gradients(zip(gradients, trainable_vars))

        # Update the metrics configured in `compile()`.
        self.compiled_metrics.update_state(y, student_predictions)

        # Return a dict of performance
        results = {m.name: m.result() for m in self.metrics}
        results.update(
            {"student_loss": student_loss, "distillation_loss": distillation_loss}
        )
        return results

    def test_step(self, data):
        # Unpack the data
        x, y = data

        # Compute predictions
        y_prediction = self.student(x, training=False)

        # Calculate the loss
        student_loss = self.student_loss_fn(y, y_prediction)

        # Update the metrics.
        self.compiled_metrics.update_state(y, y_prediction)

        # Return a dict of performance
        results = {m.name: m.result() for m in self.metrics}
        results.update({"student_loss": student_loss})
        return results

#loading teacher model
from keras.models import model_from_json
json_file = open('/content/drive/MyDrive/teacher_models/model_32.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
teacher = model_from_json(loaded_model_json)
# load weights into new model
teacher.load_weights("/content/drive/MyDrive/teacher_models/model_32.h5")
print("Loaded model from disk")