#!/usr/bin/env python
import os
import numpy as np
from forces_simf import run_forces  # Sim func from current dir

# Import libEnsemble modules
from libensemble.libE import libE
from libensemble.manager import ManagerException
from libensemble.tools import parse_args, save_libE_output, add_unique_random_streams
from libensemble import logger
from forces_support import test_libe_stats, test_ensemble_dir, check_log_exception

USE_BALSAM = False
PERSIS_GEN = False

if PERSIS_GEN:
    from libensemble.gen_funcs.persistent_uniform_sampling import persistent_uniform as gen_f
    from libensemble.alloc_funcs.start_only_persistent import only_persistent_gens as alloc_f
else:
    from libensemble.gen_funcs.sampling import uniform_random_sample as gen_f
    from libensemble.alloc_funcs.give_sim_work_first import give_sim_work_first as alloc_f


logger.set_level('INFO')  # INFO is now default

nworkers, is_manager, libE_specs, _ = parse_args()

if is_manager:
    print('\nRunning with {} workers\n'.format(nworkers))

sim_app = os.path.join(os.getcwd(), 'forces.x')

# Normally would be pre-compiled
if not os.path.isfile('forces.x'):
    if os.path.isfile('build_forces.sh'):
        import subprocess
        subprocess.check_call(['./build_forces.sh'])

# Create executor and register sim to it.
if USE_BALSAM:
    from libensemble.executors.balsam_executor import BalsamMPIExecutor
    exctr = BalsamMPIExecutor()
else:
    from libensemble.executors.mpi_executor import MPIExecutor
    exctr = MPIExecutor()
exctr.register_app(full_path=sim_app, app_name='forces')

# Note: Attributes such as kill_rate are to control forces tests, this would not be a typical parameter.

# State the objective function, its arguments, output, and necessary parameters (and their sizes)
sim_specs = {'sim_f': run_forces,         # Function whose output is being minimized
             'in': ['x'],                 # Name of input for sim_f
             'out': [('energy', float)],  # Name, type of output from sim_f
             'user': {'keys': ['seed'],
                      'cores': 2,
                      'sim_particles': 1e3,
                      'sim_timesteps': 5,
                      'sim_kill_minutes': 10.0,
                      'particle_variance': 0.2,
                      'kill_rate': 0.5,
                      'fail_on_sim': False,
                      'fail_on_submit': False}  # Won't occur if 'fail_on_sim' True
             }
# end_sim_specs_rst_tag

# State the generating function, its arguments, output, and necessary parameters.
gen_specs = {'gen_f': gen_f,                  # Generator function
             'in': [],                        # Generator input
             'out': [('x', float, (1,))],     # Name, type and size of data produced (must match sim_specs 'in')
             'user': {'lb': np.array([0]),             # Lower bound for random sample array (1D)
                      'ub': np.array([32767]),         # Upper bound for random sample array (1D)
                      'gen_batch_size': 1000,          # How many random samples to generate in one call
                      }
             }

if PERSIS_GEN:
    alloc_specs = {'alloc_f': alloc_f}
else:
    alloc_specs = {'alloc_f': alloc_f,
                   'user': {'batch_mode': True,    # If true wait for all sims to process before generate more
                            'num_active_gens': 1}  # Only one active generator at a time
                   }

libE_specs['save_every_k_gens'] = 1000  # Save every K steps
libE_specs['sim_dirs_make'] = True      # Separate each sim into a separate directory
libE_specs['profile'] = False    # Whether to have libE profile on (default False)

# Maximum number of simulations
sim_max = 8
exit_criteria = {'sim_max': sim_max}

# Create a different random number stream for each worker and the manager
persis_info = {}
persis_info = add_unique_random_streams(persis_info, nworkers + 1)

try:
    H, persis_info, flag = libE(sim_specs, gen_specs, exit_criteria,
                                persis_info=persis_info,
                                alloc_specs=alloc_specs,
                                libE_specs=libE_specs)

except ManagerException:
    if is_manager and sim_specs['user']['fail_on_sim']:
        check_log_exception()
        test_libe_stats('Exception occurred\n')
else:
    if is_manager:
        save_libE_output(H, persis_info, __file__, nworkers)
        if sim_specs['user']['fail_on_submit']:
            test_libe_stats('Task Failed\n')
        test_ensemble_dir(libE_specs, './ensemble', nworkers, sim_max)
