import numba as nb
import numpy as np

from vpso.typing import Array2d, Array3d


@nb.njit(
    nb.float64[:, :](
        nb.float64[:, :],  # x_best
        nb.bool_[:, :],  # mutation_mask
        nb.float64[:, :],  # lb
        nb.float64[:, :],  # ub
        nb.int32,  # nvec
        nb.int32,  # dim
        nb.types.NumPyRandomGeneratorType("NumPyRandomGeneratorType"),
    ),
    cache=True,
    nogil=True,
)
def polynomial_mutation(
    x_best: Array2d,
    mutation_mask: Array2d,
    lb: Array2d,
    ub: Array2d,
    nvec: int,
    dim: int,
    np_random: np.random.Generator,
) -> Array2d:
    """Performs a polynomial mutation on the best particle of each problem.

    Parameters
    ----------
    x_best : 2d array
        Best particles eligible for mutation. An array of shape `(nvec, dim)`.
    mutation_mask : 2d array
        Boolean mask indicating which particles should be mutated. An array of shape
        `(nvec, dim)`.
    lb : 2d array
        Lower bound of the search space. An array of shape `(nvec, dim)`.
    ub : 2d array
        Upper bound of the search space. An array of shape `(nvec, dim)`.
    nvec : int
        Number of vectorized problems.
    dim : int
        Dimensionality of the problem.
    np_random : np.random.Generator
        Random number generator.

    Returns
    -------
    2d array
        Modified best particles. An array of shape `(nvec, dim)`.
    """
    domain = ub - lb
    eta = np_random.uniform(6.0, 31.0, (nvec, np.int32(1)))
    mut_pow = 1.0 / eta
    xy1 = np.power((ub - x_best) / domain, eta)
    xy2 = np.power((x_best - lb) / domain, eta)
    R = np_random.random((nvec, dim))
    val1 = np.power(2.0 * R + (1.0 - 2.0 * R) * xy1, mut_pow)
    val2 = np.power(2.0 * (1.0 - R) + 2.0 * (R - 0.5) * xy2, mut_pow)
    deltaq = np.where(R <= 0.5, val1 - 1.0, 1.0 - val2)
    return np.where(mutation_mask, x_best + deltaq * domain, x_best).clip(lb, ub)


@nb.njit(
    nb.types.void(
        nb.float64[:, :, :],  # x
        nb.float64[:, :, :],  # px
        nb.float64[:, :],  # pf
        nb.float64[:, :, :],  # lb
        nb.float64[:, :, :],  # ub
        nb.int32,  # nvec
        nb.int32,  # dim
        nb.float64,  # mutation_prob
        nb.types.NumPyRandomGeneratorType("NumPyRandomGeneratorType"),
    ),
    cache=True,
    nogil=True,
)
def mutate(
    x: Array3d,
    px: Array3d,
    pf: Array2d,
    lb: Array3d,
    ub: Array3d,
    nvec: int,
    dim: int,
    mutation_prob: float,
    np_random: np.random.Generator,
) -> None:
    """Mutates in-place the best particle of each problem with a polynomial mutation.

    Parameters
    ----------
    x : 3d array
        Current positions of the particles. An array of shape `(N, M, d)`, where `N` is
        the number of vectorized problems to solve simultaneously, `M` is the number of
        particles in the swarm, and `d` is the dimension of the search space. This array
        is modified in-place.
    px : 3d array
        Best positions of the particles so far. An array of shape `(N, M, d)`.
    pf : 2d array
        Best values of the particles so far. An array of shape `(N, M)`.
    lb : 3d array
        Lower bound of the search space. An array of shape `(N, 1, d)`.
    ub : 3d array
        Upper bound of the search space. An array of shape `(N, 1, d)`.
    nvec : int
        Number of vectorized problems.
    dim : int
        Dimensionality of the problem.
    mutation_prob : float
        Probability of applying the mutation.
    np_random : np.random.Generator
        Random number generator.
    """
    # get the mutation mask for each vectorized problem, and for each dimension
    mutation_mask = np.logical_and(
        np_random.random((nvec, np.int32(1))) <= mutation_prob,
        np_random.random((nvec, dim)) <= min(0.5, 1 / dim),
    )
    if not mutation_mask.any():
        return

    # pick the best particle from each problem as target for mutation
    k = pf.argmin(1)
    # x_best = px[np.arange(nvec), k]  # cannot do this in njit
    x_best = np.empty((nvec, dim))
    for i, j in enumerate(k):
        x_best[i] = px[i, j]

    # compute mutation and apply it to the particle
    x_mutated = polynomial_mutation(
        x_best, mutation_mask, lb[:, 0], ub[:, 0], nvec, dim, np_random
    )
    for i, j in enumerate(k):
        x[i, j] = x_mutated[i]
