class MPIImportError(Exception):
    pass


class FailedMPIImport:
    """Class that can replace an mpi4py.MPI import and will throw an error if used."""

    def __init__(self, error=None, msg=None):

        self.error_obj = error

        if msg is None:
            msg = ("Failed to import MPI from mpi4py. Check your mpi4py installation if you"
                   " want to run with MPI.")
        self.msg = msg

    def __getattr__(self, attr):

        if self.error_obj is not None:
            raise MPIImportError(self.msg) from self.error_obj
        raise MPIImportError(self.msg)


def mpi_importer():
    """
    Lazy MPI import, where we only throw an error if the import failed and then we attempt to use
    the object.
    """

    try:
        from mpi4py import MPI  # pylint: disable=import-outside-toplevel
    except (ModuleNotFoundError, ImportError) as e:
        MPI = FailedMPIImport(e)

    return MPI


def mpi_numpy_decomp(MPI_N, MPI_rank, n):
    """
    Decompose a set of conditions for *MPI_N* MPI processes, where the conditions are a sequence of
    3 sequences with ordering (composition_sequence, density_sequence, temperature_sequence). This
    structure for the dataset is necessary for the vectorized reduction algorithms.
    """

    if MPI_N <= n[0]:

        comp_idx = MPI_rank
        comp_step = MPI_N
        rho_idx = T_idx = 0
        rho_step = T_step = 1

    elif MPI_N <= n[0]*n[1]:

        m = MPI_N // n[0]

        if MPI_rank <= n[0]*m:
            comp_idx = MPI_rank % n[0]
            comp_step = n[0]
            rho_idx = MPI_rank // n[0]
            rho_step = m
        else:
            comp_idx = n[0]
            comp_step = 1
            rho_idx = n[1]
            rho_step = 1

        T_idx = 0
        T_step = 1

    elif MPI_N <= n[0]*n[1]*n[2]:

        m = MPI_N // (n[0] * n[1])

        if MPI_rank <= n[0]*n[1]*m:
            comp_idx = MPI_rank % n[0]
            comp_step = n[0]
            rho_idx = (MPI_rank // n[0]) % n[1]
            rho_step = n[1]
            T_idx = MPI_rank // (n[0] * n[1])
            T_step = m
        else:
            comp_idx = n[0]
            comp_step = 1
            rho_idx = n[1]
            rho_step = 1
            T_idx = n[2]
            T_step = 1

    else:

        m = MPI_N // (n[0] * n[1])

        if MPI_rank <= n[0]*n[1]*n[2]:
            comp_idx = MPI_rank % n[0]
            comp_step = n[0]
            rho_idx = (MPI_rank // n[0]) % n[1]
            rho_step = n[1]
            T_idx = MPI_rank // (n[0] * n[1])
            T_step = n[2]
        else:
            comp_idx = n[0]
            comp_step = 1
            rho_idx = n[1]
            rho_step = 1
            T_idx = n[2]
            T_step = 1

    return comp_idx, comp_step, rho_idx, rho_step, T_idx, T_step


def to_list(x, n=1):
    """
    Convert a sequence or non-iterable to a list. In the non-iterable case, the supplied object will
    be repeated *n* times (default 1).
    """

    try:
        return list(x)
    except TypeError:
        return [x] * n
