"""
TOV Tokenizer: Токенизация через потоки массы.
Полностью на основе Теории Относительности Всего.
"""

import sys
import os
import numpy as np
import re
from collections import defaultdict
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from .graph import TokenGraph, TokenNode

class TOVTokenizerTOV:
    """
    Токенизатор, работающий через TOV-граф.
    - Иерархическое дерево токенов
    - Выбор токенов через потоки массы
    - γ-кэширование
    - λ-утечка редких токенов
    """
    
    def __init__(self, vocab_size=5000, min_frequency=2):
        self.vocab_size = vocab_size
        self.min_frequency = min_frequency
        self.graph = TokenGraph()
        
        # 👇 СПЕЦТОКЕНЫ ВСЕГДА В ПЕРВУЮ ОЧЕРЕДЬ
        self.specials = ['<pad>', '<unk>', '<sos>', '<eos>', '<sep>']
        self.token_to_id = {}
        self.id_to_token = {}
        
        for i, spec in enumerate(self.specials):
            self.token_to_id[spec] = i
            self.id_to_token[i] = spec
        
        self.stats = {
            'tokens_seen': 0,
            'unique_tokens': 0,
            'pruned_tokens': 0
        }
        
    def train(self, texts):
        """
        Обучает токенизатор на текстах.
        Строит иерархическое дерево.
        """
        print(f"\n🌳 ОБУЧЕНИЕ TOV-ТОКЕНИЗАТОРА")
        print(f"   Текстов: {len(texts)}")
        
        for i, text in enumerate(texts):
            if i % 100 == 0 and i > 0:
                print(f"   Обработано {i}/{len(texts)} текстов")
            self._add_text_to_graph(text)
        
        # Применяем λ-утечку
        pruned = self.graph.prune_rare_tokens(
            threshold=0.01, 
            time_window=max(100, len(texts)//10)
        )
        self.stats['pruned_tokens'] = pruned
        
        # Строим словарь из частотных токенов
        self._build_vocab_from_graph()
        
        print(f"✅ TOV-токенизатор готов:")
        print(f"   • Узлов в графе: {len(self.graph.nodes)}")
        print(f"   • Удалено редких: {pruned}")
        print(f"   • Размер словаря: {len(self.token_to_id)}")
    
    def _add_text_to_graph(self, text):
        """Добавляет текст в граф с иерархией"""
        # Разбиваем на предложения
        sentences = re.split(r'[.!?]+', text)
        
        for sent in sentences:
            if not sent.strip():
                continue
            
            # Добавляем предложение
            sent_node = self.graph.add_token_node(
                sent.strip(), sent[:30] + "...", "sentence"
            )
            self.graph.add_edge(str(self.graph.root.token_id), str(sent_node.token_id))
            sent_node.update_frequency(self.graph.time_step)
            
            # Разбиваем на слова
            words = re.findall(r'\w+', sent.lower())
            
            for word in words:
                if len(word) < 2:
                    continue
                    
                # Добавляем слово
                word_node = self.graph.add_token_node(word, word, "word")
                self.graph.add_edge(str(sent_node.token_id), str(word_node.token_id))
                word_node.update_frequency(self.graph.time_step)
                
                # Разбиваем на подслова (2-4 символа)
                subwords = self._generate_subwords(word)
                for subword in subwords:
                    sub_node = self.graph.add_token_node(subword, subword, "subword")
                    self.graph.add_edge(str(word_node.token_id), str(sub_node.token_id))
                    sub_node.update_frequency(self.graph.time_step)
                
                # Разбиваем на символы
                for i, char in enumerate(word):
                    char_node = self.graph.add_token_node(
                        char, char, "char"
                    )
                    self.graph.add_edge(str(word_node.token_id), str(char_node.token_id))
                    char_node.update_frequency(self.graph.time_step)
            
            self.graph.time_step += 1
    
    def _generate_subwords(self, word, min_len=2, max_len=4):
        """Генерирует подслова для BPE-like токенизации"""
        subwords = []
        word_len = len(word)
        for length in range(min_len, min(max_len, word_len) + 1):
            for i in range(word_len - length + 1):
                subword = word[i:i+length]
                subwords.append(subword)
        return list(set(subwords))  # уникальные
    
    def _build_vocab_from_graph(self):
        """Строит словарь из графа, беря самые частотные узлы"""
        if len(self.graph.nodes) <= 1:
            print("⚠️ Граф пуст, использую только спецтокены")
            return
        
        # Сортируем узлы по частоте
        nodes_by_freq = sorted(
            self.graph.nodes.values(),
            key=lambda n: n.frequency if hasattr(n, 'frequency') else 0,
            reverse=True
        )
        
        # Обычные токены (начиная после спецтокенов)
        idx = len(self.specials)
        for node in nodes_by_freq:
            if not hasattr(node, 'level') or node.level == 'root':
                continue
            if idx >= self.vocab_size:
                break
            if hasattr(node, 'token_str') and node.token_str not in self.token_to_id:
                self.token_to_id[node.token_str] = idx
                self.id_to_token[idx] = node.token_str
                idx += 1
    
    def save(self, path):
        """Сохраняет токенизатор в файл"""
        import pickle
        data = {
            'vocab_size': self.vocab_size,
            'min_frequency': self.min_frequency,
            'token_to_id': self.token_to_id,
            'id_to_token': self.id_to_token,
            'specials': self.specials,
            'stats': self.stats,
            'graph': self.graph  # сохраняем весь граф
        }
        with open(path, 'wb') as f:
            pickle.dump(data, f)
        print(f"💾 Токенизатор сохранён в {path}")

    def load(self, path):
        """Загружает токенизатор из файла"""
        import pickle
        with open(path, 'rb') as f:
            data = pickle.load(f)
        
        self.vocab_size = data['vocab_size']
        self.min_frequency = data['min_frequency']
        self.token_to_id = data['token_to_id']
        self.id_to_token = data['id_to_token']
        self.specials = data['specials']
        self.stats = data['stats']
        self.graph = data['graph']
        
        print(f"📂 Токенизатор загружен из {path}")
        print(f"   Токенов: {len(self.token_to_id)}")
        print(f"   Узлов в графе: {len(self.graph.nodes)}")

    def encode(self, text, add_special_tokens=True):
        """
        Токенизирует текст через потоки массы.
        """
        tokens = []
        
        if add_special_tokens:
            tokens.append(self.token_to_id['<sos>'])  # теперь точно есть
        
        # Разбиваем на предложения
        sentences = re.split(r'[.!?]+', text)
        
        for sent in sentences:
            if not sent.strip():
                continue
            
            # Ищем узел предложения
            sent_node = self._find_best_node(sent.strip(), 'sentence')
            if sent_node and sent_node.token_str in self.token_to_id:
                tokens.append(self.token_to_id[sent_node.token_str])
            
            # Разбиваем на слова
            words = re.findall(r'\w+', sent.lower())
            for word in words:
                # Ищем узел слова
                word_node = self._find_best_node(word, 'word')
                if word_node and word_node.token_str in self.token_to_id:
                    tokens.append(self.token_to_id[word_node.token_str])
                else:
                    # Если слова нет, разбиваем на символы
                    for char in word:
                        char_node = self._find_best_node(char, 'char')
                        if char_node and char_node.token_str in self.token_to_id:
                            tokens.append(self.token_to_id[char_node.token_str])
                        else:
                            tokens.append(1)  # <unk>
        
        if add_special_tokens:
            tokens.append(self.token_to_id['<eos>'])
        
        return tokens
    
    def _find_best_node(self, token_str, level):
        """Находит лучший узел через поток массы"""
        candidates = []
        for node_id, node in self.graph.nodes.items():
            if (hasattr(node, 'level') and node.level == level and 
                hasattr(node, 'token_str') and node.token_str == token_str):
                candidates.append(node)
        
        if not candidates:
            return None
        
        # Выбираем по потоку массы от корня
        best_node = None
        best_flow = -float('inf')
        
        for node in candidates:
            flow = node.core.a * node.core.output_mass
            if flow > best_flow:
                best_flow = flow
                best_node = node
        
        return best_node
    
    def decode(self, ids, skip_special=True):
        """
        Декодирует токены обратно в текст.
        """
        text = []
        for i in ids:
            token = self.id_to_token.get(i, '<unk>')
            if skip_special and token.startswith('<'):
                continue
            text.append(token)
        return ' '.join(text)
    
    def get_stats(self):
        """Возвращает статистику токенизатора"""
        stats = {
            'vocab_size': len(self.token_to_id),
            'graph_nodes': len(self.graph.nodes),
            'tokens_seen': self.graph.time_step,
            'pruned_tokens': self.stats['pruned_tokens']
        }
        
        # Добавляем статистику графа
        try:
            graph_stats = self.graph.get_stats()
            stats.update(graph_stats)
        except:
            pass
        
        return stats

    def __len__(self):
        """Возвращает размер словаря (нужно для len(tokenizer))"""
        return len(self.token_to_id)

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

if __name__ == "__main__":
    print("🚀 Тест TOVTokenizerTOV")
    
    # Тестовые тексты
    texts = [
        "привет, как дела?",
        "нормально, а у тебя?",
        "отлично! погнали дальше"
    ]
    
    t = TOVTokenizerTOV(vocab_size=100)
    t.train(texts)
    
    for text in texts:
        encoded = t.encode(text)
        decoded = t.decode(encoded)
        print(f"\n📝 Оригинал: {text}")
        print(f"   Токенов: {len(encoded)}")
        print(f"   Декод: {decoded}")
    
    stats = t.get_stats()
    print("\n📊 Статистика:")
    for k, v in stats.items():
        print(f"   {k}: {v}")