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


def node_replacement_mutation(chromosome: Chromosome) -> Chromosome:
    """Мутация замены операции (Node Replacement Mutation).

    Выбирает случайный узел с операцией (arity > 0) и заменяет его
    на случайную другую операцию той же арности, сохраняя поддеревья.

    Если подходящей альтернативы нет — возвращает копию без изменений.
    """
    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)
    current_arity = target_node.value.arity

    same_arity_ops = [
        op
        for op in chromosome.operations
        if op.arity == current_arity and op != target_node.value
    ]
    if not same_arity_ops:
        return chromosome

    new_operation = random.choice(same_arity_ops)

    target_node.value = new_operation

    return chromosome


def hoist_mutation(chromosome: Chromosome) -> Chromosome:
    """Hoist-мутация (анти-bloat).

    Выбирает случайное поддерево, затем внутри него — случайное поддерево меньшей глубины,
    и заменяет исходное поддерево на это внутреннее.

    В результате дерево становится короче, сохраняя часть структуры.
    """
    chromosome = chromosome.copy()

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

    outer_subtree = random.choice(operation_nodes)
    outer_nodes = outer_subtree.list_nodes()[1:]  # исключаем корень

    inner_subtree = random.choice(outer_nodes).copy_subtree()

    if outer_subtree.parent:
        outer_subtree.parent.replace_child(outer_subtree, inner_subtree)
    else:
        chromosome.root = inner_subtree

    return chromosome