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

from tensorflow.keras.preprocessing import image_dataset_from_directory

_URL = 'https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip'
path_to_zip = tf.keras.utils.get_file('cats_and_dogs.zip', origin=_URL, extract=True)
PATH = os.path.join(os.path.dirname(path_to_zip), 'cats_and_dogs_filtered')

train_dir = os.path.join(PATH, 'train')
validation_dir = os.path.join(PATH, 'validation')

BATCH_SIZE = 32
IMG_SIZE = (160, 160)

train_dataset = image_dataset_from_directory(
    train_dir,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE
)

validation_dataset = image_dataset_from_directory(
    validation_dir,
    shuffle=True,
    batch_size=BATCH_SIZE,
    image_size=IMG_SIZE
)

class_names = train_dataset.class_names

''' Function for plotting the first 9 pictures in the training set from the dataset. '''
def show_first_9_pictures_in_dataset(train_dataset, class_names):
    for images, labels in train_dataset.take(1):
        for i in range(9):
            ax = plt.subplot(3, 3, i + 1)
            plt.imshow(images[i].numpy().astype("uint8"))
            plt.title(class_names[labels[i]])
            plt.axis("off")
    plt.show()

# show_first_9_pictures_in_dataset(train_dataset=train_dataset, class_names=class_names)


