# This file is part of the pyMOR project (https://www.pymor.org).
# Copyright 2013-2021 pyMOR developers and contributors. All rights reserved.
# License: BSD 2-Clause License (https://opensource.org/licenses/BSD-2-Clause)

from itertools import chain
import os


from pymor.parallel.basic import WorkerPoolBase
from pymor.tools import mpi


class MPIPool(WorkerPoolBase):
    """|WorkerPool| based pyMOR's MPI :mod:`event loop <pymor.tools.mpi>`."""

    def __init__(self):
        super().__init__()
        self.logger.info(f'Connected to {mpi.size} ranks')
        self._payload = mpi.call(mpi.function_call_manage, _setup_worker)
        self._apply(os.chdir, os.getcwd())

    def __del__(self):
        mpi.call(mpi.remove_object, self._payload)

    def __len__(self):
        return mpi.size

    def _push_object(self, obj):
        return mpi.call(mpi.function_call_manage, _push_object, obj)

    def _apply(self, function, *args, **kwargs):
        return mpi.call(mpi.function_call, _worker_call_function, function, *args, **kwargs)

    def _apply_only(self, function, worker, *args, **kwargs):
        payload = mpi.get_object(self._payload)
        payload[0] = (function, args, kwargs)
        try:
            result = mpi.call(mpi.function_call, _single_worker_call_function, self._payload, worker)
        finally:
            payload[0] = None
        return result

    def _map(self, function, chunks, **kwargs):
        payload = mpi.get_object(self._payload)
        payload[0] = chunks
        try:
            result = mpi.call(mpi.function_call, _worker_map_function, self._payload, function, **kwargs)
        finally:
            payload[0] = None
        return result

    def _remove_object(self, remote_id):
        mpi.call(mpi.remove_object, remote_id)


def _worker_call_function(function, *args, **kwargs):
    result = function(*args, **kwargs)
    return mpi.comm.gather(result, root=0)


def _single_worker_call_function(payload, worker):
    if mpi.rank0:
        if worker == 0:
            function, args, kwargs = payload[0]
            return mpi.function_call(function, *args, **kwargs)
        else:
            mpi.comm.send(payload[0], dest=worker)
            return mpi.comm.recv(source=worker)
    else:
        if mpi.rank != worker:
            return
        (function, args, kwargs) = mpi.comm.recv(source=0)
        retval = mpi.function_call(function, *args, **kwargs)
        mpi.comm.send(retval, dest=0)


def _worker_map_function(payload, function, **kwargs):

    if mpi.rank0:
        args = list(zip(*payload[0]))
    else:
        args = None
    args = zip(*mpi.comm.scatter(args, root=0))

    result = [mpi.function_call(function, *a, **kwargs) for a in args]
    result = mpi.comm.gather(result, root=0)

    if mpi.rank0:
        return list(chain(*result))


def _setup_worker():
    # ensure that each worker starts with a different RandomState
    if not mpi.rank0:
        from pymor.tools import random
        import numpy as np
        state = random.default_random_state()
        new_state = np.random.RandomState(state.randint(0, 2**16) + mpi.rank)
        random._default_random_state = new_state
    return [None]


def _push_object(obj):
    return obj
