"""Functions and classes for the parallel environment."""

import os

import numpy as np

from tbplas.cython import interface_mpi as interface
from tbplas.base.utils import get_datetime


__all__ = ["mpienv_init", "mpienv_finalize", "MPIEnv"]


def mpienv_init() -> None:
    """
    Initialize MPI environment. Should be called on import.
    :return: None
    """
    interface.mpienv_init()


def mpienv_finalize() -> None:
    """
    Finalize MPI environment. Should be called on exit.
    :return: None
    """
    interface.mpienv_finalize()


class MPIEnv:
    """
    Wrapper over MPI4PY APIs.

    Attributes
    ----------
    _backend: 'CppBaseMPIEnv' instance
        C++ backend for MPI environment
    _rank: integer
        id of this process in mpi communicator
    _size: integer
        number of processes in mpi communicator
    _is_master: bool
        whether this is the master process
    _is_mpi_enabled: bool
        whether mpi has been enabled during compilation
    """
    def __init__(self, echo_details: bool = False) -> None:
        """
        :param echo_details: whether to report parallelization details
        :return: None
        """
        self._backend = interface.CPPBaseMPIEnv()
        self._rank = self._backend.get_rank()
        self._size = self._backend.get_size()
        self._is_master = self._backend.is_master()
        self._is_mpi_enabled = self._backend.is_mpi_enabled()

        # Print parallelization details
        if echo_details:
            spaces = " " * 2
            self.print("\nParallelization details:")
            mpi_enabled = self._backend.is_mpi_enabled()
            if mpi_enabled:
                self.print(f"{spaces}{'MPI processes':16s} : {self._size:<6d}")
            else:
                self.print(f"{spaces}{'MPI disabled':16s}")
            for env_name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS"):
                try:
                    env_value = os.environ[env_name]
                except KeyError:
                    env_value = "n/a"
                self.print(f"{spaces}{env_name:16s} : {env_value:<6s}")
            self.print()

    def print(self, text: str = "") -> None:
        """
        Print text on master process.

        NOTE: flush=True is essential for some MPI implementations,
        e.g. MPICH3.

        :param text: text to print
        :return: None
        """
        if self.is_master:
            print(text, flush=True)

    def log(self, event: str = "", fmt: str = "%x %X") -> None:
        """
        Log the date and time of event.
        :param event: notice of the event
        :param fmt: date and time format
        :return: None.
        """
        if self.is_master:
            date_time = get_datetime(fmt=fmt)
            print(f"{event} at {date_time}", flush=True)

    def bcast(self, data_local: np.ndarray) -> None:
        """
        Broadcast data_local from master to other processes.

        NOTE: data_local should be pre-allocated before calling this function,
        otherwise segfault may be raised.

        :param data_local: local results on each process
        :return: None
        """
        buffer = data_local.reshape(data_local.size)
        if data_local.dtype == np.float64:
            self._backend.bcast_double(buffer)
        elif data_local.dtype == np.complex128:
            self._backend.bcast_complex(buffer)
        else:
            raise AssertionError(f"Illegal data type {data_local.dtype}")
        # # Another implementation based on files, slow but no segfault.
        # if self.mpi_enabled:
        #     # Master process saves data while other processes waiting
        #     if self.is_master:
        #         np.save("data_local.npy", data_local)
        #     self.barrier()
        #
        #     # Other processes load data
        #     if not self.is_master:
        #         data_local = np.load("data_local.npy")
        #     self.barrier()
        #
        #     # Master process remove data
        #     if self.is_master:
        #         os.system("rm data_local.npy")
        #     self.barrier()

    def barrier(self) -> None:
        """Put a barrier for all mpi processes."""
        self._backend.barrier()

    @property
    def mpi_enabled(self) -> bool:
        """Interface for the '_is_mpi_enabled' attribute."""
        return self._is_mpi_enabled

    @property
    def is_master(self) -> bool:
        """Interface for the '_is_master' attribute."""
        return self._is_master

    @property
    def rank(self) -> int:
        """Interface for the '_rank' attribute."""
        return self._rank

    @property
    def size(self) -> int:
        """Interface for the '_size' attribute."""
        return self._size