val_batches = tf.data.experimental.cardinality(validation_dataset)
test_dataset = validation_dataset.take(val_batches // 5)
validation_dataset = validation_dataset.skip(val_batches // 5)

print('Number of validation batches: %d' % tf.data.experimental.cardinality(validation_dataset))
print('Number of test batches: %d' % tf.data.experimental.cardinality(test_dataset))

'''
    CONFIGURE THE DATASET FOR PERFORMANCE
'''
AUTOTUNE = tf.data.experimental.AUTOTUNE
train_dataset = train_dataset.prefetch(buffer_size=AUTOTUNE)
validation_dataset = validation_dataset.prefetch(buffer_size=AUTOTUNE)
test_dataset = test_dataset.prefetch(buffer_size=AUTOTUNE)

'''
    USE DATA AUGMENTATION
'''
data_augmentation = tf.keras.Sequential([
    tf.keras.layers.experimental.preprocessing.RandomFlip('horizontal'),
    tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
])

''' Plotting the data augmentation process with 1 random picture to illustrate what happens. '''
def plot_data_augmentation(train_dataset):
    for image, _ in train_dataset.take(1):
        plt.figure(figsize=(10, 10))
        first_image = image[0]
        for i in range(9):
            ax = plt.subplot(3, 3, i + 1)
            augmented_image = data_augmentation(tf.expand_dims(first_image, 0))
            plt.imshow(augmented_image[0] / 255)
            plt.axis('off')
    plt.show()

# plot_data_augmentation(train_dataset=train_dataset)


'''
RESCALE PIXEL VALUES
'''
process_input = tf.keras.applications.mobilenet_v2.preprocess_input
rescale = tf.keras.layers.experimental.preprocessing.Rescaling(1./127.5, offset= -1)

'''
    CREATE THE BASE MODEL FROM THE PRE-TRAINED CONVNETS
'''
IMG_SHAPE = IMG_SIZE + (3, )

base_model = tf.keras.applications.MobileNetV2(
    input_shape=IMG_SHAPE,
    include_top=False,
    weights='imagenet'
)

image_batch, label_batch = next(iter(train_dataset))
feature_batch = base_model(image_batch)

print('Feature_batch.shape: ', feature_batch.shape)

base_model.trainable = False

# base_model.summary()

global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
feature_batch_average = global_average_layer(feature_batch)
print(feature_batch_average.shape)

prediction_layer = tf.keras.layers.Dense(1)
prediction_batch = prediction_layer(feature_batch_average)
print(prediction_batch.shape)

'''
    BUILDING THE MODEL
'''
inputs = tf.keras.Input(shape=(160, 160, 3))
x = data_augmentation(inputs)
x = process_input(x)
x = base_model(x, training=False)
x = global_average_layer(x)
x = tf.keras.layers.Dropout(0.2)(x)
outputs = prediction_layer(x)
model = tf.keras.Model(inputs, outputs)


'''
COMPILE THE MODEL
'''
base_learning_rate = 0.0001
model.compile(
    optimizer = tf.keras.optimizers.Adam(lr=base_learning_rate),
    loss = tf.keras.losses.BinaryCrossentropy(from_logits=True),
    metrics=['accuracy']
)

model.summary() # as seen the last layer, the dense one is the only trainable becuase rest of the model is frozen.

print(len(model.trainable_variables))


'''
    TRAIN THE MODEL
'''

initial_epochs = 20

loss0, accuracy0 = model.evaluate(validation_dataset)

print("initial loss: {:.2f}".format(loss0))
print("initial accuracy: {:.2f}".format(accuracy0))

history = model.fit(
    train_dataset,
    epochs=initial_epochs,
    validation_data=validation_dataset
)


'''
PLOT ACCURACY FOR LOSS AND VALIDATION
'''

acc = history.history['accuracy']
val_acc = history.history['val_accuracy']

loss = history.history['loss']
val_loss = history.history['val_loss']

plt.figure(figsize=(8, 8))
plt.subplot(2, 1, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.legend(loc='lower right')
plt.ylabel('Accuracy')
plt.ylim([min(plt.ylim()),1])
plt.title('Training and Validation Accuracy')

plt.subplot(2, 1, 2)
plt.plot(loss, label='Training Loss')
plt.plot(val_loss, label='Validation Loss')
plt.legend(loc='upper right')
plt.ylabel('Cross Entropy')
plt.ylim([0,1.0])
plt.title('Training and Validation Loss')
plt.xlabel('epoch')

'''
    FINE TUNE THE MODEL
    
    We want to find out how many layers the models has, since we dont want to train the whole model, cause that can reverse
    a lot of the good resuts in predicitons. Therefore we need to identify how many layes then train certain portion of them.
    The later the layers are in the model, the more specific they are, visa versa the earlier layers are more genralized. Thus
    the earlier layers are more generalized we can mold them more specificly towards our current dataset, dogs and cats. Therfore
    we are targeting the earlier layers instead of the later ones.
'''

# how many layers are the in the basemodel
print('Number of layers in the base model: ', len(base_model.layers))

# The point we want to stop finetuning
fine_tune_at = 100

# enable the model to be trainable
base_model.trainable = True

# from layer 100 and onwards to layer 154 trainalbe = False
for layer in base_model.layers[:fine_tune_at]:
    layer.trainable = False


model.compile(
    loss = tf.keras.losses.BinaryCrossentropy(from_logits=True),
    optimizer = tf.keras.optimizers.RMSprop(lr=base_learning_rate/10), # 0.00001
    metrics=['accuracy']
)

model.summary()

print(len(model.trainable_variables))


''' 
TRAIN THE FINE_TUNE MODEL
'''
fine_tune_epochs = 80
total_epochs = initial_epochs + fine_tune_epochs

history_fine = model.fit(
    train_dataset,
    epochs=total_epochs,
    initial_epoch=history.epoch[-1],
    validation_data=validation_dataset
)


'''
    PLOT THE MODEL
'''

acc += history_fine.history['accuracy']
val_acc += history_fine.history['val_accuracy']

loss += history_fine.history['loss']
val_loss += history_fine.history['val_loss']

plt.figure(figsize=(8, 8))
plt.subplot(2, 1, 1)
plt.plot(acc, label='Training Accuracy')
plt.plot(val_acc, label='Validation Accuracy')
plt.ylim([0.8, 1])
plt.plot([initial_epochs-1,initial_epochs-1],
          plt.ylim(), label='Start Fine Tuning')
plt.legend(loc='lower right')
plt.title('Training and Validation Accuracy')

plt.subplot(2, 1, 2)
plt.plot(loss, label='Training Loss')
plt.plot(val_loss, label='Validation Loss')
plt.ylim([0, 1.0])
plt.plot([initial_epochs-1,initial_epochs-1],
         plt.ylim(), label='Start Fine Tuning')
plt.legend(loc='upper right')
plt.title('Training and Validation Loss')
plt.xlabel('epoch')
plt.show()


'''
EVALUATION AND PREDICTION
'''
loss, accuracy = model.evaluate(test_dataset)
print('Test accuracy :', accuracy)

# see if the model is accurate enough to predict if the image is of a cat or a dog.

# Retrieve a batch of images from the test set.
image_batch, label_batch = test_dataset.as_numpy_iterator().next()
predictions = model.predict_on_batch(image_batch).flatten()

# Apply a sigmoid since our model returns logits
predictions = tf.nn.sigmoid(predictions)
# set the values that are smaller then 0.5 to 0, and the greater is set to 1
predictions = tf.where(predictions < 0.5, 0, 1)

print('Predictions:\n', predictions.numpy())
print('Labels:\n', label_batch)

plt.figure(figsize=(10, 10))
for i in range(9):
  ax = plt.subplot(3, 3, i + 1)
  plt.imshow(image_batch[i].astype("uint8"))
  plt.title(class_names[predictions[i]])
  plt.axis("off")
plt.show()

