import random
from typing import Callable

from .chromosome import Chromosome, InitMethod

type Population = list[Chromosome]


def ramped_initialization(
    population_size: int,
    depths: list[int],
    make_chromosome: Callable[[int, InitMethod], Chromosome],  # (max_depth, method)
) -> Population:
    """Комбинация методов grow и full инициализации хромосом для инициализации начальной
    популяции.

    Начальная популяция генерируется так, чтобы в нее входили деревья с разной
    максимальной длиной примерно поровну. Для каждой глубины первая половина деревьев
    генерируется полным методом, а вторая – растущей инициализацией.
    """
    population: Population = []
    per_depth = population_size / len(depths)

    for depth in depths:
        n_full = int(per_depth / 2)
        n_grow = int(per_depth / 2)

        population.extend(make_chromosome(depth, "full") for _ in range(n_full))
        population.extend(make_chromosome(depth, "grow") for _ in range(n_grow))

    # Из-за округления хромосом может оказаться меньше заданного количества,
    # поэтому дозаполняем остаток популяции случайными хромосомами
    while len(population) < population_size:
        depth = random.choice(depths)
        method = "full" if random.random() < 0.5 else "grow"
        population.append(make_chromosome(depth, method))

    return population