#imports
import tensorflow as tf
import time
from tqdm import tqdm
from glob import glob
import cv2
from keras.utils import np_utils
from sklearn.model_selection import train_test_split
import os
from os import listdir
from os.path import isfile, join

from random import sample
import random
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import cv2
import numpy as np

import matplotlib.pyplot as plt
import numpy as np
import os
import PIL
import tensorflow as tf

from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

#imports for capsule neural network in particular
import gc
import os
import nltk
import tqdm
import numpy as np
import pandas as pd
nltk.download("punkt")

from keras.callbacks import Callback, EarlyStopping, ModelCheckpoint
from keras.engine import Layer
from keras.layers import Activation, Add, Bidirectional, Conv1D, Dense, Dropout, Embedding, Flatten
from keras.layers import concatenate, GRU, Input, LSTM, MaxPooling1D
from keras.layers import GlobalAveragePooling1D,  GlobalMaxPooling1D, SpatialDropout1D
from keras.models import Model
from keras.optimizers import Adam
from keras.preprocessing import text, sequence
from sklearn.metrics import accuracy_score, roc_auc_score, log_loss
from sklearn.model_selection import train_test_split
from keras import initializers, regularizers, constraints, optimizers, layers, callbacks
import keras.backend as K
import numpy as np
from tqdm import tqdm
import tensorflow as tf
from datetime import datetime


%load_ext tensorboard
#configuration section
imageSize = 160
imageWidth = 28 #downscale the image to x by x pixels
imageHeight = 28 #original dimensions are 640 x 480
epochNum = 50 #epochs are the number of times we pass through the data
batchSize = 8 #number of samples we work through before updating the model
maxFiles = 1000 #number to load for each category
colorDepth = 1 #3 for RGB
classes = 10 #number of classifications we want to bother with
learningRate = 0.0001 #this is how fast our model learns, lower = more accurate but takes longer
# Parameters Based on Paper
epsilon = 1e-7
m_plus = 0.9
m_minus = 0.1
lambda_ = 0.5
alpha = 0.0005
epochs = 30
no_of_secondary_capsules = 10

optimizer = tf.keras.optimizers.Adam()
params = {
    "no_of_conv_kernels": 16,
    "no_of_primary_capsules": 32,
    "no_of_secondary_capsules": 10,
    "primary_capsule_vector": 8,
    "secondary_capsule_vector": 16,
    "r":3,
}


# checkpoint_path = './logs/model/capsule'

# stamp = datetime.now().strftime("%Y%m%d-%H%M%S")

# logdir = './logs/func/%s' % stamp
# writer = tf.summary.create_file_writer(logdir)

# scalar_logdir = './logs/scalars/%s' % stamp
# file_writer = tf.summary.create_file_writer(scalar_logdir + "/metrics")
# Load the dataset from Kaggle
def get_cv2_image(path, img_size, color_type):
    # Loading as Grayscale image
    if color_type == 1:
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    # Loading as color image
    elif color_type == 3:
        img = cv2.imread(path, cv2.IMREAD_COLOR)
    # Reduce size
    img = cv2.resize(img[:500], img_size) 
    return img

def load_data(img_size , color_type):
    start_time = time.time()
    training_images = []
    training_labels = []

    # Loop over the training folder 
    for class_ in tqdm(range(classes)):
        
        print('Loading directory c{}'.format(class_))
        files = glob(os.path.join('../input/state-farm-distracted-driver-detection/imgs/train', 'c' + str(class_), '*.jpg'))
        
        for file in files:
            img = get_cv2_image(file, img_size , color_type)
            training_images.append(img)
            training_labels.append(class_) 
    
    print("Data Loaded in {} Min".format((time.time() - start_time)/60))
    return training_images, training_labels 



train_X, train_y = load_data( (imageWidth,imageHeight) , colorDepth)
# train_X = train_X[2:]
# train_y = train_y[2:]
train_X = np.array(train_X).reshape(len(train_X),imageWidth,imageHeight,colorDepth)
train_X = tf.cast(train_X, dtype=tf.float32)

