"""
This module contains various versions that evaluate the six-hump camel function.

Six-hump camel function is documented here:
  https://www.sfu.ca/~ssurjano/camel6.html

"""
__all__ = ['six_hump_camel', 'six_hump_camel_simple',
           'six_hump_camel_with_variable_resources',
           'six_hump_camel_CUDA_variable_resources',
           'persistent_six_hump_camel']

# import subprocess
import os
import sys
import numpy as np
import time
from libensemble.executors.executor import Executor
from libensemble.message_numbers import UNSET_TAG, WORKER_DONE, TASK_FAILED
from libensemble.resources.resources import Resources
from libensemble.tools.persistent_support import PersistentSupport
from libensemble.message_numbers import STOP_TAG, PERSIS_STOP, EVAL_SIM_TAG, FINISHED_PERSISTENT_SIM_TAG


def six_hump_camel(H, persis_info, sim_specs, _):
    """
    Evaluates the six hump camel function for a collection of points given in ``H['x']``.
    Additionally evaluates the gradient if ``'grad'`` is a field in
    ``sim_specs['out']`` and pauses for ``sim_specs['user']['pause_time']]`` if
    defined.

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

    batch = len(H['x'])
    H_o = np.zeros(batch, dtype=sim_specs['out'])

    for i, x in enumerate(H['x']):
        H_o['f'][i] = six_hump_camel_func(x)

        if 'grad' in H_o.dtype.names:
            H_o['grad'][i] = six_hump_camel_grad(x)

        if 'user' in sim_specs and 'pause_time' in sim_specs['user']:
            time.sleep(sim_specs['user']['pause_time'])

    return H_o, persis_info


def six_hump_camel_simple(x, persis_info, sim_specs, _):
    """
    Evaluates the six hump camel function for a single point ``x``.

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

    H_o = np.zeros(1, dtype=sim_specs['out'])

    H_o['f'] = six_hump_camel_func(x[0][0])

    if 'pause_time' in sim_specs['user']:
        time.sleep(sim_specs['user']['pause_time'])

    return H_o, persis_info


def six_hump_camel_with_variable_resources(H, persis_info, sim_specs, libE_info):
    """
    Evaluates the six hump camel for a collection of points given in ``H['x']``
    via the executor, supporting variable sized simulations/resources, as
    determined by the generator.

    .. 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
    """

    batch = len(H['x'])
    H_o = np.zeros(batch, dtype=sim_specs['out'])
    app = sim_specs['user'].get('app', 'helloworld')
    dry_run = sim_specs['user'].get('dry_run', False)  # dry_run only prints run lines in ensemble.log
    set_cores_by_rsets = True  # If True use rset count to set num procs, else use all available to this worker.
    core_multiplier = 1  # Only used with set_cores_by_rsets as a multiplier.

    exctr = Executor.executor  # Get Executor
    task_states = []
    for i, x in enumerate(H['x']):
        # If set_cores_by_rsets is set - set nprocs here (you can oversubscribe on node)
        # else let automatic resources set nodes/procs to use all available set.
        nprocs = None  # Will be as if argument is not present
        if set_cores_by_rsets:
            resources = Resources.resources.worker_resources
            nprocs = resources.num_rsets * core_multiplier
            # print('nprocs is',nprocs,flush=True)

        inpt = None  # Will be as if argument is not present
        if app == 'six_hump_camel':
            inpt = ' '.join(map(str, H['x'][i]))

        task = exctr.submit(app_name=app, app_args=inpt,
                            num_procs=nprocs,
                            stdout='out.txt', stderr='err.txt',
                            dry_run=dry_run)
        task.wait()
        # while(not task.finished):
        #     time.sleep(0.1)
        #     task.poll()

        task_states.append(task.state)

        if app == 'six_hump_camel':
            # H_o['f'][i] = float(task.read_stdout())  # Reads whole file
            with open('out.txt') as f:
                H_o['f'][i] = float(f.readline().strip())  # Read just first line
        else:
            # To return something in test
            H_o['f'][i] = six_hump_camel_func(x)

        # v = np.random.uniform(0, 10)
        # print('About to sleep for :' + str(v))
        # time.sleep(v)

    calc_status = UNSET_TAG  # Returns to worker
    if all(t == 'FINISHED' for t in task_states):
        calc_status = WORKER_DONE
    elif any(t == 'FAILED' for t in task_states):
        calc_status = TASK_FAILED

    return H_o, persis_info, calc_status


