from data.load import train_gen, x_test, y_test
from model.model import model, as_dense

from adversarial import adversarial_eval

from keras import backend as K
import tensorflow as tf
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
K.set_session(sess)

import keras
import numpy as np
m = model()

samples_per_epoch = 1000 # 1853666

optimizer = keras.optimizers.Adam()
m.compile(loss='binary_crossentropy', optimizer=optimizer, metrics=['accuracy', 'matthews_correlation'])


#tensorboard = keras.callbacks.TensorBoard(log_dir='./tensorboard', histogram_freq=1, write_graph=True)
saver = keras.callbacks.ModelCheckpoint('../models/blah.weights.{epoch:02d}-{val_loss:.2f}.hdf5', monitor='val_loss', verbose=0, save_best_only=False, save_weights_only=False, mode='auto')

epochs = 1 # 50
split = True
if split:
    history = m.fit_generator(train_gen(), samples_per_epoch=samples_per_epoch,
                    nb_epoch=epochs, pickle_safe=True, #nb_worker=4,
                    validation_data=(x_test, y_test),
                    callbacks=[saver])
else:
    history = m.fit(x_train, y_train,  batch_size=1000,
                    #train_gen(), samples_per_epoch=samples_per_epoch,
                    nb_epoch=epochs,
                    validation_data=(x_test, y_test),
                    callbacks=[saver])

#from IPython import embed; embed()

if True:
    y_pred = m.predict(x_test)
    y_truth = y_test

    y_round = np.copy(y_pred)
    y_round[y_round < 0.5] = 0
    y_round[y_round >= 0.5] = 1

    import matplotlib
    matplotlib.use('Agg') 
    import matplotlib.pyplot as plt

    from sklearn.metrics import classification_report, roc_curve, auc, confusion_matrix

    print (confusion_matrix(y_truth, y_round))
    print(classification_report(y_truth, y_round, target_names=['fine', 'defect']))

    false_positive_rate, true_positive_rate, thresholds = roc_curve(y_truth, y_pred, pos_label=1)
    roc_auc = auc(false_positive_rate, true_positive_rate)

    plt.title('Receiver Operating Characteristic')
    plt.plot(false_positive_rate, true_positive_rate, 'b', label='AUC = %0.2f'% roc_auc)
    plt.legend(loc='lower right')
    plt.plot([0,1],[0,1],'r--')
    plt.xlim([-0.1,1.2])
    plt.ylim([-0.1,1.2])
    plt.ylabel('True Positive Rate')
    plt.xlabel('False Positive Rate')
    plt.savefig('../figures/roc.png')

m2 = as_dense(m)

adversarial_eval(m2)