train_y = np.array(train_y)
train_X = train_X / 255.0


dataset = tf.data.Dataset.from_tensor_slices((train_X, train_y))

dataset = dataset.shuffle(buffer_size=len(dataset), reshuffle_each_iteration=True)




train_dataset = dataset.take(int(len(dataset)*0.85))
test_dataset = dataset.skip(int(len(dataset)*0.85))

training_dataset_size = len(dataset)
train_size = len(train_dataset)
test_size = len(test_dataset)

dataset = dataset.batch(batch_size=32)
train_dataset = train_dataset.batch(batch_size=32)
test_dataset = test_dataset.batch(batch_size=32)



count = 0
for a, b in train_dataset:
    count += 1
    print(b)
    if(count > 10):
        break




print((training_dataset_size))
print(train_size)



class CapsuleNetwork(tf.keras.Model):
    def __init__(self, no_of_conv_kernels, no_of_primary_capsules, primary_capsule_vector, no_of_secondary_capsules, secondary_capsule_vector, r):
        super(CapsuleNetwork, self).__init__()
        self.no_of_conv_kernels = no_of_conv_kernels
        self.no_of_primary_capsules = no_of_primary_capsules
        self.primary_capsule_vector = primary_capsule_vector
        self.no_of_secondary_capsules = no_of_secondary_capsules
        self.secondary_capsule_vector = secondary_capsule_vector
        self.r = r
        
        
        with tf.name_scope("Variables") as scope:
            self.convolution = tf.keras.layers.Conv2D(self.no_of_conv_kernels, [9,9], strides=[1,1], name='ConvolutionLayer', activation='relu')
            self.primary_capsule = tf.keras.layers.Conv2D(self.no_of_primary_capsules * self.primary_capsule_vector, [9,9], strides=[2,2], name="PrimaryCapsule")
            self.w = tf.Variable(tf.random_normal_initializer()(shape=[1, 1152, self.no_of_secondary_capsules, self.secondary_capsule_vector, self.primary_capsule_vector]), dtype=tf.float32, name="PoseEstimation", trainable=True)
            self.dense_1 = tf.keras.layers.Dense(units = 256, activation='relu')
            self.dense_2 = tf.keras.layers.Dense(units = 512, activation='relu')
            self.dense_3 = tf.keras.layers.Dense(units = 784, activation='sigmoid', dtype='float32')
        
    def build(self, input_shape):
        pass
        
    def squash(self, s):
        with tf.name_scope("SquashFunction") as scope:
            s_norm = tf.norm(s, axis=-1, keepdims=True)
            return tf.square(s_norm)/(1 + tf.square(s_norm)) * s/(s_norm + epsilon)
    
    @tf.function
    def call(self, inputs):
        input_x, y = inputs
        # input_x.shape: (None, 28, 28, 1)
        # y.shape: (None, 10)
        
        x = self.convolution(input_x) # x.shape: (None, 20, 20, 256)
        x = self.primary_capsule(x) # x.shape: (None, 6, 6, 256)
        
        with tf.name_scope("CapsuleFormation") as scope:
            u = tf.reshape(x, (-1, self.no_of_primary_capsules * x.shape[1] * x.shape[2], 8)) # u.shape: (None, 1152, 8)
            u = tf.expand_dims(u, axis=-2) # u.shape: (None, 1152, 1, 8)
            u = tf.expand_dims(u, axis=-1) # u.shape: (None, 1152, 1, 8, 1)
            u_hat = tf.matmul(self.w, u) # u_hat.shape: (None, 1152, 10, 16, 1)
            u_hat = tf.squeeze(u_hat, [4]) # u_hat.shape: (None, 1152, 10, 16)

        
        with tf.name_scope("DynamicRouting") as scope:
            b = tf.zeros((input_x.shape[0], 1152, self.no_of_secondary_capsules, 1)) # b.shape: (None, 1152, 10, 1)
            for i in range(self.r): # self.r = 3
                c = tf.nn.softmax(b, axis=-2) # c.shape: (None, 1152, 10, 1)
                s = tf.reduce_sum(tf.multiply(c, u_hat), axis=1, keepdims=True) # s.shape: (None, 1, 10, 16)
                v = self.squash(s) # v.shape: (None, 1, 10, 16)
                agreement = tf.squeeze(tf.matmul(tf.expand_dims(u_hat, axis=-1), tf.expand_dims(v, axis=-1), transpose_a=True), [4]) # agreement.shape: (None, 1152, 10, 1)
                # Before matmul following intermediate shapes are present, they are not assigned to a variable but just for understanding the code.
                # u_hat.shape (Intermediate shape) : (None, 1152, 10, 16, 1)
                # v.shape (Intermediate shape): (None, 1, 10, 16, 1)
                # Since the first parameter of matmul is to be transposed its shape becomes:(None, 1152, 10, 1, 16)
                # Now matmul is performed in the last two dimensions, and others are broadcasted
                # Before squeezing we have an intermediate shape of (None, 1152, 10, 1, 1)
                b += agreement
                
        with tf.name_scope("Masking") as scope:
            y = tf.expand_dims(y, axis=-1) # y.shape: (None, 10, 1)
            y = tf.expand_dims(y, axis=1) # y.shape: (None, 1, 10, 1)
            mask = tf.cast(y, dtype=tf.float32) # mask.shape: (None, 1, 10, 1)
            v_masked = tf.multiply(mask, v) # v_masked.shape: (None, 1, 10, 16)
            
        with tf.name_scope("Reconstruction") as scope:
            v_ = tf.reshape(v_masked, [-1, self.no_of_secondary_capsules * self.secondary_capsule_vector]) # v_.shape: (None, 160)
            reconstructed_image = self.dense_1(v_) # reconstructed_image.shape: (None, 512)
            reconstructed_image = self.dense_2(reconstructed_image) # reconstructed_image.shape: (None, 1024)
            reconstructed_image = self.dense_3(reconstructed_image) # reconstructed_image.shape: (None, 784)
        
        return v, reconstructed_image

    @tf.function
    def predict_capsule_output(self, inputs):
        x = self.convolution(inputs) # x.shape: (None, 20, 20, 256)
        x = self.primary_capsule(x) # x.shape: (None, 6, 6, 256)
        
        with tf.name_scope("CapsuleFormation") as scope:
            u = tf.reshape(x, (-1, self.no_of_primary_capsules * x.shape[1] * x.shape[2], 8)) # u.shape: (None, 1152, 8)
            u = tf.expand_dims(u, axis=-2) # u.shape: (None, 1152, 1, 8)
            u = tf.expand_dims(u, axis=-1) # u.shape: (None, 1152, 1, 8, 1)
            u_hat = tf.matmul(self.w, u) # u_hat.shape: (None, 1152, 10, 16, 1)
            u_hat = tf.squeeze(u_hat, [4]) # u_hat.shape: (None, 1152, 10, 16)

        
        with tf.name_scope("DynamicRouting") as scope:
            b = tf.zeros((inputs.shape[0], 1152, self.no_of_secondary_capsules, 1)) # b.shape: (None, 1152, 10, 1)
            for i in range(self.r): # self.r = 3
                c = tf.nn.softmax(b, axis=-2) # c.shape: (None, 1152, 10, 1)
                s = tf.reduce_sum(tf.multiply(c, u_hat), axis=1, keepdims=True) # s.shape: (None, 1, 10, 16)
                v = self.squash(s) # v.shape: (None, 1, 10, 16)
                agreement = tf.squeeze(tf.matmul(tf.expand_dims(u_hat, axis=-1), tf.expand_dims(v, axis=-1), transpose_a=True), [4]) # agreement.shape: (None, 1152, 10, 1)
                # Before matmul following intermediate shapes are present, they are not assigned to a variable but just for understanding the code.
                # u_hat.shape (Intermediate shape) : (None, 1152, 10, 16, 1)
                # v.shape (Intermediate shape): (None, 1, 10, 16, 1)
                # Since the first parameter of matmul is to be transposed its shape becomes:(None, 1152, 10, 1, 16)
                # Now matmul is performed in the last two dimensions, and others are broadcasted
                # Before squeezing we have an intermediate shape of (None, 1152, 10, 1, 1)
                b += agreement
        return v

    @tf.function
    def regenerate_image(self, inputs):
        with tf.name_scope("Reconstruction") as scope:
            v_ = tf.reshape(inputs, [-1, self.no_of_secondary_capsules * self.secondary_capsule_vector]) # v_.shape: (None, 160)
            reconstructed_image = self.dense_1(v_) # reconstructed_image.shape: (None, 512)
            reconstructed_image = self.dense_2(reconstructed_image) # reconstructed_image.shape: (None, 1024)
            reconstructed_image = self.dense_3(reconstructed_image) # reconstructed_image.shape: (None, 784)
        return reconstructed_image