def six_hump_camel_CUDA_variable_resources(H, persis_info, sim_specs, libE_info):
    """Launches an app setting GPU resources

    The standard test apps do not run on GPU, but demonstrates accessing resource
    information to set ``CUDA_VISIBLE_DEVICES``, and typical run configuration.
    """
    x = H['x'][0]
    H_o = np.zeros(1, dtype=sim_specs['out'])

    # Interrogate resources available to this worker
    resources = Resources.resources.worker_resources
    if resources.even_slots:  # Need same slots on each node
        resources.set_env_to_slots("CUDA_VISIBLE_DEVICES")  # Use convenience function.
        num_nodes = resources.local_node_count
        cores_per_node = resources.slot_count  # One CPU per GPU
        print('CUDA_VISIBLE_DEVICES={}  \tnodes {} ppn {}'
              .format(os.environ["CUDA_VISIBLE_DEVICES"], num_nodes, cores_per_node), flush=True)
    else:
        # Unless use a matching sub-set, but usually you probably don't want this
        print('Error: Cannot set CUDA_VISIBLE_DEVICES when uneven slots on nodes {}'.format(resources.slots))

    # Create application input file
    inpt = ' '.join(map(str, x))
    exctr = Executor.executor  # Get Executor

    # Launch application via system MPI runner, using assigned resources.
    task = exctr.submit(app_name='six_hump_camel',
                        app_args=inpt,
                        num_nodes=num_nodes,
                        procs_per_node=cores_per_node,
                        stdout='out.txt',
                        stderr='err.txt')

    task.wait()  # Wait for run to complete

    # Access app output
    with open('out.txt') as f:
        H_o['f'] = float(f.readline().strip())  # Read just first line

    calc_status = WORKER_DONE if task.state == 'FINISHED' else 'FAILED'
    return H_o, persis_info, calc_status


def persistent_six_hump_camel(H, persis_info, sim_specs, libE_info):
    """
    Similar to ``six_hump_camel``, but runs in persistent mode.
    """

    ps = PersistentSupport(libE_info, EVAL_SIM_TAG)

    # Either start with a work item to process - or just start and wait for data
    if H.size > 0:
        tag = None
        Work = None
        calc_in = H
    else:
        tag, Work, calc_in = ps.recv()

    while tag not in [STOP_TAG, PERSIS_STOP]:

        # calc_in: This should either be a function (unpack_work ?) or included/unpacked in ps.recv/ps.send_recv.
        if Work is not None:
            persis_info = Work.get('persis_info', persis_info)
            libE_info = Work.get('libE_info', libE_info)

        # Call standard six_hump_camel sim
        H_o, persis_info = six_hump_camel(calc_in, persis_info, sim_specs, libE_info)

        tag, Work, calc_in = ps.send_recv(H_o)

    final_return = None

    # Overwrite final point - for testing only
    if sim_specs['user'].get('replace_final_fields', 0):
        calc_in = np.ones(1, dtype=[('x', float, (2,))])
        H_o, persis_info = six_hump_camel(calc_in, persis_info, sim_specs, libE_info)
        final_return = H_o

    return final_return, persis_info, FINISHED_PERSISTENT_SIM_TAG


def six_hump_camel_func(x):
    """
    Definition of the six-hump camel
    """
    x1 = x[0]
    x2 = x[1]
    term1 = (4-2.1*x1**2+(x1**4)/3) * x1**2
    term2 = x1*x2
    term3 = (-4+4*x2**2) * x2**2

    return term1 + term2 + term3


def six_hump_camel_grad(x):
    """
    Definition of the six-hump camel gradient
    """

    x1 = x[0]
    x2 = x[1]
    grad = np.zeros(2)

    grad[0] = 2.0*(x1**5 - 4.2*x1**3 + 4.0*x1 + 0.5*x2)
    grad[1] = x1 + 16*x2**3 - 8*x2

    return grad


if __name__ == "__main__":
    x = (float(sys.argv[1]), float(sys.argv[2]))
    result = six_hump_camel_func(x)
    print(result)
