import random
from typing import Optional

from .chromosome import Chromosome
from .operation import Operation
from .terminal import Terminal


def mutate_grow(parent: Chromosome, max_depth: int, max_tries: int = 64) -> Chromosome:
    """
    Растущая мутация (subtree-growing mutation).

    Аргументы:
        parent    : исходная хромосома (не изменяется).
        max_depth : верхняя граница глубины мутанта.
        max_tries : ограничение попыток подбора узла/поддерева.

    Возвращает:
        Новый экземпляр Chromosome (мутант).
    """

    # ---------- Вспомогательные ----------
    def enumerate_nodes_with_meta(
        root: Chromosome.Node,
    ) -> list[
        tuple[Chromosome.Node, Optional[Chromosome.Node], Optional[int], list[int], int]
    ]:
        """
        (node, parent, index_in_parent, path_from_root, depth_from_root)
        depth_from_root: 1 для корня, 2 для детей корня и т.д.
        """
        out: list[
            tuple[
                Chromosome.Node,
                Optional[Chromosome.Node],
                Optional[int],
                list[int],
                int,
            ]
        ] = []

        def dfs(
            n: Chromosome.Node,
            p: Optional[Chromosome.Node],
            idx: Optional[int],
            path: list[int],
            depth: int,
        ) -> None:
            out.append((n, p, idx, path, depth))
            for i, ch in enumerate(n.children):
                dfs(ch, n, i, path + [i], depth + 1)

        dfs(root, None, None, [], 1)
        return out

    def get_parent_and_index_by_path(
        root: Chromosome.Node, path: list[int]
    ) -> tuple[Optional[Chromosome.Node], Optional[int], Chromosome.Node]:
        if not path:
            return None, None, root
        parent = root
        for i in path[:-1]:
            parent = parent.children[i]
        idx = path[-1]
        return parent, idx, parent.children[idx]

    def build_depth_limited_subtree(
        chromo: Chromosome, max_depth_limit: int, current_depth: int = 1
    ) -> Chromosome.Node:
        """
        Строит случайное поддерево с ограничением по глубине.
        current_depth: текущая глубина узла (1 для корня поддерева).
        max_depth_limit: максимальная допустимая глубина поддерева.
        """
        # Если достигли максимальной глубины — обязательно терминал
        if current_depth >= max_depth_limit:
            term = random.choice(chromo.terminals)
            return Chromosome.Node(term, [])

        # Иначе случайно выбираем между операцией и терминалом
        # С большей вероятностью выбираем операцию, если глубина позволяет
        if random.random() < 0.7:  # 70% вероятность операции
            op = random.choice(chromo.operations)
            children = [
                build_depth_limited_subtree(chromo, max_depth_limit, current_depth + 1)
                for _ in range(op.arity)
            ]
            return Chromosome.Node(op, children)
        else:
            term = random.choice(chromo.terminals)
            return Chromosome.Node(term, [])

    # ---------- Подготовка ----------
    # Если в дереве только терминал — мутация невозможна (нужен нетерминал)
    if isinstance(parent.root.value, Terminal):
        return parent.clone()

    # Работаем на клоне
    child = parent.clone()

    # Список нетерминальных узлов с путями и глубинами
    nodes = enumerate_nodes_with_meta(child.root)
    internal = [
        (n, p, i, path, depth)
        for (n, p, i, path, depth) in nodes
        if isinstance(n.value, Operation)
    ]

    if not internal:
        # На всякий случай: если всё терминалы
        return child

    # ---------- Основной цикл подбора позиции ----------
    for _ in range(max_tries):
        node, _, _, path, node_depth = random.choice(internal)

        # Вычисляем максимальную допустимую глубину для нового поддерева
        # max_depth - node_depth + 1 (так как node_depth начинается с 1)
        allowed_subtree_depth = max_depth - node_depth + 1

        if allowed_subtree_depth < 1:
            # Этот узел слишком глубоко — попробуем другой
            continue

        # Строим новое поддерево с ограничением по глубине
        new_sub = build_depth_limited_subtree(child, allowed_subtree_depth)

        # Вставляем его на место узла мутации
        parent_node, idx, _ = get_parent_and_index_by_path(child.root, path)
        if parent_node is None:
            child.root = new_sub
        else:
            parent_node.children[idx] = new_sub  # type: ignore[index]

        # Проверяем, что не превысили максимальную глубину
        if child.get_depth() <= max_depth:
            return child
        else:
            # Откат: пересоздаём клон
            child = parent.clone()

    # Если не удалось подобрать подходящее место/поддерево — вернём немутированного клона
    return parent.clone()