#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Dec 28 18:18:18 2020

@author: marcin
"""

import os
import tensorflow as tf
from keras.preprocessing.image import ImageDataGenerator
import pandas as pd
import numpy as np
from tensorflow.keras.layers import Dense, Flatten, Conv2D, Dropout
from tensorflow.keras.models import Sequential
import matplotlib.pyplot as plt
from keras.preprocessing.image import ImageDataGenerator
import json

gpu_options = tf.compat.v1.GPUOptions(allow_growth=True)
sess = tf.compat.v1.Session(config=tf.compat.v1.ConfigProto(gpu_options=gpu_options))


def smooth_curve(points, factor=0.8):
    smoothed_points = []
    for point in points:
        if smoothed_points:
            previous = smoothed_points[-1]
            smoothed_points.append(previous * factor + point * (1 - factor))
        else:
            smoothed_points.append(point)
    return smoothed_points


conv_base = tf.keras.applications.VGG16(weights='imagenet',
                                        include_top=False,
                                        input_shape=(150, 150, 3))

# 150 x 150 pixels image generator
trainDataGen = ImageDataGenerator(
    rescale=1./255,
    rotation_range=40,
    width_shift_range=0.2,
    height_shift_range=0.2,
    shear_range=0.2,
    zoom_range=0.2,
    horizontal_flip=True,
    fill_mode='nearest')

testDataGen = ImageDataGenerator(rescale=1./255)

baseDir = "/home/marcin/Pictures/dogs_vs_cats/cats_and_dogs_small"
trainDir = os.path.join(baseDir, 'train')
validationDir = os.path.join(baseDir, 'validation')
testDir = os.path.join(baseDir, 'test')

trainGenerator = trainDataGen.flow_from_directory(trainDir,
                                                  target_size=(150, 150),
                                                  batch_size=20,
                                                  class_mode='binary')

validationGenerator = testDataGen.flow_from_directory(validationDir,
                                                  target_size=(150, 150),
                                                  batch_size=20,
                                                  class_mode='binary')

# freezing conv_base   ESSENTIAL!
conv_base.trainable = False
print("weights:", len(conv_base.weights))
print("trainable_weights:", len(conv_base.trainable_weights))
print("non_trainable_weights:", len(conv_base.non_trainable_weights))

# building densed model
model = Sequential()
model.add(conv_base)  # added conv_base 
model.add(Flatten())
model.add(Dense(256, activation='relu'))
model.add(Dense(1, activation='sigmoid'))


model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=2e-5),
              loss='binary_crossentropy',
              metrics=['acc'])

# creating callbacks folder
checkpoint_path = "/home/marcin/Documents/TensorflowProjects/callbacks_dogs_vs_cats"
checkpoint_dir = os.path.dirname(checkpoint_path)

# creating callback
cp_callback_1 = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)

cp_callback_2 = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_path,
                                                 save_weights_only=True,
                                                 verbose=1)

history = model.fit_generator(trainGenerator,
                             steps_per_epoch=100,
                             epochs=30,
                             validation_data=validationGenerator,
                             validation_steps=50,
                             callbacks=[cp_callback_1])



acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(acc) + 1)
plt.plot(epochs, acc, 'bo', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy')
plt.legend()
plt.figure()
plt.plot(epochs, loss, 'bo', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show()

# unfreezing conv_base
conv_base.trainable = True

# setting specific layers as trainable
set_trainable = True

for layer in conv_base.layers:
    if layer.name =='block5_conv1':
        set_trainable = True
    if set_trainable:
        layer.trainable = True
    else:
        layer.trainable = False
        
model.compile(loss='binary_crossentropy',
              optimizer=tf.keras.optimizers.RMSprop(lr=1e-5),
              metrics=['acc'])

history = model.fit_generator(
    trainGenerator,
    steps_per_epoch=100,
    epochs=100,
    validation_data=validationGenerator,
    validation_steps=50,
    callbacks=[cp_callback_2])

model.save("/home/marcin/Documents/TensorflowProjects/cats_and_dogs_small_4_fine_tunning.h5")



acc = history.history['acc']
val_acc = history.history['val_acc']
loss = history.history['loss']
val_loss = history.history['val_loss']
epochs = range(1, len(acc) + 1)
plt.plot(epochs, acc, 'bo', label='Training acc')
plt.plot(epochs, val_acc, 'b', label='Validation acc')
plt.title('Training and validation accuracy with fine tunning')
plt.legend()
plt.figure()
plt.plot(epochs, loss, 'bo', label='Training loss')
plt.plot(epochs, val_loss, 'b', label='Validation loss')
plt.title('Training and validation loss with fine tunning')
plt.legend()
plt.show()

# smoothed graphs
plt.plot(epochs,
smooth_curve(acc), 'bo', label='Smoothed training acc')
plt.plot(epochs,
smooth_curve(val_acc), 'b', label='Smoothed validation acc')
plt.title('Training and validation accuracy')
plt.legend()
plt.figure()
plt.plot(epochs,
smooth_curve(loss), 'bo', label='Smoothed training loss')
plt.plot(epochs,
smooth_curve(val_loss), 'b', label='Smoothed validation loss')
plt.title('Training and validation loss')
plt.legend()
plt.show()

# model evaluation
test_generator = testDataGen.flow_from_directory(
    testDir,
    target_size=(150, 150),
    batch_size=20,
    class_mode='binary')

test_loss, test_acc = model.evaluate(test_generator, steps=50)
print('Test accuracy =', test_acc)