tf.summary.trace_on(graph=True, profiler=True)
keras.backend.clear_session()
model = CapsuleNetwork(**params)


def safe_norm(v, axis=-1, epsilon=1e-7):
    v_ = tf.reduce_sum(tf.square(v), axis = axis, keepdims=True)
    return tf.sqrt(v_ + epsilon)

def loss_function(v, reconstructed_image, y, y_image):
    prediction = safe_norm(v)
    prediction = tf.reshape(prediction, [-1, no_of_secondary_capsules])
    
    left_margin = tf.square(tf.maximum(0.0, m_plus - prediction))
    right_margin = tf.square(tf.maximum(0.0, prediction - m_minus))
    
    l = tf.add(y * left_margin, lambda_ * (1.0 - y) * right_margin)
    
    margin_loss = tf.reduce_mean(tf.reduce_sum(l, axis=-1))
    
    y_image_flat = tf.reshape(y_image, [-1, 784])
    reconstruction_loss = tf.reduce_mean(tf.square(y_image_flat - reconstructed_image))
    
    loss = tf.add(margin_loss, alpha * reconstruction_loss)
    
    return loss
def train(x,y):
    y_one_hot = tf.one_hot(y, depth=10)
    with tf.GradientTape() as tape:
        v, reconstructed_image = model([x, y_one_hot])
        loss = loss_function(v, reconstructed_image, y_one_hot, x)
    grad = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grad, model.trainable_variables))
    return loss

