class PipelineMPI:

    """PipelineMPI class, it orchestrates the instantiation of MPI jobs
    and distributes the workload among the workers.


    """

    def __init__(
        self,
        exec: callable = None,
        extra_params: dict = None,
        collect: bool = None,
        show_log: bool = True,
    ) -> None:
        self.exec = exec
        self.show_log = show_log

        if extra_params is not None:
            self.extra_params = extra_params
        else:
            self.extra_params = {}

        self.collect = collect

        self.comm = MPI.COMM_WORLD
        self.n_procs = self.comm.Get_size()

        self.status = (self.n_procs - 1) * [False]
        self.status_dict = dict()

    def _check_kwargs_consistency(self, kwargs: dict = None) -> int:
        """It checks if the kwargs provided for each worker
        have the same length.

        Args:
            kwargs (dict, optional): a dictionary containing the kwargs of all (Default value = None)

        Returns:
            int: Length of the batch sent for each worker.

        """

        types = [type(value) for value in kwargs.values()]
        lengths = [len(value) for value in kwargs.values()]

        assert all([t == list for t in types]), (
            f"All the elements in kwargs must be list," f" but received {types}."
        )

        assert len(set(lengths)) == 1, (
            f"All the elements in kwargs must be the same length,"
            f" but received {lengths}"
        )

        print("kwargs is alright.")

        return lengths[0]

    def _split_kwargs(
        self, kwargs: dict, rank: int, size: int, total_size: int
    ) -> Tuple[dict, int]:
        """It allows the workload be executed serially in each worker node

        Args:
            kwargs (dict): A dictionary containing kwargs, which will be distributed for all the workers.
            rank (int): The index of the rank.
            size (int): The number of available workers.
            total_size (int): The total number of elements to be distributed among the workers.

        Returns:
            kwargs_batch: A dictionary containing the kwargs to be sent for each worker.
            batch_size: The batch size, which corresponds to the number of elements
            to be sent for each worker.

        """

        # Decrement rank and size by 1, because they are usually 0-indexed in Python
        size -= 1
        rank -= 1

        # Calculate batch size and remainder using divmod() function
        batch_size, remainder = divmod(total_size, size)

        # If rank is less than remainder, calculate kwargs_batch using batch size + 1
        if rank < remainder:
            kwargs_batch = {
                key: value[rank * (batch_size + 1) : (rank + 1) * (batch_size + 1)]
                for key, value in kwargs.items()
            }
            return kwargs_batch, batch_size + 1
        # If rank is not less than remainder, calculate kwargs_batch using batch size
        else:
            kwargs_batch = {
                key: value[
                    remainder * (batch_size + 1)
                    + (rank - remainder)
                    * batch_size : (rank - remainder + 1)
                    * batch_size
                ]
                for key, value in kwargs.items()
            }

        return kwargs_batch, batch_size

    def _attribute_dict_output(self, dicts: list = None) -> None:
        root = dict()
        for e in dicts:
            root.update(e)

        for key, value in root.items():
            self.status_dict[key] = value

    @staticmethod
    def inner_type(obj: list = None):
        types_list = [type(o) for o in obj]
        assert len(set(types_list)) == 1, "Composed types are not supported."

        return types_list[0]

    def _exec_wrapper(self, kwargs: dict, total_size: int) -> None:
        """A wrapper method around exec to facilitate the
        instantiation of each worker.

        Args:
            kwargs (dict): A dictionary containing kwargs for the worker.
            total_size (int): The total number of elements.

        """

        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()
        size = comm.Get_size()
        size_ = size

        # Rank 0 is the 'master' node
        # The worker nodes execute their workload and send a message to
        # master

        if rank != 0:
            print(f"Executing rank {rank}.")
            kwargs_batch, batch_size = self._split_kwargs(
                kwargs, rank, size_, total_size
            )

            kwargs_batch_list = [
                {key: value[j] for key, value in kwargs_batch.items()}
                for j in range(batch_size)
            ]

            out = list()
            for i in kwargs_batch_list:
                print(f"Executing batch {i['key']} in rank {rank}")
                # Concatenate the rank to the extra parameters
                i.update(self.extra_params)
                # Appending the result of the operation self.exec to the partial list
                out.append(self.exec(**i))

            if self.collect is True:
                msg = out
            else:
                msg = 1

            if self.show_log:
                print(f"Sending the output {msg} to rank 0")

            comm.send(msg, dest=0)

            print(f"Execution concluded for rank {rank}.")

        # The master awaits the responses of each worker node
        elif rank == 0:
            for r in range(1, size):
                msg = comm.recv(source=r)
                self.status[r - 1] = msg

                if self.inner_type(msg) == dict:
                    self._attribute_dict_output(dicts=msg)

                if self.show_log:
                    print(f"Rank 0 received {msg} from rank {r}")

        comm.barrier()

    @property
    def success(self) -> bool:
        """It returns True if the entire process worked without issues.

        """

        return all(self.status)

    def run(self, kwargs: dict = None) -> None:
        """It runs the MPI job

        Args:
            kwargs (dict, optional): A kwargs dictionary containing chunks of input arguments to be sent for each worker. (Default value = None)

        """

        comm = MPI.COMM_WORLD
        rank = comm.Get_rank()

        total_size = 0

        # Checking if the datasets dimensions are in accordance with the expected ones
        if rank == 0:
            total_size = self._check_kwargs_consistency(kwargs=kwargs)

        total_size = comm.bcast(total_size, root=0)

        comm.barrier()

        # Executing a wrapper containing the parallelized operation
        self._exec_wrapper(kwargs, total_size)

        comm.barrier()