import random
from typing import Sequence

from .node import Node
from .primitive import Primitive


class Chromosome:
    def __init__(
        self,
        terminals: Sequence[Primitive],
        operations: Sequence[Primitive],
        root: Node,
    ):
        self.terminals = terminals
        self.operations = operations
        self.root = root

    def copy(self) -> Chromosome:
        return Chromosome(self.terminals, self.operations, self.root.copy_subtree())

    def prune(self, max_depth: int) -> None:
        self.root.prune(self.terminals, max_depth)

    def __str__(self) -> str:
        """Строковое представление хромосомы в виде формулы в инфиксной форме."""
        return str(self.root)

    @classmethod
    def full_init(
        cls,
        terminals: Sequence[Primitive],
        operations: Sequence[Primitive],
        max_depth: int,
    ) -> Chromosome:
        """Полная инициализация.

        В полном методе при генерации дерева, пока не достигнута максимальная глубина,
        допускается выбор только функциональных символов, а на последнем уровне
        (максимальной глубины) выбираются только терминальные символы.
        """

        def build(level: int) -> Node:
            # Если достигнута максимальная глубина — выбираем терминал
            if level == max_depth:
                return Node(random.choice(terminals))

            # Иначе выбираем операцию и создаём потомков
            op = random.choice(operations)
            node = Node(op)
            for _ in range(op.arity):
                node.add_child(build(level + 1))
            return node

        return cls(terminals, operations, build(1))

    @classmethod
    def grow_init(
        cls,
        terminals: Sequence[Primitive],
        operations: Sequence[Primitive],
        max_depth: int,
        # min_depth: int, # ???
        terminal_probability: float = 0.5,
    ) -> Chromosome:
        """Растущая инициализация.

        В растущей инициализации генерируются нерегулярные деревья с различной глубиной
        листьев вследствие случайного на каждом шаге выбора функционального
        или терминального символа. Здесь при выборе терминального символа рост дерева
        прекращается по текущей ветви и поэтому дерево имеет нерегулярную структуру.
        """

        def build(level: int) -> Node:
            # Если достигнута максимальная глубина, либо сыграла заданная вероятность
            # — выбираем терминал
            if level == max_depth or random.random() < terminal_probability:
                return Node(random.choice(terminals))

            # Иначе выбираем случайную операцию и создаём потомков
            op = random.choice(operations)
            node = Node(op)
            for _ in range(op.arity):
                node.add_child(build(level + 1))
            return node

        return cls(terminals, operations, build(1))