import os
import random
import time

import numpy as np
import tensorflow as tf
from GPUtil import GPUtil as Gpu
from typeguard import typechecked

from baconian.common import files as file
from baconian.common.logging import Logger, ConsoleLogger
from baconian.config.global_config import GlobalConfig
from copy import deepcopy
import tracemalloc


def _reset_global_seed(seed):
    """

    :param seed: seed to reset global variables
    :type seed: int
    """
    sess = tf.get_default_session()
    if sess:

        sess.__exit__(None, None, None)
    tf.reset_default_graph()
    tf.set_random_seed(seed)
    np.random.seed(seed)
    random.seed(seed)


@typechecked
def single_exp_runner(task_fn, auto_choose_gpu_flag=False, gpu_id: int = 0, seed=None, del_if_log_path_existed=False,
                      **task_fn_kwargs):

    """

    :param task_fn: task function defined bu users
    :type task_fn: method
    :param auto_choose_gpu_flag: auto choose gpu, default False
    :type auto_choose_gpu_flag: bool
    :param gpu_id: gpu id, default 0
    :type gpu_id: int
    :param seed: seed generated by system time
    :type seed: int
    :param del_if_log_path_existed:delete obsolete log file path if existed, by default False
    :type del_if_log_path_existed: bool
    :param task_fn_kwargs:
    :type task_fn_kwargs:
    :return:
    :rtype:
    """
    os.environ['CUDA_DEVICE_ORDER'] = "PCI_BUS_ID"
    if auto_choose_gpu_flag is True:
        DEVICE_ID_LIST = Gpu.getFirstAvailable()
        DEVICE_ID = DEVICE_ID_LIST[0]
        os.environ["CUDA_VISIBLE_DEVICES"] = str(DEVICE_ID)
    else:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
    if not seed:
        seed = int(round(time.time() * 1000)) % (2 ** 32 - 1)
    _reset_global_seed(seed)
    print("create log path at {}".format(GlobalConfig().DEFAULT_LOG_PATH), flush=True)

    file.create_path(path=GlobalConfig().DEFAULT_LOG_PATH, del_if_existed=del_if_log_path_existed)
    Logger().init(config_or_config_dict=dict(),
                  log_path=GlobalConfig().DEFAULT_LOG_PATH,
                  log_level=GlobalConfig().DEFAULT_LOG_LEVEL)
    ConsoleLogger().init(to_file_flag=GlobalConfig().DEFAULT_WRITE_CONSOLE_LOG_TO_FILE_FLAG,
                         to_file_name=os.path.join(GlobalConfig().DEFAULT_LOG_PATH,
                                                   GlobalConfig().DEFAULT_CONSOLE_LOG_FILE_NAME),
                         level=GlobalConfig().DEFAULT_LOG_LEVEL,
                         logger_name=GlobalConfig().DEFAULT_CONSOLE_LOGGER_NAME)

    task_fn(**task_fn_kwargs)


@typechecked
def duplicate_exp_runner(num, task_fn, auto_choose_gpu_flag=False, gpu_id: int = 0, seeds: list = None,
                         del_if_log_path_existed=False,
                         **task_fn_kwargs):

    """

    :param num: the number of multiple experiments
    :type num: int
    :param task_fn: task function, defined by users
    :type task_fn: method
    :param auto_choose_gpu_flag: auto choose gpu, default False
    :type auto_choose_gpu_flag: bool
    :param gpu_id: gpu id, default 0
    :type gpu_id: int
    :param seeds: seeds generated by system time
    :type seeds: list
    :param del_if_log_path_existed: delete the existing log path, default False
    :type del_if_log_path_existed: bool
    :param task_fn_kwargs:
    :type task_fn_kwargs:
    :return:
    :rtype:
    """
    # tracemalloc.start(100)
    if seeds:
        assert len(seeds) == num
    base_log_path = deepcopy(GlobalConfig().DEFAULT_LOG_PATH)

    for i in range(num):
        GlobalConfig().set('DEFAULT_LOG_PATH', os.path.join(base_log_path, 'exp_{}'.format(i)))
        single_exp_runner(task_fn=task_fn, auto_choose_gpu_flag=auto_choose_gpu_flag,
                          del_if_log_path_existed=del_if_log_path_existed,
                          gpu_id=gpu_id, seed=seeds[i] if seeds else None, **task_fn_kwargs)
