# Copyright 2019 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""generate gaussian random array"""

import numpy as np
import os
import random
import logging
import sys
import time
import functools
from multiprocessing import Pool
from itertools import repeat
from akg.utils.kernel_exec import func_time_required
from akg.utils.kernel_exec import get_profiling_mode

RANDOM_SEED_NUM = 20
PROF_ERROR_CODE = 9999999999


def func(size_, miu_=0, sigma_=8, seed_=None):
    """
    Select random func according to RANDOM_FUNC_MODE and randint, calculated by the length of the random_func_list.

    Args:
        size_ (int): Size of data.
        miu_ (int): Mean value. Default: 0.
        sigma_ (int): Standard deviation. Default: 8.
        seed_ (int): seed for random.

    Returns:
        Random func, from random_func_list.
    """
    size_ = (size_ + RANDOM_SEED_NUM - 1) // RANDOM_SEED_NUM
    random_func_list = [
        np.random.RandomState(seed_).normal(miu_, sigma_, size_),
        np.random.RandomState(seed_).logistic(miu_, sigma_, size_),
        np.random.RandomState(seed_).laplace(miu_, sigma_, size_),
        np.random.RandomState(seed_).uniform(miu_, sigma_, size_),
        np.random.RandomState(seed_).tomaxint(size_),
    ]
    env_dic = os.environ
    if not env_dic.get('RANDOM_FUNC_MODE'):
        func_idx = 0
    else:
        func_idx = np.random.RandomState(None).randint(len(random_func_list))
    res = random_func_list[func_idx]
    return res


@func_time_required
def random_gaussian(size, miu=0, sigma=8, epsilon=0, seed=None):
    """Generate random array with absolution value obeys gaussian distribution."""
    random_data_disk_path = None
    if os.environ.get("RANDOM_DATA_DISK_PATH") is not None:
        random_data_disk_path = os.environ.get("RANDOM_DATA_DISK_PATH") + "/random_data_%s_%s.bin" % (str(miu), str(sigma))

    if random_data_disk_path is None or (not os.path.exists(random_data_disk_path)):
        if sigma <= 0:
            sys.stderr.write("Error: Expect positive sigmal for gaussian distribution. but get %f\n" % sigma)
            sys.exit(1)
        size_c = 1
        for x in size:
            size_c = size_c * x

        if seed is None:
            seed_ = []
            for i in range(RANDOM_SEED_NUM):
                now = int(time.time() % 10000 * 10000) + random.randint(i, 100)
                seed_.append(now)
        else:
            seed_ = [seed] * RANDOM_SEED_NUM
        logging.debug("random_gaussian seeds: {}".format(seed_))
        # In the profiling scenario, when a new process is used to run test cases, data generated by multiple processes
        # stops responding. To locate the fault, please set this parameter gen_data_multi_process to False.
        gen_data_multi_process = not bool(get_profiling_mode())
        if gen_data_multi_process:
            with Pool(processes=8) as pool:
                ret = np.array(pool.starmap(func, zip(repeat(size_c), repeat(miu), repeat(sigma), seed_)))
        else:
            numbers = list()
            for s in seed_:
                numbers.extend(func(size_c, miu, sigma, s))
            ret = np.array(numbers)
        ret = ret.flatten()
        return ret[:size_c].reshape(size) + epsilon

    data_len = functools.reduce(lambda x, y: x * y, size)
    data_pool = np.fromfile(random_data_disk_path)
    if data_len % len(data_pool) != 0:
        copy_num = (data_len // len(data_pool)) + 1
    else:
        copy_num = data_len // len(data_pool)
    data_copy = np.copy(data_pool)
    data_copy_list = []
    for _ in range(copy_num):
        np.random.shuffle(data_copy)
        data_copy_list.append(data_copy)
    data_pool = np.concatenate(tuple(data_copy_list), axis=0)
    return data_pool[0:data_len].reshape(size) + epsilon

def gen_epsilon(dtype):
    """Generate suggested epsilon according to data type."""
    return 1e-7 if dtype == np.float32 else 1e-3