def get_loss(x,y):
    y_one_hot = tf.one_hot(y, depth=10)
    with tf.GradientTape() as tape:
        v, reconstructed_image = model([x, y_one_hot])
        loss = loss_function(v, reconstructed_image, y_one_hot, x)
    return loss


print(train_X.shape)
print(train_y.shape)

_ = train(train_X[:1],train_y[:1])
# with writer.as_default():
#     tf.summary.trace_export(name="my_func_trace", step=0, profiler_outdir=logdir)

# tf.summary.trace_off()
model.save_weights('my_model_weights.h5')
model.summary()
def predict(model, x):
    pred = safe_norm(model.predict_capsule_output(x))
    pred = tf.squeeze(pred, [1])
    return np.argmax(pred, axis=1)[:,0]

# checkpoint = tf.train.Checkpoint(model=model)
# %tensorboard --logdir ./logs

count = 10
train_dataset = train_dataset.shuffle(64)
for x,y in train_dataset:
    print(predict(model, x))
    #print(x)
    print(y)
    count -= 1
    if count < 0:
        break
def reset_weights(model):
    session = tensorflow.keras.backend.get_session()
    for layer in model.layers: 
        if hasattr(layer, 'kernel_initializer'):
            layer.kernel.initializer.run(session=session)

    
losses = []
accuracy = []
test_acc = []

