"""
TOV Token Graph: Иерархическое дерево токенов.
Предложение -> Слово -> Символ/Подслово.
"""

import sys
import os
import numpy as np
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# Вместо импорта внешних классов, создадим минимальные версии прямо здесь
class TOVCore:
    """Минимальное ядро для токенизатора"""
    def __init__(self, name="layer"):
        self.name = name
        self.a = 0.3
        self.b = 0.4
        self.c = 0.3
        self.output_mass = 0.0
        self.accumulated_mass = 0.0
        self.gamma = 0.0
        self.beta = 0.5
        
    def add_learned_mass(self, delta):
        self.accumulated_mass += delta
        self.output_mass = self.accumulated_mass * self.beta

class TOVTreeNode:
    """Минимальный узел для токенизатора"""
    def __init__(self, token_id, embedding):
        self.token_id = token_id
        self.embedding = embedding
        self.children = []
        self.parent = None
        self.core = TOVCore(f"node_{token_id}")
        self.frequency = 0
        self.last_seen = 0
        
    def add_child(self, child_node):
        self.children.append(child_node)
        child_node.parent = self

class TokenNode(TOVTreeNode):
    """Узел-токен с дополнительной информацией"""
    def __init__(self, token_id, token_str, level, embedding=None):
        super().__init__(token_id, embedding or np.zeros(20))
        self.token_str = token_str
        self.level = level  # 'root', 'sentence', 'word', 'subword', 'char'
        self.frequency = 0
        self.last_seen = 0

    def update_frequency(self, time_step):
        self.frequency += 1
        self.last_seen = time_step
        # Масса растёт с частотой
        self.core.add_learned_mass(0.1)

class TOVGraph:
    """Минимальный граф для токенизатора"""
    def __init__(self):
        self.nodes = {}
        self.root = None
        
    def add_node(self, name, embedding=None):
        """Добавляет узел в граф"""
        if embedding is None:
            embedding = np.zeros(20)
        node = TOVTreeNode(name, embedding)
        self.nodes[name] = node
        return node
        
    def add_edge(self, parent_id, child_id):
        """Добавляет связь между узлами"""
        if parent_id in self.nodes and child_id in self.nodes:
            self.nodes[parent_id].add_child(self.nodes[child_id])

class TokenGraph:
    """Граф для иерархической токенизации"""
    def __init__(self):
        self.nodes = {}
        self.token_map = {}  # token_str -> node_id
        self.time_step = 0
        self.root = self.add_token_node("<root>", "<root>", "root")

    def add_token_node(self, token_str, display_str, level):
        """Добавляет узел-токен в граф"""
        token_id = hash(token_str) % 1000000  # простой хеш
        node = TokenNode(token_id, display_str, level)
        self.nodes[str(token_id)] = node
        self.token_map[token_str] = str(token_id)
        return node

    def add_edge(self, parent_id, child_id):
        """Добавляет связь между узлами"""
        if parent_id in self.nodes and child_id in self.nodes:
            self.nodes[parent_id].add_child(self.nodes[child_id])

    def add_sequence(self, tokens, parent_node, level):
        """Добавляет последовательность токенов как дочерние узлы"""
        nodes = []
        for token in tokens:
            if token not in self.token_map:
                node = self.add_token_node(token, token, level)
            else:
                node_id = self.token_map[token]
                node = self.nodes[node_id]

            self.add_edge(str(parent_node.token_id), str(node.token_id))
            node.update_frequency(self.time_step)
            nodes.append(node)

        self.time_step += 1
        return nodes

    def prune_rare_tokens(self, threshold=0.01, time_window=1000):
        """Удаляет редкие токены (λ-утечка)"""
        current_time = self.time_step
        to_prune = []

        for node_id, node in self.nodes.items():
            if not hasattr(node, 'level') or node.level == 'root':
                continue

            # Если токен не видели давно и частота низкая
            age = current_time - node.last_seen
            if age > time_window and node.frequency < threshold * time_window:
                to_prune.append(node_id)

        for node_id in to_prune:
            # Удаляем из nodes
            if node_id in self.nodes:
                node = self.nodes[node_id]
                # Удаляем из token_map
                for token_str, tid in list(self.token_map.items()):
                    if tid == node_id:
                        del self.token_map[token_str]
                # Удаляем сам узел
                del self.nodes[node_id]

        return len(to_prune)

    def get_stats(self):
        """Возвращает статистику графа"""
        stats = {
            'total_nodes': len(self.nodes),
            'unique_tokens': len(self.token_map),
            'time_step': self.time_step
        }

        # Считаем по уровням
        level_counts = {}
        for node in self.nodes.values():
            if hasattr(node, 'level'):
                level_counts[node.level] = level_counts.get(node.level, 0) + 1
        stats['levels'] = level_counts

        return stats

# ===================== ТЕСТ =====================

if __name__ == "__main__":
    print("🚀 Тест TokenGraph")
    g = TokenGraph()
    print(f"✅ Граф создан")
    print(f"   Узлов: {len(g.nodes)}")
    print(f"   Корень: {g.root.token_id}")

    stats = g.get_stats()
    print(f"\n📊 Статистика:")
    for k, v in stats.items():
        print(f"   {k}: {v}")