from keras.layers import Conv2D, MaxPooling2D, Input, Dense, Flatten, Activation, add, Lambda
from keras.layers.normalization import BatchNormalization
from keras.layers.pooling import GlobalAveragePooling2D
from keras.optimizers import RMSprop
from keras.backend import tf as ktf
from keras.models import Model, Sequential, load_model
from keras.callbacks import ModelCheckpoint, EarlyStopping
from keras.applications.resnet50 import ResNet50
from lib.data_utils import get_MNIST_data

"""
Explanation: CNN transfer learning - Keras+TensorFlow
This is for CNN models transferred from pretrained model, using Keras based on TensorFlow. First, some preparation work.
End of explanation
"""


data = get_MNIST_data(num_validation=0, fit=True)

# see if we get the data correctly
print('image size: ', data['X_train'].shape)

"""
Explanation: Read the MNIST data. Notice that we assume that it's 'kaggle-DigitRecognizer/data/train.csv', and we use helper function to read into a dictionary.
End of explanation
"""


# build the model
# preprocess to (28,28,3), then build a resize layer using tf.resize_images() to (224,224,3) as input
inputs = Input(shape=(28,28,3))
inputs_resize = Lambda(lambda img: ktf.image.resize_images(img, (224,224)))(inputs) # resize layer
resnet50 = ResNet50(include_top=False, input_tensor=inputs_resize, input_shape=(224,224,3), pooling='avg')
x = resnet50.output
#x = Dense(units=1024, activation='relu')(x)
predictions = Dense(units=10, activation='softmax')(x)

# connect the model
freezemodel = Model(inputs=inputs, outputs=predictions)
#freezemodel.summary()

# freeze all ResNet50 layers
for layer in resnet50.layers:
    layer.trainable = False

# set the loss and optimizer
freezemodel.compile(optimizer='rmsprop', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# fit the model
checkpoint = ModelCheckpoint('../models/freezeResNet_{epoch:02d}-{loss:.2f}.h5',
                             monitor='loss',
                             save_best_only=True)
freezemodel.fit(data['X_train'], data['y_train'].reshape(-1,1),
                batch_size=16, epochs=10, callbacks=[checkpoint], initial_epoch=1)

# test the model and see accuracy
score = freezemodel.evaluate(data['X_test'], data['y_test'].reshape(-1, 1), batch_size=32)
print(score)

# save the model: 0.96
freezemodel.save('ResNet50_freeze.h5')

# continue the model training
freezemodel = load_model('../models/ResNet50_freeze.h5', custom_objects={'ktf': ktf})

# set the loss and optimizer
rmsprop = RMSprop(lr=0.0001)
freezemodel.compile(optimizer=rmsprop, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# fit the model
checkpoint = ModelCheckpoint('../models/freezeResNet_{epoch:02d}-{loss:.2f}.h5',
                             monitor='loss',
                             save_best_only=True)
freezemodel.fit(data['X_train'], data['y_train'].reshape(-1, 1),
                batch_size=16, epochs=10, callbacks=[checkpoint], initial_epoch=4)

"""
Explanation: Freeze-weights transfer
We would use ResNet50 provided in Keras. In this section, the pretrained model would all be freezed, and new output layer would be attatched to the model, and only this output layer would be trained.
End of explanation
"""


# build the model
# preprocess to (28,28,3), then build a resize layer using tf.resize_images() to (224,224,3) as input
inputs = Input(shape=(28,28,3))
inputs_resize = Lambda(lambda img: ktf.image.resize_images(img, (224,224)))(inputs) # resize layer
resnet50 = ResNet50(include_top=False, input_tensor=inputs_resize, input_shape=(224,224,3), pooling='avg')
x = resnet50.output
#x = Dense(units=1024, activation='relu')(x)
predictions = Dense(units=10, activation='softmax')(x)

# connect the model
tunemodel = Model(inputs=inputs, outputs=predictions)
#freezemodel.summary()

# set the loss and optimizer
rmsprop = RMSprop(lr=0.0001)
tunemodel.compile(optimizer=rmsprop, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# fit the model
checkpoint = ModelCheckpoint('../models/tuneResNet_{epoch:02d}-{loss:.2f}.h5',
                             monitor='loss',
                             save_best_only=True)
tunemodel.fit(data['X_train'], data['y_train'].reshape(-1, 1),
                batch_size=16, epochs=10, callbacks=[checkpoint], initial_epoch=0)

# test the model and see accuracy
score = tunemodel.evaluate(data['X_test'], data['y_test'].reshape(-1, 1), batch_size=32)
print(score)

"""
Explanation: Fine-tune transfer
In this section, the model is the same as before, but all weights are trained along with the final layer using smaller learning rate.
End of explanation
"""


# build the model
# preprocess to (28,28,3), then build a resize layer using tf.resize_images() to (224,224,3) as input
inputs = Input(shape=(28,28,3))
inputs_resize = Lambda(lambda img: ktf.image.resize_images(img, (224,224)))(inputs) # resize layer
resnet50 = ResNet50(include_top=False, input_tensor=inputs_resize, input_shape=(224,224,3), pooling='avg')
x = resnet50.output
predictions = Dense(units=10, activation='softmax')(x)

# connect the model
tunemodel = Model(inputs=inputs, outputs=predictions)

# set the loss and optimizer
rmsprop = RMSprop(lr=0.0001)
tunemodel.compile(optimizer=rmsprop, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# fit the model
checkpoint = ModelCheckpoint('../models/tuneResNet_early_{epoch:02d}-{loss:.2f}.h5',
                             monitor='loss',
                             save_best_only=True)
earlystop = EarlyStopping(min_delta=0.0001, patience=1)
tunemodel.fit(data['X_train'], data['y_train'].reshape(-1, 1),
                batch_size=16, epochs=10, callbacks=[checkpoint, earlystop], initial_epoch=0)

# test the model and see accuracy
score = tunemodel.evaluate(data['X_test'], data['y_test'].reshape(-1, 1),
                           batch_size=16, validation_data=(data['X_test'], data['y_test'].reshape(-1, 1)))
print(score)

"""
Explanation: Fine-tune transfer with early stopping
Based on the previous section, the test set is used as the validation set, so as to monitor for early stopping.
End of explanation
"""


from lib.data_utils import create_submission
from keras.models import load_model

# for freeze ResNet50 model (3 epochs)
simple_CNN = load_model('../models/freezeResNet_03-0.09.h5', custom_objects={'ktf': ktf})
print('Load model successfully.')
create_submission(simple_CNN, '../data/test.csv', '../submission/submission_freezeResNet_03.csv', 16, fit=True)

"""
Explanation: Create submissions
Load the saved trained models and produce predictions for submission on Kaggle.
End of explanation
"""


