import os
import datetime
from timeit import default_timer as timer
from typing import Callable
import cProfile
import pstats
import io
import numpy as np
import random
import inspect
import sys
from itertools import combinations

LINES = '-' * 80
CONST_COLOR_MAP = {
    'ResetAll': "\033[0m",
    'Bold': "\033[1m",
    'Dim': "\033[2m",
    'Underlined ': "\033[4m",
    'Blink      ': "\033[5m",
    'Reverse    ': "\033[7m",
    'Hidden     ': "\033[8m",

    'ResetBold': "\033[21m",
    'ResetDim': "\033[22m",
    'ResetUnderlined': "\033[24m",
    'ResetBlink': "\033[25m",
    'ResetReverse': "\033[27m",
    'ResetHidden': "\033[28m",

    'Default': "\033[39m",
    'Black': "\033[97m",
    'Red': "\033[31m",
    'Green': "\033[32m",
    'Yellow': "\033[33m",
    'Blue': "\033[34m",
    'Magenta': "\033[35m",
    'Cyan': "\033[36m",
    'LightGray': "\033[37m",
    'DarkGray': "\033[90m",
    'LightRed': "\033[91m",
    'LightGreen': "\033[92m",
    'LightYellow': "\033[93m",
    'LightBlue': "\033[94m",
    'LightMagenta': "\033[95m",
    'LightCyan': "\033[96m",
    'White': "\033[30m",

    'BackgroundDefault': "\033[49m",
    'BackgroundBlack': "\033[107m",
    'BackgroundRed': "\033[41m",
    'BackgroundGreen': "\033[42m",
    'BackgroundYellow': "\033[43m",
    'BackgroundBlue': "\033[44m",
    'BackgroundMagenta': "\033[45m",
    'BackgroundCyan': "\033[46m",
    'BackgroundLightGray': "\033[47m",
    'BackgroundDarkGray': "\033[100m",
    'BackgroundLightRed': "\033[101m",
    'BackgroundLightGreen': "\033[102m",
    'BackgroundLightYellow': "\033[103m",
    'BackgroundLightBlue': "\033[104m",
    'BackgroundLightMagenta': "\033[105m",
    'BackgroundLightCyan': "\033[106m",
    'BackgroundWhite': "\033[40m"
}


def get_timer_delta(start_timer: float) -> datetime.timedelta:
    """
    :param start_timer: begin time
    :return:
    time = get_timer()
    print('Total run time {}'.format(get_timer_delta(start_timer=time)))
    "Total run time 0:00:00.000026"
    """
    end_timer = get_timer()
    d = datetime.timedelta(seconds=(end_timer - start_timer))
    return d


def get_timer() -> float:
    return timer()


def get_current_date_hour() -> str:
    """
    running start time - good for long programs
    :return:
    """
    now = datetime.datetime.now()
    current_time = now.strftime('%d-%m-%Y %H:%M:%S')
    return current_time


def get_pc_name() -> str:
    """
    :return: pc name as str
    """
    import platform
    return platform.uname()[1]


def get_mac_address() -> str:
    """
    :return: pc mac address. e.g. 70:4D:7B:8A:65:EE
    """
    from uuid import getnode as get_mac
    mac = get_mac()
    mac = ':'.join(("%012X" % mac)[i:i + 2] for i in range(0, 12, 2))
    return mac


def get_cuda_version() -> str:
    """
    :return: cuda version if found on environment variables
    """
    if 'CUDA_PATH' in os.environ:
        cuda_v = os.path.basename(os.environ['CUDA_PATH'])
    else:
        cuda_v = 'No CUDA_PATH found'
    return cuda_v


def make_cuda_invisible() -> None:
    """
        disable(the -1) gpu 0
        TODO support hiding multiple gpus
    """
    os.environ['CUDA_VISIBLE_DEVICES'] = '-1, 0'
    return


def start_profiler() -> cProfile.Profile:
    """
    starts profiling
    :return: profiling object that is needed for end_profiler()
    """
    pr = cProfile.Profile()
    pr.enable()
    return pr


def end_profiler(pr: cProfile.Profile, rows: int = 10) -> str:
    """
    profiling output
    :param pr: object returned from start_profiler()
    :param rows: how many rows to print sorted by 'cumulative' run time
    :return:
    """
    pr.disable()
    s = io.StringIO()
    ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
    ps.print_stats(rows)
    return s.getvalue()


def set_seed(seed: int = 42) -> None:
    """
    :param seed: setting numpy and random seeds
    :return:
    """
    np.random.seed(seed)
    random.seed(seed)
    return


