import random
from abc import ABC, abstractmethod
from typing import Sequence

from .chromosome import Chromosome


class BaseMutation(ABC):
    @abstractmethod
    def mutate(self, chromosome: Chromosome) -> Chromosome: ...
    def __call__(self, chromosome: Chromosome) -> Chromosome:
        chromosome = chromosome.copy()
        return self.mutate(chromosome)


class ShrinkMutation(BaseMutation):
    """Усекающая мутация. Заменяет случайно выбранную операцию на случайный терминал."""

    def mutate(self, chromosome: Chromosome) -> Chromosome:
        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


class GrowMutation(BaseMutation):
    """Растущая мутация. Заменяет случайно выбранный узел на случайное поддерево."""

    def __init__(self, max_depth: int):
        self.max_depth = max_depth

    def mutate(self, chromosome: Chromosome) -> Chromosome:
        target_node = random.choice(chromosome.root.list_nodes())

        max_subtree_depth = self.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


class NodeReplacementMutation(BaseMutation):
    """Мутация замены операции (Node Replacement Mutation).

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

    Если подходящей альтернативы нет — возвращает копию без изменений.
    """

    def mutate(self, chromosome: Chromosome) -> Chromosome:
        target_node = random.choice(chromosome.root.list_nodes())
        current_arity = target_node.value.arity

        same_arity = [
            op
            for op in list(chromosome.operations) + list(chromosome.terminals)
            if op.arity == current_arity and op != target_node.value
        ]
        if not same_arity:
            return chromosome

        new_operation = random.choice(same_arity)

        target_node.value = new_operation

        return chromosome


class HoistMutation(BaseMutation):
    def mutate(self, chromosome: Chromosome) -> Chromosome:
        """Hoist-мутация (анти-bloat).

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

        В результате дерево становится короче, сохраняя часть структуры.
        """
        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


class CombinedMutation(BaseMutation):
    """Комбинированная мутация.

    Принимает список (или словарь) мутаций и случайно выбирает одну из них
    для применения. Можно задать веса вероятностей.
    """

    def __init__(
        self, mutations: Sequence[BaseMutation], probs: Sequence[float] | None = None
    ):
        if probs is not None:
            assert abs(sum(probs) - 1.0) < 1e-8, (
                "Сумма вероятностей должна быть равна 1"
            )
            assert len(probs) == len(mutations), (
                "Число вероятностей должно совпадать с числом мутаций"
            )
        self.mutations = mutations
        self.probs = probs

    def mutate(self, chromosome: Chromosome) -> Chromosome:
        mutation = random.choices(self.mutations, weights=self.probs, k=1)[0]
        return mutation(chromosome)