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

def multiple_models(data_location, model_name, ddh = 0, checkpoints = 0):
    
    if checkpoints == 1:
        comd = "cp model_checkpoint_%s.h5 model_checkpoint_%s_copy.h5"%(model_name, model_name)
        model_name = model_name+"-moretrain"
        os.system(comd)
        
    train, test, train_targ, test_targ, train_ID, test_ID = data_utils.read_data_folders(data_location, 10000,
                                                                          2000, ddh = ddh)
    def diff_nowGN(data):
        np.random.seed(45)
        noise = np.random.normal(0,1,51*51*len(data)).reshape(len(data),51,51)
        traingn = np.concatenate((noise,data[:,:,51:]), axis=2)
        return traingn

    # -- feautres need to have an extra axis on the end (for mini-batching)
    if "3CHANNEL" in model_name.split("_")[0]: 
        feat_tr2 = train.reshape(len(train), 51, 51, 3)
        feat_te2 = test.reshape(len(test), 51, 51, 3)

    elif "2CHANNEL" in model_name.split("_")[0]:
        train = train[:,:,51:]
        test = test[:, :,51:]
        feat_tr2 = train.reshape(len(train), 51, 51, 2)
        feat_te2 = test.reshape(len(test), 51, 51, 2)

    elif "Gauss" in model_name.split("_")[0]:
        train = diff_nowGN(train)
        test = diff_nowGN(test)
        feat_tr2 = train.reshape(len(train), train.shape[1], train.shape[2], 1)
        feat_te2 = test.reshape(len(test), test.shape[1], test.shape[2], 1)
    else:
        # -- feautres need to have an extra axis on the end (for mini-batching)
        feat_tr2 = train.reshape(len(train), train.shape[1], train.shape[2], 1)
        feat_te2 = test.reshape(len(test), test.shape[1], test.shape[2], 1)

    #train = diff_nowGN(train)
    #test = diff_nowGN(test)

    #train = train.reshape(len(train), 51, 51, 3)
    #test = test.reshape(len(test), 51, 51, 3)

    print(feat_tr2.shape, feat_te2.shape)
    train = feat_tr2.copy()
    test = feat_te2.copy()
    feat_tr2 = None
    feat_te2 = None
    del feat_tr2
    del feat_te2
    
    print(train.shape, test.shape)
    if ddh == 0:
        history, model, times = model_utils.keras_model_AAD_and_save(model_name, train, test, 
                                                train_targ, test_targ, checkpoints = checkpoints)
    else:
        history, model, times = model_utils.keras_model_CC_and_save(model_name, train, test, 
                                                train_targ, test_targ, checkpoints = checkpoints)
        #history, model, times = model_utils.keras_model_CC_CONVextra(model_name, train, test, 
         #                                        train_targ, test_targ, checkpoints = checkpoints)

    '''
    pd.DataFrame(history.history["loss"]).to_csv("../outputs/trainc_loss%s"%model_name)
    pd.DataFrame(history.history["val_loss"]).to_csv("../outputs/testc_loss%s"%model_name)
    pd.DataFrame(history.history["accuracy"]).to_csv("../outputs/trainc_acc%s"%model_name)
    pd.DataFrame(history.history["val_accuracy"]).to_csv("../outputs/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/lossaccC_model%s.pdf"%model_name,bbox_inches="tight")
    else:
        plt.savefig("../outputs/lossaccC_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


if __name__ == "__main__":
    os.environ['PYTHONHASHSEED'] = str(76)
    random.seed(45)
    np.random.seed(1)
    tf.random.set_seed(346)
    name = sys.argv[1]
    model_name = name
    print(model_name)
    div = 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 = diff_nowGN(train)
    #test = diff_nowGN(test)

    #train = train.reshape(len(train), 51, 51, 3)
    #test = test.reshape(len(test), 51, 51, 3)

    #print(feat_tr2.shape, feat_te2.shape)
        
    train, test, train_targ, test_targ, train_ID, test_ID = multiple_models(data_location, model_name, ddh = ddhs, checkpoints = 0)