test_losses = []
#keras.backend.clear_session()
#model.load_weights('my_model_weights.h5')
K.get_session().close()
K.set_session(tf.Session())
K.get_session().run(tf.global_variables_initializer())
for i in range(1, epochs+1, 1):

    loss = 0
    with tqdm(total=len(train_dataset)) as pbar:
        
        description = "Epoch " + str(i) + "/" + str(epochs)
        pbar.set_description_str(description)

        for X_batch, y_batch in train_dataset:

            loss += train(X_batch,y_batch)
            pbar.update(1)

        loss /= len(train_dataset)
        losses.append(loss.numpy())
        
        training_sum = 0

        print_statement = "Loss :" + str(loss.numpy()) + " Evaluating Accuracy ..."
        pbar.set_postfix_str(print_statement)

        for X_batch, y_batch in train_dataset:
            training_sum += sum(predict(model, X_batch)==y_batch.numpy())
        accuracy.append(training_sum/train_size)
        
        training_sum = 0
        test_loss = 0
        
        for X_batch, y_batch in test_dataset:
            training_sum += sum(predict(model, X_batch)==y_batch.numpy())
            test_loss += get_loss(X_batch, y_batch)
        
        test_loss /= len(test_dataset)
        test_losses.append(test_loss)
        test_acc.append(training_sum/test_size)
            

#         with file_writer.as_default():
#             tf.summary.scalar('Loss', data=loss.numpy(), step=i)
#             tf.summary.scalar('Accuracy', data=accuracy[-1], step=i)
#             tf.summary.scalar('Test Acc', data = test_acc[-1], step=i)
        
        print_statement = "Loss :" + str(loss.numpy()) + " Accuracy :" + str(accuracy[-1])  + " Test Acc :" + str(test_acc[-1])

#         if i % 10 == 0:
#             print_statement += ' Checkpoint Saved'
#             checkpoint.save(checkpoint_path)
        
        pbar.set_postfix_str(print_statement)
import matplotlib.pyplot as plt
import numpy as np

plt.plot(losses, label = "train")
plt.plot(test_losses, label = "test")
plt.title('Loss over Time')
plt.xlabel('Loss')
plt.ylabel('Epoch')
plt.legend()
plt.show()
import matplotlib.pyplot as plt
import numpy as np

plt.plot(accuracy, label = "train")
plt.plot(test_acc, label = "test")
plt.title('Accuracy over Time')
plt.xlabel('Accuracy')
plt.ylabel('Epoch')
plt.legend()
plt.show()
"""
OLD TRAINING, IGNORE PAST THIS POINT
"""
losses = []
accuracy = []
for i in range(1, epochs+1, 1):

    loss = 0
    with tqdm(total=len(dataset)) as pbar:
        
        description = "Epoch " + str(i) + "/" + str(epochs)
        pbar.set_description_str(description)

        for X_batch, y_batch in dataset:

            loss += train(X_batch,y_batch)
            pbar.update(1)

        loss /= len(dataset)
        losses.append(loss.numpy())
        
        training_sum = 0

        print_statement = "Loss :" + str(loss.numpy()) + " Evaluating Accuracy ..."
        pbar.set_postfix_str(print_statement)

        for X_batch, y_batch in dataset:
            training_sum += sum(predict(model, X_batch)==y_batch.numpy())
        accuracy.append(training_sum/training_dataset_size)

        with file_writer.as_default():
            tf.summary.scalar('Loss', data=loss.numpy(), step=i)
            tf.summary.scalar('Accuracy', data=accuracy[-1], step=i)
        
        print_statement = "Loss :" + str(loss.numpy()) + " Accuracy :" + str(accuracy[-1])

#         if i % 10 == 0:
#             print_statement += ' Checkpoint Saved'
#             checkpoint.save(checkpoint_path)
        
        pbar.set_postfix_str(print_statement)
# Load the dataset from Kaggle
def get_cv2_image(path, img_size, color_type):
    # Loading as Grayscale image
    if color_type == 1:
        img = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
    # Loading as color image
    elif color_type == 3:
        img = cv2.imread(path, cv2.IMREAD_COLOR)
    # Reduce size
    img = cv2.resize(img[:500], img_size) 
    return img

