# importing required libraries
from tensorflow import keras
import os 
from pathlib import Path

# downloading the ImageNet model with output layer removed and input shape (224, 224, 3)
base_model = keras.applications.VGG16(
    weights = 'imagenet',
    input_shape = (224, 224, 3),
    include_top = False)
 
base_model.summary()

# Freeze base model
base_model.trainable = False

# Create inputs with correct shape
inputs = keras.Input(shape=(224, 224, 3))

x = base_model(inputs, training=False)

# Add pooling layer or flatten layer
x = keras.layers.GlobalAveragePooling2D()(x)

# Add final dense layer
outputs = keras.layers.Dense(6, activation = 'softmax')(x)

# Combine inputs and outputs to create model
model = keras.Model(inputs, outputs)

model.summary()

# compile the model
model.compile(loss='categorical_crossentropy', metrics=['accuracy'])

# Augment the Data 
datagen = keras.preprocessing.image.ImageDataGenerator(
        samplewise_center = True,  # set each sample mean to 0
        rotation_range = 10,  # randomly rotate images in the range (degrees, 0 to 180)
        zoom_range = 0.1, # Randomly zoom image 
        width_shift_range = 0.1,  # randomly shift images horizontally (fraction of total width)
        height_shift_range = 0.1,  # randomly shift images vertically (fraction of total height)
        horizontal_flip = True,  # randomly flip images
        vertical_flip = True)

# load and iterate training dataset
dirPath = os.path.join(Path(os.path.realpath("__file__")).parent.parent, "assets/data")
train_it = datagen.flow_from_directory(dirPath + "/train", 
                                       target_size = (224, 224), 
                                       color_mode = 'rgb', 
                                       class_mode = "categorical",
                                       batch_size = 8)
# load and iterate validation dataset
valid_it = datagen.flow_from_directory(dirPath + "/valid", 
                                       target_size = (224, 224), 
                                       color_mode = 'rgb', 
                                       class_mode = "categorical",
                                       batch_size = 8)

# Train the Model
model.fit(train_it,
          validation_data = valid_it,
          steps_per_epoch = train_it.samples/train_it.batch_size,
          validation_steps = valid_it.samples/valid_it.batch_size,
          epochs = 15)

# # Unfreeze the base model
# base_model.trainable = True

# # Compile the model with a low learning rate
# model.compile(optimizer=keras.optimizers.RMSprop(learning_rate = 0.0001),
#               loss='categorical_crossentropy', metrics=['accuracy'])

# model.fit(train_it,
#           validation_data = valid_it,
#           steps_per_epoch = train_it.samples/train_it.batch_size,
#           validation_steps = valid_it.samples/valid_it.batch_size,
#           epochs = 8)

# saving the trained model
dirPath = os.path.join(Path(os.path.realpath("__file__")).parent.parent, "assets/model")
if not os.path.exists(dirPath): os.makedirs(dirPath)
    
model.save(dirPath + "/hand_gesture_classify_model.h5")