"""
This module contains multiple generation functions for sampling a domain. All
use (and return) a random stream in ``persis_info``, given by the allocation
function.
"""
import numpy as np

__all__ = ['uniform_random_sample',
           'uniform_random_sample_with_variable_resources',
           'uniform_random_sample_obj_components',
           'latin_hypercube_sample',
           'uniform_random_sample_cancel']


def uniform_random_sample(H, persis_info, gen_specs, _):
    """
    Generates ``gen_specs['user']['gen_batch_size']`` points uniformly over the domain
    defined by ``gen_specs['user']['ub']`` and ``gen_specs['user']['lb']``.

    .. seealso::
        `test_uniform_sampling.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/regression_tests/test_uniform_sampling.py>`_ # noqa
    """
    ub = gen_specs['user']['ub']
    lb = gen_specs['user']['lb']

    n = len(lb)
    b = gen_specs['user']['gen_batch_size']

    H_o = np.zeros(b, dtype=gen_specs['out'])

    H_o['x'] = persis_info['rand_stream'].uniform(lb, ub, (b, n))

    return H_o, persis_info


def uniform_random_sample_with_variable_resources(H, persis_info, gen_specs, _):
    """
    Generates points uniformly over the domain defined by ``gen_specs['user']['ub']`` and
    ``gen_specs['user']['lb']``. Also randomly requests a different number of resource
    sets to be used in the evaluation of the generated points after the initial batch.

    .. seealso::
        `test_uniform_sampling_with_variable_resources.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/regression_tests/test_uniform_sampling_with_variable_resources.py>`_ # noqa
    """
    ub = gen_specs['user']['ub']
    lb = gen_specs['user']['lb']
    n = len(lb)

    if len(H) == 0:
        b = gen_specs['user']['initial_batch_size']

        H_o = np.zeros(b, dtype=gen_specs['out'])
        for i in range(0, b):
            # x= i*np.ones(n)
            x = persis_info['rand_stream'].uniform(lb, ub, (1, n))
            H_o['x'][i] = x
            H_o['resource_sets'][i] = 1
            H_o['priority'] = 1

    else:
        H_o = np.zeros(1, dtype=gen_specs['out'])
        # H_o['x'] = len(H)*np.ones(n)  # Can use a simple count for testing.
        H_o['x'] = persis_info['rand_stream'].uniform(lb, ub)
        H_o['resource_sets'] = persis_info['rand_stream'].randint(1, gen_specs['user']['max_resource_sets']+1)
        H_o['priority'] = 10*H_o['resource_sets']
        # print('Created sim for {} workers'.format(H_o['resource_sets']), flush=True)

    return H_o, persis_info


def uniform_random_sample_obj_components(H, persis_info, gen_specs, _):
    """
    Generates points uniformly over the domain defined by ``gen_specs['user']['ub']``
    and ``gen_specs['user']['lb']`` but requests each ``obj_component`` be evaluated
    separately.

    .. seealso::
        `test_uniform_sampling_one_residual_at_a_time.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/regression_tests/test_uniform_sampling_one_residual_at_a_time.py>`_ # noqa
    """
    ub = gen_specs['user']['ub']
    lb = gen_specs['user']['lb']

    n = len(lb)
    m = gen_specs['user']['components']
    b = gen_specs['user']['gen_batch_size']

    H_o = np.zeros(b*m, dtype=gen_specs['out'])
    for i in range(0, b):
        x = persis_info['rand_stream'].uniform(lb, ub, (1, n))
        H_o['x'][i*m:(i+1)*m, :] = np.tile(x, (m, 1))
        H_o['priority'][i*m:(i+1)*m] = persis_info['rand_stream'].uniform(0, 1, m)
        H_o['obj_component'][i*m:(i+1)*m] = np.arange(0, m)

        H_o['pt_id'][i*m:(i+1)*m] = len(H)//m+i

    return H_o, persis_info


def uniform_random_sample_cancel(H, persis_info, gen_specs, _):
    """
    Similar to uniform_random_sample but with immediate cancellation of
    selected points for testing.

    """
    ub = gen_specs['user']['ub']
    lb = gen_specs['user']['lb']

    n = len(lb)
    b = gen_specs['user']['gen_batch_size']

    H_o = np.zeros(b, dtype=gen_specs['out'])
    for i in range(b):
        if i % 10 == 0:
            H_o[i]['cancel_requested'] = True

    H_o['x'] = persis_info['rand_stream'].uniform(lb, ub, (b, n))

    return H_o, persis_info


def latin_hypercube_sample(H, persis_info, gen_specs, _):
    """
    Generates ``gen_specs['user']['gen_batch_size']`` points in a Latin
    hypercube sample over the domain defined by ``gen_specs['user']['ub']`` and
    ``gen_specs['user']['lb']``.

    .. seealso::
        `test_1d_sampling.py <https://github.com/Libensemble/libensemble/blob/develop/libensemble/tests/regression_tests/test_1d_sampling.py>`_ # noqa
    """

    ub = gen_specs['user']['ub']
    lb = gen_specs['user']['lb']

    n = len(lb)
    b = gen_specs['user']['gen_batch_size']

    H_o = np.zeros(b, dtype=gen_specs['out'])

    A = lhs_sample(n, b, persis_info['rand_stream'])

    H_o['x'] = A*(ub-lb)+lb

    return H_o, persis_info


def lhs_sample(n, k, stream):

    # Generate the intervals and random values
    intervals = np.linspace(0, 1, k+1)
    rand_source = stream.uniform(0, 1, (k, n))
    rand_pts = np.zeros((k, n))
    sample = np.zeros((k, n))

    # Add a point uniformly in each interval
    a = intervals[:k]
    b = intervals[1:]
    for j in range(n):
        rand_pts[:, j] = rand_source[:, j]*(b-a) + a

    # Randomly perturb
    for j in range(n):
        sample[:, j] = rand_pts[stream.permutation(k), j]

    return sample
