import random

from .chromosome import Chromosome


def shrink_mutation(chromosome: Chromosome) -> Chromosome:
    """Усекающая мутация. Заменяет случайно выбранную операцию на случайный терминал."""
    chromosome = chromosome.copy()

    operation_nodes = [n for n in chromosome.root.list_nodes() if n.value.arity > 0]

    if not operation_nodes:
        return chromosome

    target_node = random.choice(operation_nodes)

    target_node.prune(chromosome.terminals, max_depth=1)

    return chromosome


def grow_mutation(chromosome: Chromosome, max_depth: int) -> Chromosome:
    """Растущая мутация. Заменяет случайно выбранный узел на случайное поддерево."""
    chromosome = chromosome.copy()

    target_node = random.choice(chromosome.root.list_nodes())

    max_subtree_depth = max_depth - target_node.get_level() + 1

    subtree = Chromosome.grow_init(
        chromosome.terminals, chromosome.operations, max_subtree_depth
    ).root

    if target_node.parent:
        target_node.parent.replace_child(target_node, subtree)
    else:
        chromosome.root = subtree

    return chromosome