"""
Define various types of solvers.

Example:
    To initialize one of the solver classes:

    >>> from llg3d.parameters import RunParameters
    >>> from llg3d.solvers.numpy import NumpySolver
    >>> solver = NumpySolver(**RunParameters(solver="numpy").as_dict())
"""

from typing import TYPE_CHECKING

import os
import importlib.util

if TYPE_CHECKING:
    from .base import BaseSolver


def get_size() -> int:
    """
    Return the number of parallel MPI processes.

    Use environment variables to avoid initializing MPI unnecessarily.

    Returns:
        Number of MPI processes if in an MPI environment, else 1.
    """
    # Open MPI
    if "OMPI_COMM_WORLD_SIZE" in os.environ:
        return int(os.environ["OMPI_COMM_WORLD_SIZE"])

    # MPICH/Intel MPI
    if "PMI_SIZE" in os.environ:
        return int(os.environ["PMI_SIZE"])

    # SLURM
    if "SLURM_NTASKS" in os.environ:
        return int(os.environ["SLURM_NTASKS"])

    return 1


def get_rank() -> int:
    """
    Return the rank of the current MPI process.

    Use environment variables to avoid initializing MPI unnecessarily.

    Returns:
        Rank of the current MPI process if in an MPI environment, else 0.
    """
    # PMIx
    if "PMIX_RANK" in os.environ:
        return int(os.environ["PMIX_RANK"])

    # Open MPI
    if "OMPI_COMM_WORLD_RANK" in os.environ:
        return int(os.environ["OMPI_COMM_WORLD_RANK"])

    # MPICH/Intel MPI
    if "PMI_RANK" in os.environ:
        return int(os.environ["PMI_RANK"])

    # SLURM
    if "SLURM_PROCID" in os.environ:
        return int(os.environ["SLURM_PROCID"])

    return 0


__all__ = [
    "rank",
    "size",
    "comm",
    "status",
    "mpi_initialized",
    "get_size",
    "get_rank",
    "get_solver_class",
    "LIB_AVAILABLE",
]

LIB_AVAILABLE: dict[str, bool] = {}

# Check for other solver availability
for lib in "pyopencl", "jax", "mpi4py":
    LIB_AVAILABLE[lib] = importlib.util.find_spec(lib, package=__package__) is not None


# MPI library: initialize dummy variables at first:
# it prevents from initializing the MPI communicator if not needed
class _DummyComm:
    pass


class _DummyStatus:
    pass


comm = _DummyComm()
rank = get_rank()
size = get_size()
status = _DummyStatus()
mpi_initialized = False


def get_solver_class(solver_name: str) -> "type[BaseSolver]":
    """
    Get the solver class based on the solver name.

    Args:
        solver_name: Name of the solver ("mpi", "numpy", "opencl", "jax")

    Returns:
        The solver class

    Raises:
        ValueError: If the selected solver is not compatible with MPI
            or if the solver name is unknown

    Example:
        >>> Solver = get_solver_class("numpy")
        >>> Solver.__name__
        "NumpySolver"
    """
    if size > 1 and solver_name != "mpi":
        raise ValueError(f"Solver method '{solver_name}' is not compatible with MPI.")

    Solver: type[BaseSolver]
    if solver_name == "mpi":
        from .mpi import MPISolver

        Solver = MPISolver
    elif solver_name == "numpy":
        from .numpy import NumpySolver

        Solver = NumpySolver

    elif solver_name == "opencl":
        from .opencl import OpenCLSolver

        Solver = OpenCLSolver

    elif solver_name == "jax":
        from .experimental.jax import JaxSolver

        Solver = JaxSolver
    else:
        raise ValueError(f"Unknown solver method '{solver_name}'.")

    return Solver
