# Copyright 2020, 2021 Francesco Biscani (bluescarni@gmail.com), Dario Izzo (dario.izzo@gmail.com)
#
# This file is part of the heyoka.py library.
#
# This Source Code Form is subject to the terms of the Mozilla
# Public License v. 2.0. If a copy of the MPL was not distributed
# with this file, You can obtain one at http://mozilla.org/MPL/2.0/.

# NOTE: this is a small helper to splat a 1D grid into the
# appropriate shape for a batch integrator. If ta is a scalar integrator,
# the original grid will be returned unchanged.
def _splat_grid(arg, ta):
    if hasattr(ta, "batch_size"):
        import numpy as np

        return np.repeat(arg, ta.batch_size).reshape((-1, ta.batch_size))
    else:
        return arg


# Thread-based implementation.
def _ensemble_propagate_thread(tp, ta, arg, n_iter, gen, **kwargs):
    from concurrent.futures import ThreadPoolExecutor
    from copy import deepcopy

    # Pop the multithreading options from kwargs.
    max_workers = kwargs.pop("max_workers", None)

    # The worker function.
    def func(i):
        # Create the local integrator.
        local_ta = gen(deepcopy(ta), i)

        # Run the propagation.
        if tp == "until":
            loc_ret = local_ta.propagate_until(arg, **kwargs)
        elif tp == "for":
            loc_ret = local_ta.propagate_for(arg, **kwargs)
        else:
            loc_ret = local_ta.propagate_grid(_splat_grid(arg, ta), **kwargs)

        # Return the results.
        # NOTE: in batch mode, loc_ret will be single
        # value rather than a tuple, hence the branch.
        if isinstance(loc_ret, tuple):
            return (local_ta, ) + loc_ret
        else:
            return (local_ta, loc_ret)

    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        ret = list(executor.map(func, range(n_iter)))

    return ret


# The worker function used in the multiprocessing implementation.
def _mp_propagate(tup):
    from . import _s11n_backend_map

    tp, ta, gen, arg, kwargs, i, s11n_str = tup

    # Fetch the s11n backend from its
    # str representation.
    s11n_be = _s11n_backend_map[s11n_str]

    # Unpickle the other arguments.
    ta = s11n_be.loads(ta)
    gen = s11n_be.loads(gen)
    kwargs = s11n_be.loads(kwargs)

    # Create the local integrator.
    local_ta = gen(ta, i)

    # Run the propagation.
    if tp == "until":
        loc_ret = local_ta.propagate_until(arg, **kwargs)
    elif tp == "for":
        loc_ret = local_ta.propagate_for(arg, **kwargs)
    else:
        loc_ret = local_ta.propagate_grid(_splat_grid(arg, ta), **kwargs)

    # Return the results.
    # NOTE: in batch mode, loc_ret will be single
    # value rather than a tuple, hence the branch.
    if isinstance(loc_ret, tuple):
        return s11n_be.dumps((local_ta, ) + loc_ret)
    else:
        return s11n_be.dumps((local_ta, loc_ret))


# Process-based implementation.
def _ensemble_propagate_process(tp, ta, arg, n_iter, gen, **kwargs):
    from concurrent.futures import ProcessPoolExecutor
    import multiprocessing as mp
    from . import get_serialization_backend, _s11n_backend_inv_map

    # Fetch the currently active s11n backend.
    s11n_be = get_serialization_backend()

    # Fetch its string counterpart.
    s11n_str = _s11n_backend_inv_map[s11n_be]

    # NOTE: ensure the processes are started with
    # the 'spawn' method.
    ctx = mp.get_context("spawn")

    # Pop the multiprocessing options from kwargs.
    max_workers = kwargs.pop("max_workers", None)
    chunksize = kwargs.pop("chunksize", 1)

    with ProcessPoolExecutor(max_workers=max_workers, mp_context=ctx) as executor:
        ret = list(executor.map(
            _mp_propagate, zip([tp]*n_iter, [s11n_be.dumps(ta)]*n_iter, [s11n_be.dumps(gen)]*n_iter,
                               [arg]*n_iter, [s11n_be.dumps(kwargs)] *
                               n_iter, range(n_iter),
                               [s11n_str]*n_iter), chunksize=chunksize))

    return [s11n_be.loads(_) for _ in ret]
