import random
from typing import Literal, Sequence

from .operation import Operation
from .terminal import Terminal

type InitMethod = Literal["full", "grow"]


class Chromosome:
    def __init__(
        self,
        operations: Sequence[Operation],
        terminals: Sequence[Terminal],
        max_depth: int,
        init_method: InitMethod,
        terminal_probability: float = 0.5,
    ):
        self.operations = operations
        self.terminals = terminals
        self.root = (
            self._init_full(max_depth)
            if init_method == "full"
            else self._init_grow(max_depth, terminal_probability)
        )
        self.depth = self._compute_depth(self.root)

    def _compute_depth(self, node: "Chromosome.Node") -> int:
        """Вычисляет глубину дерева. Дерево из одного только корня имеет глубину 1."""
        return (
            max(self._compute_depth(child) for child in node.children) + 1
            if node.children
            else 1
        )

    def _random_terminal(self) -> Terminal:
        return random.choice(self.terminals)

    def _init_full(self, max_depth: int) -> "Chromosome.Node":
        """Полная инициализация.

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

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

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

        return build(1)

    def _init_grow(
        self, max_depth: int, terminal_probability: float
    ) -> "Chromosome.Node":
        """Растущая инициализация.

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

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

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

        return build(1)

    def eval(self, values: list[float]) -> float:
        """Вычисляет значение хромосомы для заданных значений терминалов."""
        for terminal, value in zip(self.terminals, values):
            terminal._value = value

        return self.root._eval()

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

    def _tree_lines(
        self, node: "Chromosome.Node", prefix: str = "", is_last: bool = True
    ) -> list[str]:
        connector = "└── " if is_last else "├── "
        lines = [prefix + connector + node.value.name]
        child_prefix = prefix + ("    " if is_last else "│   ")
        for i, child in enumerate(node.children):
            last = i == len(node.children) - 1
            lines.extend(self._tree_lines(child, child_prefix, last))
        return lines

    def str_tree(self) -> str:
        """Строковое представление древовидной структуры формулы."""
        lines = [self.root.value.name]
        for i, child in enumerate(self.root.children):
            last = i == len(self.root.children) - 1
            lines.extend(self._tree_lines(child, "", last))
        return "\n".join(lines)

    class Node:
        def __init__(
            self, value: Operation | Terminal, children: list["Chromosome.Node"]
        ):
            self.value = value
            self.children = children

        def __str__(self) -> str:
            """Рекурсивный перевод древовидного вида формулы в строку в инфиксной форме."""
            if isinstance(self.value, Terminal):
                return self.value.name

            if self.value.arity == 2:
                return f"({self.children[0]} {self.value.name} {self.children[1]})"

            return (
                f"{self.value.name}({', '.join(str(child) for child in self.children)})"
            )

        def _eval(self) -> float:
            """Рекурсивно вычисляет значение поддерева. Значения терминалов должны быть
            заданы предварительно."""
            if isinstance(self.value, Terminal):
                return self.value._value  # type: ignore

            return self.value._eval([child._eval() for child in self.children])