import logging
import numpy as np
import os
import PIL
import PIL.Image
import tensorflow as tf

from tensorflow.keras.layers import Layer, Conv2D, MaxPool2D, Dense, Flatten, Dropout, GlobalAveragePooling2D
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras import layers
from tensorflow.keras import Model

img_height = 224
img_width = 224
batch_size = 64

data_dir = './100-bird-species/'
data_dir_train = os.path.join(data_dir, 'train')
data_dir_valid = os.path.join(data_dir, 'valid')
data_dir_test = os.path.join(data_dir, 'test')

train_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir_train,
    label_mode='categorical',
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

valid_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir_valid,
    label_mode='categorical',
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)

test_ds = tf.keras.utils.image_dataset_from_directory(
    data_dir_test,
    label_mode='categorical',
    seed=123,
    image_size=(img_height, img_width),
    batch_size=batch_size)


def normalize(img, label):
    return img / 255.0, label


data_augmentation = tf.keras.Sequential([
    tf.keras.layers.RandomFlip("horizontal"),
    tf.keras.layers.RandomRotation(0.2),
    tf.keras.layers.RandomZoom(0.2)
])

train_dataset = (train_ds
                 .map(normalize)
                 .map(lambda x, y: (data_augmentation(x), y))
                 .prefetch(tf.data.AUTOTUNE))

valid_dataset = valid_ds.map(normalize)
test_dataset = test_ds.map(normalize)


def get_birds_mobilenet():
    pre_trained_model = MobileNetV2(
        include_top=False,
        input_shape=(img_height, img_width, 3),
        classifier_activation='softmax'
    )

    for layer in pre_trained_model.layers:
        layer.trainable = False

    last_layer = pre_trained_model.output
    last_layer.trainable = True

    x = GlobalAveragePooling2D()(last_layer)
    x = Dense(1024, activation='relu')(x)
    x = layers.Dense(325, activation='softmax')(x)

    model = Model(pre_trained_model.input, x)
    return model


model = get_birds_mobilenet()
model.summary()

model.compile(loss='categorical_crossentropy',
              optimizer='adam', metrics=['accuracy'])
checkpoint_path = "./checkpoints/birds_mobilenet/"

model.load_weights(checkpoint_path)

model_history = model.fit(
    train_dataset,
    validation_data=valid_dataset,
    epochs=200,
    callbacks=[
        #tf.keras.callbacks.EarlyStopping(patience=5),
        tf.keras.callbacks.ModelCheckpoint(
            filepath=checkpoint_path, verbose=0, save_freq="epoch")
    ])

tensorflow.keras.utils.image_dataset_from_directory tensorflow.keras.callbacks.ModelCheckpoint tensorflow.keras.layers.GlobalAveragePooling2D tensorflow.keras.layers.RandomFlip tensorflow.keras.layers.Dense tensorflow.keras.Model tensorflow.keras.layers.RandomZoom tensorflow.keras.layers.RandomRotation tensorflow.keras.applications.MobileNetV2
[(18, 'os.path.join', 'os.path.join', (['data_dir', '"""train"""'], {}), False, 'import os\n'), (19, 'os.path.join', 'os.path.join', (['data_dir', '"""valid"""'], {}), False, 'import os\n'), (20, 'os.path.join', 'os.path.join', (['data_dir', '"""test"""'], {}), False, 'import os\n'), (22, 'tensorflow.keras.utils.image_dataset_from_directory', 'tf.keras.utils.image_dataset_from_directory', (['data_dir_train'], {'label_mode': '"""categorical"""', 'seed': '(123)', 'image_size': '(img_height, img_width)', 'batch_size': 'batch_size'}), True, 'import tensorflow as tf\n'), (29, 'tensorflow.keras.utils.image_dataset_from_directory', 'tf.keras.utils.image_dataset_from_directory', (['data_dir_valid'], {'label_mode': '"""categorical"""', 'seed': '(123)', 'image_size': '(img_height, img_width)', 'batch_size': 'batch_size'}), True, 'import tensorflow as tf\n'), (36, 'tensorflow.keras.utils.image_dataset_from_directory', 'tf.keras.utils.image_dataset_from_directory', (['data_dir_test'], {'label_mode': '"""categorical"""', 'seed': '(123)', 'image_size': '(img_height, img_width)', 'batch_size': 'batch_size'}), True, 'import tensorflow as tf\n'), (64, 'tensorflow.keras.applications.MobileNetV2', 'MobileNetV2', ([], {'include_top': '(False)', 'input_shape': '(img_height, img_width, 3)', 'classifier_activation': '"""softmax"""'}), False, 'from tensorflow.keras.applications import MobileNetV2\n'), (80, 'tensorflow.keras.Model', 'Model', (['pre_trained_model.input', 'x'], {}), False, 'from tensorflow.keras import Model\n'), (49, 'tensorflow.keras.layers.RandomFlip', 'tf.keras.layers.RandomFlip', (['"""horizontal"""'], {}), True, 'import tensorflow as tf\n'), (50, 'tensorflow.keras.layers.RandomRotation', 'tf.keras.layers.RandomRotation', (['(0.2)'], {}), True, 'import tensorflow as tf\n'), (51, 'tensorflow.keras.layers.RandomZoom', 'tf.keras.layers.RandomZoom', (['(0.2)'], {}), True, 'import tensorflow as tf\n'), (76, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'GlobalAveragePooling2D', ([], {}), False, 'from tensorflow.keras.layers import Layer, Conv2D, MaxPool2D, Dense, Flatten, Dropout, GlobalAveragePooling2D\n'), (77, 'tensorflow.keras.layers.Dense', 'Dense', (['(1024)'], {'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Layer, Conv2D, MaxPool2D, Dense, Flatten, Dropout, GlobalAveragePooling2D\n'), (78, 'tensorflow.keras.layers.Dense', 'layers.Dense', (['(325)'], {'activation': '"""softmax"""'}), False, 'from tensorflow.keras import layers\n'), (99, 'tensorflow.keras.callbacks.ModelCheckpoint', 'tf.keras.callbacks.ModelCheckpoint', ([], {'filepath': 'checkpoint_path', 'verbose': '(0)', 'save_freq': '"""epoch"""'}), True, 'import tensorflow as tf\n')]