def get_tensorflow_version() -> str:
    try:
        import tensorflow as tf
        string = '* TensorFlow Version {}'.format(tf.__version__)
    except (ImportError, ModuleNotFoundError, NameError) as err:
        string = '* {}'.format(err)
    return string


def main_wrapper(
        main_function: Callable,
        cuda_off: bool = False,
        torch_v: bool = False,
        tf_v: bool = False,
        cv2_v: bool = False,
        with_profiler: bool = False
) -> None:
    """
    :param main_function: the function to run
    :param cuda_off: make gpu invisible and force run on cpu
    :param torch_v: print torch version
    :param tf_v: print tensorflow version
    :param cv2_v: print opencv version
    :param with_profiler: run profiler
    :return:
    """
    print(LINES)
    start_timer = get_timer()

    # make_cuda_invisible()
    print('main_wrapper:')
    print('* Run started at {}'.format(get_current_date_hour()))
    print('* Python Version {}'.format(sys.version))
    print('* Working Dir: {}'.format(os.getcwd()))
    print('* Computer Name: {}'.format(get_pc_name()))
    print('* Computer Mac: {}'.format(get_mac_address()))
    cuda_msg = '* CUDA Version: {}'.format(get_cuda_version())
    if cuda_off:
        make_cuda_invisible()
        cuda_msg += ' (Turned off)'
    print(cuda_msg)

    if torch_v:
        from wizzi_utils.torch_tools import get_torch_version
        print(get_torch_version())
    if tf_v:
        print(get_tensorflow_version())
    if cv2_v:
        from wizzi_utils.open_cv_tools import get_cv_version
        print(get_cv_version())

    print('Function {} started:'.format(main_function))
    print(LINES)
    pr = start_profiler() if with_profiler else None
    main_function()
    if with_profiler:
        print(end_profiler(pr))
    print(LINES)
    print('Total run time {}'.format(get_timer_delta(start_timer)))
    return


def add_data(var_l, data_chars_l: int) -> str:
    """
    Aux function for to_str()
    :param var_l:
    :param data_chars_l:
    :return:
    """
    data_str_raw = str(var_l).replace('\n', '').replace('  ', '')
    data_str = ''
    if data_chars_l < 0:  # all data
        data_str = ': {}'.format(data_str_raw)
    elif data_chars_l > 0:  # first 'data_chars' characters
        data_str = ': {}'.format(data_str_raw[:data_chars_l])
        if len(data_str_raw) > data_chars_l:
            data_str += ' ...too long'
    return data_str


def to_str(var, title: str, data_chars: int = 100) -> str:
    """
    :param var: the variable
    :param title: the title (usually variable name)
    :param data_chars: how many char to print.
        -1: all
         0: none
        +0: maximum 'data_chars' (e.g. data_chars=50 and |str(var)|=100 - first 50 chars)

    examples:
        print(to_str(var=3, title='my_int'))
        print(to_str(var=3.2, title='my_float'))
        print(to_str(var='a', title='my_str'))
        print(to_str(var=[], title='my_empty_list'))
        print(to_str(var=[1, 3, 4], title='1d list of ints'))
        print(to_str(var=[1, 3], title='1d list of ints no data', data_chars=0))  # no data
        print(to_str(var=[15] * 1000, title='1d long list'))
        print(to_str(var=(19, 3, 9), title='1d tuple'))
        print(to_str(var=[[11.2, 15.9], [3.0, 7.55]], title='2d list'))
        print(to_str(var=[(11.2, 15.9), (3.0, 7.55)], title='2d list of tuples'))
        b = np.array([[11.2, 15.9], [3.0, 7.55]])
        print(to_str(var=b, title='2d np array'))
        cv_img = np.zeros(shape=[480, 640, 3], dtype=np.uint8)
        print(to_str(var=cv_img, title='cv_img'))
        print(to_str(var={'a': [1, 2]}, title='dict of lists'))
        print(to_str(var={'a': [{'k': [1, 2]}, {'c': [7, 2]}]}, title='nested dict'))
    :return: informative string of the variable
    """

    type_s = str(type(var)).replace('<class \'', '').replace('\'>', '')  # clean type name
    string = '{}({})'.format(title, type_s)  # base: title and type

    if isinstance(var, (int, float, str)):
        if hasattr(var, "__len__"):
            string = string.replace(')', ',len={}'.format(var.__len__()))
        string += add_data(var, data_chars)

    elif isinstance(var, (list, tuple)):
        string = string.replace(')', ',len={})'.format(var.__len__()))
        string += add_data(var, data_chars)
        if len(var) > 0:  # recursive call
            string += '\n\t{}'.format(to_str(var=var[0], title='{}[0]'.format(title), data_chars=data_chars))

    elif isinstance(var, np.ndarray):
        string = string.replace(')', ',shape={},dtype={})'.format(var.shape, var.dtype))
        string += add_data(var.tolist(), data_chars)
        if len(var) > 0:  # recursive call
            string += '\n\t{}'.format(to_str(var=var[0], title='{}[0]'.format(title), data_chars=data_chars))

    elif isinstance(var, dict):
        string = string.replace(')', ',len={},keys={})'.format(var.__len__(), list(var.keys())))
        string += add_data(var, data_chars)
        if len(var) > 0:  # recursive call
            first_key = next(iter(var))
            string += '\n\t{}'.format(
                to_str(var=var[first_key], title='{}[{}]'.format(title, first_key), data_chars=data_chars))

    else:  # all unidentified elements get default print (title(type): data)
        string += add_data(var, data_chars)
    return string


