"""
Runs libEnsemble with APOSMM with the NLopt local optimizer.

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

When running with the above commands, the number of concurrent evaluations of
the objective function will be 2, as one of the three workers will be the
persistent generator.
"""

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

import sys
import numpy as np

# Import libEnsemble items for this test
from libensemble.libE import libE
from math import gamma, pi, sqrt
from libensemble.sim_funcs.six_hump_camel import six_hump_camel as sim_f

import libensemble.gen_funcs

libensemble.gen_funcs.rc.aposmm_optimizers = 'nlopt'
from libensemble.gen_funcs.persistent_aposmm import aposmm as gen_f

from libensemble.alloc_funcs.persistent_aposmm_alloc import persistent_aposmm_alloc as alloc_f
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
from time import time

nworkers, is_manager, libE_specs, _ = parse_args()

if is_manager:
    start_time = time()

if nworkers < 2:
    sys.exit("Cannot run with a persistent worker if only one worker -- aborting...")

n = 2
sim_specs = {
    'sim_f': sim_f,
    'in': ['x'],
    'out': [('f', float)],
}

gen_out = [
    ('x', float, n),
    ('x_on_cube', float, n),
    ('sim_id', int),
    ('local_min', bool),
    ('local_pt', bool),
]

gen_specs = {
    'gen_f': gen_f,
    'persis_in': ['f'] + [n[0] for n in gen_out],
    'out': gen_out,
    'user': {
        'initial_sample_size': 100,
        'sample_points': np.round(minima, 1),
        'localopt_method': 'LN_BOBYQA',
        'rk_const': 0.5 * ((gamma(1 + (n / 2)) * 5) ** (1 / n)) / sqrt(pi),
        'xtol_abs': 1e-6,
        'ftol_abs': 1e-6,
        'dist_to_bound_multiple': 0.5,
        'max_active_runs': 6,
        'lb': np.array([-3, -2]),
        'ub': np.array([3, 2]),
    },
}

alloc_specs = {'alloc_f': alloc_f}

persis_info = add_unique_random_streams({}, nworkers + 1)

exit_criteria = {'sim_max': 1000}

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

if is_manager:
    print('[Manager]:', H[np.where(H['local_min'])]['x'])
    print('[Manager]: Time taken =', time() - start_time, flush=True)

    tol = 1e-5
    for m in minima:
        # The minima are known on this test problem.
        # We use their values to test APOSMM has identified all minima
        print(np.min(np.sum((H[H['local_min']]['x'] - m) ** 2, 1)), flush=True)
        assert np.min(np.sum((H[H['local_min']]['x'] - m) ** 2, 1)) < tol

    save_libE_output(H, persis_info, __file__, nworkers)
