import data_utils
import model_utils
import plot_utils
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import sys
import os
import tensorflow as tf
import numpy as np
import random
from tensorflow import keras


os.environ['PYTHONHASHSEED'] = str(76)
random.seed(45)
np.random.seed(1)
tf.random.set_seed(346)
model_name = sys.argv[1]
#model_name = "resnet_3s3DH"
print(model_name)
div = model_name.split("_")[1]
if div[0:2] == "3D" or div[0:2] == "2D":
    data_location = "../data/data_split/"
if div[0:2] == "n3" or div[0:2] == "n2":
    data_location = "../data/data_split_n/"
if div[0:2] == "3s":
    data_location = "../data/data_split_3s/"
print(data_location)
#data_location = "../data/data_split_n/"
if div[-3:] == "3DH":
    ddhs = 0
if div[-3:] == "2DH":
    ddhs = 1

train, test, train_targ, test_targ, train_ID, test_ID = data_utils.read_data_folders(data_location, 10000, 2000, ddh = ddhs)

train = train.reshape(len(train), 51, 51, 3)
test = test.reshape(len(test), 51, 51, 3)
print(train.shape, test.shape)

def premade(model_name, train, test, train_targ, test_targ, checkpoints=0):

    random.seed(45)
    np.random.seed(1)
    tf.random.set_seed(346)

    # premade = keras.applications.EfficientNetB7(include_top = False, weights = None)
    # layer = keras.layers.ZeroPadding2D((197-51, 197-51), input_shape=(51, 51, 3))
    premade = keras.applications.ResNet50V2(include_top = False, weights = None, input_shape=(51, 51, 3))
    #premade = keras.applications.VGG16(include_top = False, weights = None, input_shape=(51, 51, 3))
    layer7 = keras.layers.Dropout(0.4)
    layer8 = keras.layers.Flatten()
    layer9 = keras.layers.Dense(32, activation="relu")
    layer10 = keras.layers.Dense(2, activation="softmax")
    layers = [premade, layer7, layer8, layer9, layer10]

    model = keras.Sequential(layers)

    opt = keras.optimizers.SGD(learning_rate=0.01)
    model.compile(optimizer=opt, loss="sparse_categorical_crossentropy", metrics=["accuracy"])

    filepath = "../outputs/models/model_premade_%s.h5"%model_name
    print(train.shape, test.shape)
    print('unique test: {}'.format(np.unique(test_targ, return_counts=True)))
    feat_te2 = test.copy()
    feat_tr2 = train.copy()
    
    model.summary()
    
    # -- define the checkpoint
    checkpoint = tf.keras.callbacks.ModelCheckpoint(filepath, monitor='val_accuracy', verbose=1,save_best_only=True, mode='max')
    earlyst = tf.keras.callbacks.EarlyStopping(monitor='val_loss', mode='min', verbose=0, patience=100)
    # -- fit the model
    if checkpoints == 0:
        history = model.fit(feat_tr2, train_targ, validation_split=0.20, epochs=650, batch_size = 200, 
                            callbacks=[checkpoint, earlyst])

    # -- print the accuracy
    loss_tr, acc_tr = model.evaluate(feat_tr2, train_targ, verbose=0)
    loss_te, acc_te = model.evaluate(feat_te2, test_targ, verbose=0)

    print("Training accuracy : {0:.4f}".format(acc_tr))
    print("Testing accuracy  : {0:.4f}".format(acc_te))
    
    pd.DataFrame(history.history["loss"]).to_csv("trainc_loss%s"%model_name)
    pd.DataFrame(history.history["val_loss"]).to_csv("testc_loss%s"%model_name)
    pd.DataFrame(history.history["accuracy"]).to_csv("trainc_acc%s"%model_name)
    pd.DataFrame(history.history["val_accuracy"]).to_csv("testc_acc%s"%model_name)

    fig, ax = plt.subplots(2,1,figsize=(8, 8))
    ax[0].plot(history.history["loss"], color = "#00441b")
    ax[0].plot(history.history["val_loss"], color = "#40004b")
    ax[0].legend(["train", "validation"], loc="upper right")
    #ax[0].set_xlabel("epoch", fontsize=15)
    ax[0].set_xticks([])
    ax[0].set_ylabel("Loss", fontsize=15)

    ax[1].plot(history.history["accuracy"], color = "#00441b")
    ax[1].plot(history.history["val_accuracy"], color = "#40004b")
    ax[1].legend(["train", "validation"], loc="lower right")
    ax[1].set_xlabel("Epoch", fontsize=15)
    ax[1].set_ylabel("Accuracy", fontsize=15)
    if checkpoints == 0:
        plt.savefig("../outputs/lossacc_model%s.pdf"%model_name,bbox_inches="tight")
    else:
        plt.savefig("../outputs/lossacc_model%s_checkpoint.pdf"%model_name, bbox_inches="tight")

    y_pred_test = model_utils.predict_data(test, model, 0, test_ID,test_targ)
    if checkpoints == 0:
        plot_utils.create_confusion_matrix("test_%s"%model_name, y_pred_test, test_targ)
    else:
        plot_utils.create_confusion_matrix("test_%s_checkpoint"%model_name, y_pred_test, test_targ)

    y_pred_train = model_utils.predict_data(train, model, 0, train_ID, train_targ)
    if checkpoints == 0:
        plot_utils.create_confusion_matrix("train_%s"%model_name, y_pred_train, train_targ)
    else:
        plot_utils.create_confusion_matrix("train_%s_checkpoint"%model_name, y_pred_train, train_targ)

    return train, test, train_targ, test_targ, train_ID, test_ID
    
    
train, test, train_targ, test_targ, train_ID, test_ID = premade(model_name, train, test, train_targ, test_targ, checkpoints=0)