def load_data(img_size , color_type, training):
    start_time = time.time()
    training_images = []
    training_labels = []

    # Loop over the training folder 
    for class_ in tqdm(range(classes)):
        
        print('Loading directory c{}'.format(class_))
        if(training):
            files = glob(os.path.join('../input/state-farm-distracted-driver-detection/imgs/train', 'c' + str(class_), '*.jpg'))
        else:
            files = glob(os.path.join('../input/state-farm-distracted-driver-detection/imgs/test', 'c' + str(class_), '*.jpg'))
        
        for file in files:
            img = get_cv2_image(file, img_size , color_type)
            training_images.append(img)
            training_labels.append(class_) 
    
    print("Data Loaded in {} Min".format((time.time() - start_time)/60))
    return training_images, training_labels 

train_X, train_y = load_data( (imageWidth,imageHeight) , colorDepth, True)
test_X, test_y = load_data( (imageWidth,imageHeight) , colorDepth, False)

train_y = np_utils.to_categorical(train_y, classes)
test_y = np_utils.to_categorical(test_y, classes)

#now we prep our data
npLabels = np_utils.to_categorical(labels, classes) #we're casting the list of ints as a list of class labels
trainingImages, validationImages, trainingLabels, validationLabels = train_test_split(images, npLabels, test_size=0.15,shuffle=True)

#transform our images into array that can be fed into the machine learning model
trainingImages = np.array(trainingImages).reshape(len(trainingImages),imageWidth,imageHeight,colorDepth)
validationImages = np.array(validationImages).reshape(len(validationImages),imageWidth,imageHeight,colorDepth)

print(type(trainingImages))
resNet  = tf.keras.applications.resnet.ResNet50(include_top = False,weights = 'imagenet',input_shape = (imageWidth,imageHeight,colorDepth))
resNet.summary()
#resnet50 has some random output, we need to add a final layer that fits the output of our problem
#preppedOutput = resNet.output #set this layer to the output layer of initial model
#preppedOutput = tf.keras.layers.Flatten()(preppedOutput) #this takes our many dimensional output, and converts it to a longer 1 dimensional output

#add a dense fully connected output with a number of nodes equal to our classes
#ourOutput =tf.keras.layers.Dense(classes,activation = tf.nn.softmax)(preppedOutput)
#finalModel = tf.keras.models.Model(inputs=resNet.inputs, outputs=ourOutput)

#let's use a simple bread and butter model model and compare the two!
simpleModel = Sequential([
  layers.experimental.preprocessing.Rescaling(1./255, input_shape=(100, 100, colorDepth)),
  layers.Conv2D(16, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Conv2D(32 * 8, [9,9], strides=[2,2], name="PrimaryCapsule"),
  layers.Conv2D(32, 3, padding='same', activation='relu'),
  layers.MaxPooling2D(),
  layers.Conv2D(64, 3, padding='same', activation='relu'),
  layers.Dropout(0.5),
  layers.MaxPooling2D(),
  layers.Flatten(),
  layers.Dense(128, activation='relu'),
  layers.Dense(classes)
])



#compile the model with the learning rate specified above, and standard loss config from resnet example
#finalModel.compile(optimizer=tf.keras.optimizers.Adam(learningRate),
#              loss=tf.keras.losses.CategoricalCrossentropy(from_logits = False),
#              metrics=['accuracy'])

simpleModel.compile(optimizer=tf.keras.optimizers.Adam(learningRate),
              loss=tf.keras.losses.CategoricalCrossentropy(from_logits = False),
              metrics=['accuracy'])

#print model details
#finalModel.summary()
simpleModel.summary()
model2 = Sequential()
model2.add(Dense(500, input_dim=2, activation='relu'))

model2.add(Dense(1, activation='sigmoid'))
model2.compile(loss='binary_crossentropy', optimizer='adam', metrics=['accuracy'])
model2.summary()
#train resnet model
trainingResNet = finalModel.fit(
      x = train_X,y=train_y,#training data
      validation_data=(test_X,test_y), #validation data
      batch_size = batchSize,
      epochs=epochNum,
      verbose=1)
#train our simple model
trainingSimple = simpleModel.fit(
      x = train_X,y=train_y,#training data
      validation_data=(test_X,test_y), #validation data
      batch_size = batchSize,
      epochs=epochNum,
      verbose=1)