def save_np(t: np.array, path: str, ack_print: bool = True, tabs: int = 0):
    """
    to save a dict - see np.savez(path, name1=a, name2=b ... ) - loading the same
    :param t: numpy array
    :param path: suffix '.npy' added automatically
    :param ack_print:
    :param tabs:
    :return:
    """
    np.save(path, t)
    if ack_print:
        print('{}Saved to: {}'.format(tabs * '\t', path))
    return


def load_np(path: str, ack_print: bool = True, tabs: int = 0) -> np.array:
    """
    :param path:
    :param ack_print:
    :param tabs:
    :return: numpy array
    """
    t = np.load(path)
    if ack_print:
        print('{}Loaded: {}'.format(tabs * '\t', path))
    return t


def get_uniform_dist_by_dim(A):
    lows = np.min(A, axis=0)
    highs = np.max(A, axis=0)
    return lows, highs


def get_normal_dist_by_dim(A):
    means = np.mean(A, axis=0)
    stds = np.std(A, axis=0)
    return means, stds


def np_normal(shape: tuple, mius: list, stds: list) -> np.array:
    """
    e.g.
        d = 2
        A = np.zeros((3, d))
        A[0][0], A[0][1] = 0, 1000
        A[1][0], A[1][1] = 1, 2000
        A[2][0], A[2][1] = 5, 4000
        means, stds = get_normal_dist_by_dim(A)
        print(means, stds)
        A_2 = np_normal(shape=(500, d), mius=means, stds=stds)
        print(A_2.shape)
        print(get_normal_dist_by_dim(A_2))
    """
    ret = np.random.normal(loc=mius, scale=stds, size=shape)
    return ret


def np_uniform(shape: tuple, lows: list, highs: list) -> np.array:
    """
    e.g.
        d=2
        A = np.zeros((3, d))
        A[0][0], A[0][1] = 0, 1000
        A[1][0], A[1][1] = 1, 2000
        A[2][0], A[2][1] = 5, 4000
        lows, highs = get_uniform_dist_by_dim(A)
        print(lows, highs)
        A_2 = np_uniform(shape=(10, d), lows=lows, highs=highs)
        print(A_2.shape)
        print(get_normal_dist_by_dim(A_2))
    """
    ret = np.random.uniform(low=lows, high=highs, size=shape)
    return ret


def augment_x_y_numpy(X: np.array, y: np.array) -> np.array:
    """ creates A=X|y """
    assert X.shape[0] == y.shape[0], 'row count must be the same'
    if len(X.shape) == 1:  # change x size()=[n,] to size()=[n,1]
        X = X.reshape(X.shape[0], 1)
    if len(y.shape) == 1:  # change y size()=[n,] to size()=[n,1]
        y = y.reshape(y.shape[0], 1)
    A = np.column_stack((X, y))
    return A


def de_augment_numpy(A: np.array) -> (np.array, np.array):
    """ creates X|y=A """
    if len(A.shape) == 1:  # A is 1 point. change from size (n) to size (1,n)
        A = A.reshape(1, A.shape[0])
    X, y = A[:, :-1], A[:, -1]
    if len(X.shape) == 1:  # change x size()=[n,] to size()=[n,1]
        X = X.reshape(X.shape[0], 1)
    if len(y.shape) == 1:  # change y size()=[n,] to size()=[n,1]
        y = y.reshape(y.shape[0], 1)
    return X, y


