import os
import keras
import psutil
import configparser
import os
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import scripts.tools.utils as tools_utils
import argparse
import json
class LossHistory(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.losses = []
        self.acc = []
        self.mem_info_list = []

    def on_batch_end(self, batch, logs=None):
        process = psutil.Process()
        self.losses.append(logs.get('loss'))
        self.acc.append(logs.get('acc'))
        self.mem_info_list.append(process.memory_info().rss / (1024**3))

def save_log_txt(model, path, name, bk, x_train, y_train):
    history_loss = LossHistory()
    model.fit(x_train, y_train, epochs=2, batch_size=1024, validation_split=0.2, verbose=1, callbacks=[history_loss])

    model_name = name.split("/")[-1].split("_")[0]
    method_name = name.split("/")[-1].split("_")[1]
    valid_acc = [acc for acc in history_loss.acc if acc is not None]
    result_data = {
        "model": model_name,
        "method": method_name,
        "result": {
            "Losses": sum(history_loss.losses)/len(history_loss.losses),
            "Accuracy": sum(valid_acc)/len(valid_acc),
            "MemoryInfoList": sum(history_loss.mem_info_list)/len(history_loss.mem_info_list)
        }
    }

    if os.path.exists(path):
        with open(path, 'r') as json_file:
            data = json.load(json_file)
    else:
        data = []
    data.append(result_data)
    
    with open(path, 'w') as json_file:
        json.dump(data, json_file, indent=4)
        

def custom_objects():
    def no_activation(x):
        return x
    def leakyrelu(x):
        import keras.backend as K
        return K.relu(x, alpha=0.01)
    objects = {}
    objects['no_activation'] = no_activation
    objects['leakyrelu'] = leakyrelu
    return objects

def model_to_txt1(model_path, bk):
    
    cur_model = keras.models.load_model(model_path, custom_objects=custom_objects())
    cur_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

    data = tools_utils.DataUtils
    
    
    if 'svhn' in model_path or 'fashion2' in model_path:
        txt_path = model_path.split("\\")[-1][:-5].split("mut_model")[0]+bk+".json"
    else:
        txt_path = model_path.split("\\")[-1][:-3].split("mut_model")[0]+bk+".json"
    if 'svhn' in model_path or 'fashion2' in model_path:
        model_path = model_path.split("\\")[-1][:-5]
    else:
        model_path = model_path.split("\\")[-1][:-3]
    data_path = model_path.split("/")[-1]
    x_test, y_test = data.get_data_by_exp(data_path)
    save_log_txt(cur_model,txt_path,model_path,bk,x_test, y_test)


if __name__ == "__main__":

    """Parser of command args"""
    parse = argparse.ArgumentParser()
    parse.add_argument("--backend", type=str, help="name of backends")
    parse.add_argument("--model_path", type=str, help="redis db port")
    parse.add_argument("--root_dir", type=str, help="redis db port")
    flags, unparsed = parse.parse_known_args(sys.argv[1:])

    """Load Configuration"""
    lemon_cfg = configparser.ConfigParser()
    # lemon_cfg.read(f"./config/{flags.config_name}")
    
    conf_path = os.path.join(os.path.dirname(os.getcwd()), "config", "demo.conf")
    lemon_cfg.read(conf_path)
    parameters = lemon_cfg['parameters']
    gpu_ids = parameters['gpu_ids']
    gpu_list = parameters['gpu_ids'].split(",")

    """Init cuda"""
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids

    """Switch backend"""
    bk_list = ['tensorflow', 'theano', 'cntk', 'mxnet']
    bk = flags.backend
    print('.........................',type(bk))
    os.environ['KERAS_BACKEND'] = bk
    os.environ['PYTHONHASHSEED'] = '0'

    if bk == 'tensorflow':
        os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'  # 只显示 warning 和 Error
        import tensorflow as tf

    if bk == 'theano':
        if len(gpu_list) == 2:
            os.environ[
                'THEANO_FLAGS'] = f"device=cuda,contexts=dev{gpu_list[0]}->cuda{gpu_list[0]};dev{gpu_list[1]}->cuda{gpu_list[1]}," \
                                  f"force_device=True,floatX=float32,lib.cnmem=1"
        else:
            os.environ['THEANO_FLAGS'] = f"device=cuda,contexts=dev{gpu_list[0]}->cuda{gpu_list[0]}," \
                                         f"force_device=True,floatX=float32,lib.cnmem=1"
        batch_size = 32
        import theano as th

        mylogger.info(th.__version__)
    if bk == "cntk":
        batch_size = 32
        from cntk.device import try_set_default_device, gpu
        try_set_default_device(gpu(int(gpu_list[0])))
        import cntk as ck

    if bk == "mxnet":
        batch_size = 32
        import mxnet as mxnet

    from keras import backend as K
    import keras
    print("Using {} as backend for states extraction| {} is wanted".format(K.backend(),bk))
    """Get model hidden output on selected_index data on specific backend"""
    model_to_txt1(flags.model_path, bk)