import numpy as np
import sys
import time as t
# Disables Tensorflow's warning about not utilising AVX/FMA
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
#from keras.models import Sequential
from keras.layers import Dense, Dropout, Activation, Flatten, Input
from keras.layers import Conv2D, MaxPooling2D, ZeroPadding2D
from keras.models import Model 

from sklearn import svm, tree, naive_bayes, ensemble
from _image_classifier import ImageClassifier

from keras.optimizers import Adadelta
from keras.callbacks import ModelCheckpoint
from keras import backend as K
K.set_image_dim_ordering('th')
np.random.seed(7)

'''
Model definition define the network structure
'''
def FCN():
    ## List of model layers
    inputs = Input((3, 64, 64))

    conv1 = Conv2D(8, (2, 2), activation='relu', padding='same', input_shape=(64, 64, 3))(inputs)
    m_pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(16, (2, 2), activation='relu', padding='same')(m_pool1)
    drop1 = Dropout(0.2)(conv2) # Drop some portion of features to prevent overfitting
    m_pool2 = MaxPooling2D(pool_size=(2, 2))(drop1)

    conv3 = Conv2D(32, (2, 2), activation='relu', padding='same')(m_pool2)
    drop2 = Dropout(0.2)(conv3) # Drop some portion of features to prevent overfitting
    m_pool2 = MaxPooling2D(pool_size=(2, 2))(drop2)
    
    conv4 = Conv2D(164, (2, 2), activation='relu', padding='same')(m_pool2)

    flat = Flatten()(conv4)                             # Makes data 1D
    dense = Dense(64, activation='relu')(flat)         # Fully connected layer
    drop3 = Dropout(0.2)(dense)                         
    classif = Dense(2, activation='softmax')(drop3)    # Final layer to classify

    ## Define the model structure
    model = Model(inputs=inputs, outputs=classif)
    # Optimizer recommended Adadelta values (lr=0.01)
    model.compile(optimizer=Adadelta(lr=0.1), loss='sparse_categorical_crossentropy', metrics=['accuracy'])  

    return model


## Open data
im_train = np.load('Waldo_train_data.npy')
lbl_train = np.load('Waldo_train_lbl.npy')
im_test = np.load('Waldo_test_data.npy')
lbl_test = np.load('Waldo_test_lbl.npy')

## Define model
model = FCN()
svm_iclf = ImageClassifier(svm.SVC)
tree_iclf = ImageClassifier(tree.DecisionTreeClassifier)
naive_bayes_iclf = ImageClassifier(naive_bayes.GaussianNBd)
ensemble_iclf = ImageClassifier(ensemble.RandomForestClassifier)

## Define training parameters
epochs = 20             # an epoch is one forward pass and back propogation of all training data 
batch_size = 5
#lrate = 0.01
#decay = lrate/epochs
# epoch - one forward pass and one backward pass of all training data
# batch size - number of training example used in one forward/backward pass
# (higher batch size uses more memory)
# learning rate - controls magnitude of weight changes in training the NN

## Train model
# Purely superficial output
sys.stdout.write("\nFitting model")
sys.stdout.flush()
for i in range(0, 3):
    t.sleep(0.8)
    sys.stdout.write('.')
    sys.stdout.flush()
print()

# Outputs the model structure
for i in range(0, len(model.layers)):
    print("Layer {}: {}".format(i, model.layers[i].output))
print('-'*30)

filepath = "checkpoint.hdf5"                                            # Defines the model checkpoint file
checkpoint = ModelCheckpoint(filepath, verbose=1, save_best_only=False) # Defines the checkpoint process
callbacks_list = [checkpoint]               # Adds the checkpoint process to the list of action performed during training
start = t.time()                            # Records time before training

# Fits model based on initial parameters
model.fit(im_train, lbl_train, epochs=epochs, batch_size=batch_size,
          verbose=2, shuffle=True, callbacks=callbacks_list)
# If getting a value error here, output of network and corresponding lbl_train
# data probably don't match
end = t.time()                              # Records time after tranining

print('Training Duration: {}'.format(end-start))
print('-'*30)
print("*** Saving FCN model and weights ***")

'''
# *To save model and weights separately:
# save model as json file
model_json = model.to_json()
with open("UNet_model.json", "w") as json_file:
    json_file.write(model_json)
# save weights as h5 file
model.save_weights("UNet_weights.h5")
print("\nModel weights and structure have been saved.\n")
'''
# Save model as one file
model.save('Waldo.h5')
print("\nModel weights and structure have been saved.\n")

## Testing the model
# Show data stats
print('*'*30)
print(im_test.shape)
print(lbl_test.shape)
print('*'*30)
start = t.time()
# Passes the dataset through the model
pred_lbl = model.predict(im_test, verbose=1, batch_size=batch_size)
end = t.time()
print("Images generated in {} seconds".format(end - start))
np.save('predicted_results.npy', pred_lbl)