import numpy as np

from .types import Fitnesses, Population


def roulette_selection(population: Population, fitnesses: Fitnesses) -> Population:
    """Селекция методом рулетки.

    Чем больше значение фитнеса, тем больше вероятность выбора особи. Для минимизации
    значения фитнеса нужно предварительно инвертировать.
    """
    # Чтобы работать с отрицательными f, сдвигаем значения фитнес функции на минимальное
    # значение в популяции. Вычитаем min_fit, т. к. min_fit может быть отрицательным.
    min_fit = np.min(fitnesses)
    shifted_fitnesses = fitnesses - min_fit + 1e-12

    # Получаем вероятности для каждой особи
    probs = shifted_fitnesses / np.sum(shifted_fitnesses)
    cum = np.cumsum(probs)

    # Выбираем особей методом рулетки
    selected = []
    for _ in population:
        r = np.random.random()
        idx = int(np.searchsorted(cum, r, side="left"))
        selected.append(population[idx])

    return selected


def tournament_selection(
    population: Population,
    fitnesses: Fitnesses,
    k: int = 3,
) -> Population:
    """Турнирная селекция.

    В каждом турнире случайно выбирается k особей, и побеждает та,
    у которой лучшее (наибольшее) значение фитнеса. Для минимизации
    значения фитнеса нужно предварительно инвертировать.

    Args:
        population: список особей (Population)
        fitnesses: список или массив фитнесов (Fitnesses)
        k: размер турнира

    Returns:
        Новая популяция того же размера
    """
    size = len(population)
    selected = []
    for _ in range(size):
        idxs = np.random.choice(size, size=k, replace=False)

        fits = fitnesses[idxs]

        winner_idx = idxs[np.argmax(fits)]
        selected.append(population[winner_idx])
    return selected


def stochastic_tournament_selection(
    population: Population,
    fitnesses: Fitnesses,
    k: int = 3,
    p_best: float = 0.75,
) -> Population:
    """Стохастическая турнирная селекция.

    Побеждает лучший в турнире с вероятностью p_best, иначе выбирается
    случайный участник турнира.
    """
    size = len(population)
    selected = []

    for _ in range(size):
        idxs = np.random.choice(size, size=k, replace=False)
        fits = fitnesses[idxs]
        order = np.argsort(-fits)

        if np.random.random() < p_best:
            winner_idx = idxs[order[0]]
        else:
            winner_idx = np.random.choice(idxs[1:]) if k > 1 else idxs[0]

        selected.append(population[winner_idx])

    return selected