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 keras import metrics

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 CNN():
    ## 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)
    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()(m_pool2)                             # 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 start and end
    model = Model(inputs=inputs, outputs=classif)
    # Optimizer recommended Adadelta values (lr=0.01)
    model.compile(optimizer=Adam(), loss='binary_crossentropy', metrics=['accuracy'])  

    return model

'''
Model definition for a fully convolutional (no dense layers) 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)
    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()(conv4)                             # Makes data 1D
    drop3 = Dropout(0.2)(flat)
    classif = Dense(2, activation='sigmoid')(drop3)    # Final layer to classify

    ## Define the model start and end
    model = Model(inputs=inputs, outputs=classif)
    # Optimizer recommended Adadelta values (lr=0.01)
    model.compile(optimizer=Adam(), loss='binary_crossentropy', metrics=['accuracy'])  

    return model


'''
Model definition for the network structure of LeNet
Note: LeNet was designed to classify into 10 classes, but we are only performing binary classification
'''
def LeNet():
    ## List of model layers
    inputs = Input((3, 64, 64))

    conv1 = Conv2D(6, (5, 5), activation='relu', padding='valid', input_shape=(64, 64, 3))(inputs)
    m_pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)

    conv2 = Conv2D(16, (5, 5), activation='relu', padding='valid')(m_pool1)
    m_pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)

    flat = Flatten()(m_pool2)                               # Makes data 1D
    
    dense1 = Dense(120, activation='relu')(flat)          # Fully connected layer
    dense2 = Dense(84, activation='relu')(dense1)         # Fully connected layer
    drop3 = Dropout(0.2)(dense2)
    classif = Dense(2, activation='sigmoid')(drop3)    # Final layer to classify

    ## Define the model start and end
    model = Model(inputs=inputs, outputs=classif)
    model.compile(optimizer=Adam(), loss='binary_crossentropy', metrics=['accuracy'])  

    return model

'''
AlexNet architecture
'''
def AlexNet():
    #inputs = Input(shape=(3, 64, 64))

    
    return model


def f1(y_true, y_pred):
    def recall(y_true, y_pred):
        """Recall metric.

        Only computes a batch-wise average of recall.

        Computes the recall, a metric for multi-label classification of
        how many relevant items are selected.
        """
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
        recall = true_positives / (possible_positives + K.epsilon())
        return recall

    def precision(y_true, y_pred):
        """Precision metric.

        Only computes a batch-wise average of precision.

        Computes the precision, a metric for multi-label classification of
        how many selected items are relevant.
        """
        true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
        predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
        precision = true_positives / (predicted_positives + K.epsilon())
        return precision
    precision = precision(y_true, y_pred)
    recall = recall(y_true, y_pred)
    return 2*((precision*recall)/(precision+recall+K.epsilon()))
# def precision(y_true, y_pred):
#     y_pred = K.round(y_pred)
#     num = K.sum(tf.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(tf.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 = CNN()
model = FCN()
#model = LeNet()
# 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 = 25             # 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)