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 sklearn.metrics import accuracy_score
from _image_classifier import ImageClassifier

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

from keras.utils import to_categorical
'''
Model definition define the network structure
'''
def FCN():
    ## List of model layers
    inputs = Input((3, 64, 64))

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

    conv2 = Conv2D(32, (3, 3), 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))(conv2)

    conv3 = Conv2D(32, (3, 3), 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(64, (2, 2), activation='relu', padding='same')(m_pool2)

    flat = Flatten()(conv3)                             # Makes data 1D
    dense = Dense(64, activation='relu')(flat)         # Fully connected layer
    drop3 = Dropout(0.2)(dense)
    classif = Dense(2, activation='sigmoid')(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=Adam(), loss='binary_crossentropy', metrics=['accuracy'])

    return model

def precision(y_true, y_pred):
    y_pred = K.round(y_pred)
    num = K.sum(K.logical_and(y_true, y_pred))
    den = K.sum(y_pred)
    return K.divide(num, den)

def recall(y_true, y_pred):
    y_pred = K.round(y_pred)
    num = K.sum(K.logical_and(y_true, y_pred))
    den = K.sum(y_true)
    return K.divide(num, den)

def f_measure(y_true, y_pred):
    p = precision(y_true, y_pred)
    r = recall(y_true, y_pred)
    return 2 * p * r / (p + r)

## 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')

lbl_train = to_categorical(lbl_train)       # One hot encoding the labels
lbl_test = to_categorical(lbl_test)

## 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 = 150        # batch size - number of training example used in one forward/backward pass
# (higher batch size uses more memory, smaller batch size takes more time)
#lrate = 0.01           # Learning rate of the model - controls magnitude of weight changes in training the NN
#decay = lrate/epochs   # Decay rate of the model

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

# Outputs the model structure
print(model.summary())

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
start = t.time()
# Passes the dataset through the model
pred_lbl = model.predict(im_test, verbose=1, batch_size=batch_size)
end = t.time()
pred_lbl = np.round(pred_lbl)
accuracy = accuracy_score(lbl_test, pred_lbl)
print("Accuracy: " + str(accuracy))
print("Images generated in {} seconds".format(end - start))
np.save('predicted_results.npy', pred_lbl)