"""
Runs libEnsemble with a non-persistent generator performing uniform random
sampling.

Execute via one of the following commands (e.g. 3 workers):
   mpiexec -np 4 python3 test_uniform_sampling.py
   python3 test_uniform_sampling.py --nworkers 3 --comms local
   python3 test_uniform_sampling.py --nworkers 3 --comms tcp

The number of concurrent evaluations of the objective function will be 4-1=3.
"""

# Do not change these lines - they are parsed by run-tests.sh
# TESTSUITE_COMMS: mpi local tcp
# TESTSUITE_NPROCS: 2 4

import numpy as np

# Import libEnsemble items for this test
from libensemble.libE import libE
from libensemble.sim_funcs.six_hump_camel import six_hump_camel
from libensemble.gen_funcs.sampling import uniform_random_sample
from libensemble.tools import parse_args, save_libE_output, add_unique_random_streams
from libensemble.tests.regression_tests.support import six_hump_camel_minima as minima

nworkers, is_manager, libE_specs, _ = parse_args()
libE_specs['save_every_k_sims'] = 400
libE_specs['save_every_k_gens'] = 300

sim_specs = {
    'sim_f': six_hump_camel,  # Function whose output is being minimized
    'in': ['x'],  # Keys to be given to sim_f
    'out': [('f', float)],  # Name of the outputs from sim_f
}
# end_sim_specs_rst_tag

gen_specs = {
    'gen_f': uniform_random_sample,  # Function generating sim_f input
    'out': [('x', float, (2,))],  # Tell libE gen_f output, type, size
    'user': {
        'gen_batch_size': 500,  # Used by this specific gen_f
        'lb': np.array([-3, -2]),  # Used by this specific gen_f
        'ub': np.array([3, 2]),  # Used by this specific gen_f
    },
}
# end_gen_specs_rst_tag

persis_info = add_unique_random_streams({}, nworkers + 1)

exit_criteria = {'gen_max': 501, 'elapsed_wallclock_time': 300}

# Perform the run
H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria, persis_info, libE_specs=libE_specs)

if is_manager:
    assert flag == 0

    tol = 0.1
    for m in minima:
        assert np.min(np.sum((H['x'] - m) ** 2, 1)) < tol

    print("\nlibEnsemble found the 6 minima within a tolerance " + str(tol))
    save_libE_output(H, persis_info, __file__, nworkers)