def nCk(n: int, k: int, as_int: bool = False):
    """
    n choose k
    :return: if as_int True: the result of nCk, else the combinations of nCk

    e.g.
    A = np.random.randint(low=-10, high=10, size=(3, 2))
    print('A={}'.format(A.tolist()))

    # let's iterate on every 2 different indices of A
    combs_count = utils.nCk(len(A), k=2, as_int=True)
    print('|A| choose 2={}:'.format(combs_count)) # result is 3

    combs_list = utils.nCk(len(A), k=2)  # result is [[0, 1], [0, 2], [1, 2]]
    for i, comb in enumerate(combs_list):
        print('\tcomb {}={}. A[comb]={}'.format(i, comb, A[comb].tolist()))
    """
    range_list = np.arange(0, n, 1)
    combs = list(combinations(range_list, k))
    combs = [list(comb) for comb in combs]
    if as_int:
        combs = len(combs)  #
    return combs


def redirect_std_start():
    """
    redirect all prints to summary_str
    e.g.
        old_stdout, summary_str = redirect_std_start()
        print('blablabla')
        string = redirect_std_finish(old_stdout, summary_str)
        print('captured output: {}'.format(string))
    :return:
    """
    old_stdout = sys.stdout
    sys.stdout = summary_str = io.StringIO()
    return old_stdout, summary_str


def redirect_std_finish(old_stdout, summary_str) -> str:
    """ redirect all prints back to std out and return a string of what was captured"""
    sys.stdout = old_stdout
    return summary_str.getvalue()


def get_line_number(depth: int = 1) -> str:
    ret_val = ''
    try:
        scope_1_back = inspect.stack()[depth]  # stack()[0] is this function
        ret_val = '{}'.format(scope_1_back.lineno)
    except IndexError:
        pass
    return ret_val


def get_function_name(depth: int = 1) -> str:
    ret_val = ''
    try:
        scope_1_back = inspect.stack()[depth]  # stack()[0] is this function
        ret_val = '{}'.format(scope_1_back.function)
    except IndexError:
        pass
    return ret_val


def get_file_name(depth: int = 1) -> str:
    ret_val = ''
    try:
        scope_1_back = inspect.stack()[depth]  # stack()[0] is this function
        ret_val = '{}'.format(scope_1_back.filename)
    except IndexError:
        pass
    return ret_val


def get_function_name_and_line(depth: int = 1) -> str:
    # +1 because of this function
    ret_val = '{} line {}'.format(get_function_name(depth + 1), get_line_number(depth + 1))
    return ret_val


def add_color(string: str, color: str = 'Red') -> str:
    """
    :param string:
    :param color: from color map
    # todo add option for combination (like bold, underlined ...)
    # todo change to lower case
    # todo color numbers maybe wrong
    e.g.
        my_str = 'hello colorful world!'
        print(my_str)
        print(add_color(my_str, color='Red'))
        print(add_color(my_str, color='Blue'))
        print(add_color(my_str, color='Bold'))
    to see all colors:
        for k, v in CONST_COLOR_MAP.items():
            print('{}{}{}'.format(v, k, CONST_COLOR_MAP['ResetAll']))
    """

    if color in CONST_COLOR_MAP:
        # concat (color, string, reset tag)
        string = '{}{}{}'.format(CONST_COLOR_MAP[color], string, CONST_COLOR_MAP['ResetAll'])
    return string


def init_logger(logger_path: str = None):
    """
    example:
        init_logger('./log.txt')
        log_print('line 1')
        flush_logger()
        log_print('line 2')
        log_print('line 3')
        close_logger()
    """
    global logger
    logger = open(file=logger_path, mode='w', encoding='utf-8') if logger_path is not None else None
    print(type(logger))
    return logger


def flush_logger() -> None:
    """ good for loops - writes every iteration if used """
    global logger
    if logger is not None:
        logger.flush()
    return


def log_print(string: str) -> None:
    global logger
    print(string)
    if logger is not None:
        logger.write('{}\n'.format(string))
    return


def log_print_dict(my_dict) -> None:
    for k, v in my_dict.items():
        log_print('\t{}: {}'.format(k, v))
    return


def close_logger() -> None:
    global logger
    if logger is not None:
        logger.close()
    return


def create_dir_if_not_exists(data_root: str, ack: bool = True, tabs: int = 1):
    if os.path.exists(data_root):
        if ack:
            print('{}{} exists'.format(tabs * '\t', data_root))
    else:
        os.mkdir(data_root)
        if ack:
            print('{}{} created'.format(tabs * '\t', data_root))
    return


def main():
    return


if __name__ == '__main__':
    main_wrapper(
        main_function=main,
        cuda_off=True,
        torch_v=True,
        tf_v=False,
        cv2_v=True,
        with_profiler=False
    )
