import os
import numpy as np
import argparse
import shutil
import gc

import keras
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF

from model.FeatureNet import build_FeatureNet
from model.DataGenerator import kFoldGenerator
from utlis.Utils import *

if __name__ == "__main__":

    print(128 * '#')
    print('Start to train FeatureNet.')

    # # 1. Get configuration

    # ## 1.1. Read .config file

    # command line parameters -c -g
    parser = argparse.ArgumentParser()
    parser.add_argument("-c", type = str, help = "configuration file" ,default="./config/ISRUC.config")
    parser.add_argument("-g", type = str, help = "GPU number to use, set '-1' to use CPU" , default= 0)
    args = parser.parse_args()
    args.c = "./config/ISRUC.config"
    args.g = "0"
    Path, cfgFeature, _, _ = ReadConfig(args.c)

    # set GPU number or use CPU only
    os.environ["CUDA_VISIBLE_DEVICES"] = args.g
    if args.g != "-1":
        config = tf.ConfigProto()
        config.gpu_options.allow_growth=True
        sess = tf.Session(config=config)
        KTF.set_session(sess)
        print("Use GPU #"+args.g)
    else:
        print("Use CPU only")

    # ## 1.2. Analytic parameters

    # [train] parameters ('_f' means FeatureNet)
    channels   = int(cfgFeature["channels"])
    fold       = int(cfgFeature["fold"])
    num_epochs_f = int(cfgFeature["epoch_f"])
    batch_size_f = int(cfgFeature["batch_size_f"])
    optimizer_f  = cfgFeature["optimizer_f"]
    learn_rate_f = float(cfgFeature["learn_rate_f"])


    # ## 1.3. Parameter check and enable

    # Create save pathand copy .config to it
    if not os.path.exists(Path['Save']):
        os.makedirs(Path['Save'])
    shutil.copyfile(args.c, Path['Save']+"last.config")


    # # 2. Read data and process data

    # ## 2.1. Read data
    # Each fold corresponds to one subject's data (ISRUC-S3 dataset)
    ReadList = np.load(Path['data'], allow_pickle=True)
    Fold_Num   = ReadList['Fold_len']    # Num of samples of each fold
    Fold_Data  = ReadList['Fold_data']   # Data of each fold
    Fold_Label = ReadList['Fold_label']  # Labels of each fold

    print("Read data successfully")
    print('Number of samples: ',np.sum(Fold_Num))

    # ## 2.2. Build kFoldGenerator or DominGenerator
    # 只是把data和label读出来了
    DataGenerator = kFoldGenerator(Fold_Data, Fold_Label)


    # # 3. Model training (cross validation)

    # k-fold cross validation
    all_scores = []
    for i in range(fold):
        print(128*'_')
        print('Fold #', i)

        # Instantiation optimizer
        opt_f = Instantiation_optim(optimizer_f, learn_rate_f) # optimizer of FeatureNet

        # get i th-fold data
        train_data,train_targets,val_data,val_targets = DataGenerator.getFold(i)

        ## build FeatureNet & train
        featureNet, featureNet_p = build_FeatureNet(opt_f, channels) # '_p' model is without the softmax layer
        history_fea = featureNet.fit(
            x = train_data,
            y = train_targets,
            epochs = num_epochs_f,
            batch_size = batch_size_f,
            shuffle = True,
            validation_data = (val_data, val_targets),
            verbose = 2,
            callbacks=[keras.callbacks.ModelCheckpoint(Path['Save']+'FeatureNet_Best_'+str(i)+'.h5',
                                                       monitor='val_acc',
                                                       verbose=0,
                                                       save_best_only=True,
                                                       save_weights_only=False,
                                                       mode='auto',
                                                       period=1 )])
        # Save training information
        if i==0:
            fit_loss = np.array(history_fea.history['loss'])*Fold_Num[i]
            fit_acc = np.array(history_fea.history['acc'])*Fold_Num[i]
            fit_val_loss = np.array(history_fea.history['val_loss'])*Fold_Num[i]
            fit_val_acc = np.array(history_fea.history['val_acc'])*Fold_Num[i]
        else:
            fit_loss = fit_loss+np.array(history_fea.history['loss'])*Fold_Num[i]
            fit_acc = fit_acc+np.array(history_fea.history['acc'])*Fold_Num[i]
            fit_val_loss = fit_val_loss+np.array(history_fea.history['val_loss'])*Fold_Num[i]
            fit_val_acc = fit_val_acc+np.array(history_fea.history['val_acc'])*Fold_Num[i]

        # load the weights of best performance
        featureNet.load_weights(Path['Save']+'FeatureNet_Best_'+str(i)+'.h5')

        # get and save the learned feature
        train_feature = featureNet_p.predict(train_data)
        val_feature = featureNet_p.predict(val_data)
        print('Save feature of Fold #' + str(i) + ' to' + Path['Save']+'Feature_'+str(i) + '.npz')
        np.savez(Path['Save']+'Feature_'+str(i)+'.npz',
            train_feature = train_feature,
            val_feature = val_feature,
            train_targets = train_targets,
            val_targets = val_targets
        )

        saveFile = open(Path['Save'] + "Result_FeatureNet.txt", 'a+')
        print('Fold #'+str(i), file=saveFile)
        print(history_fea.history, file=saveFile)
        saveFile.close()

        # Fold finish
        keras.backend.clear_session()
        del featureNet, featureNet_p, train_data, train_targets, val_data, val_targets, train_feature, val_feature
        gc.collect()

    print(128 * '_')

    print('End of training FeatureNet.')
    print(128 * '#')
