# createsonline/nlp/tokenizer.py
"""
CREATESONLINE Tokenizers - Pure Python Implementation

Byte-Pair Encoding (BPE), Word-level, and Character-level tokenizers.
Zero external dependencies.
"""

import re
import json
import os
from typing import Dict, List, Optional, Tuple, Set
from collections import Counter, defaultdict

from .vocabulary import Vocabulary


class CharTokenizer:
    """
    Character-level tokenizer.
    Splits text into individual characters.
    
    Usage:
        tok = CharTokenizer()
        tok.train(["Hello world"])
        ids = tok.encode("Hello")
        text = tok.decode(ids)
    """
    
    def __init__(self, max_vocab_size: int = 500):
        self.vocab = Vocabulary(max_size=max_vocab_size, min_frequency=1)
    
    def train(self, texts: List[str]):
        """Build character vocabulary from texts"""
        counter = Counter()
        for text in texts:
            for char in text:
                counter[char] += 1
        self.vocab.build_from_counter(counter)
    
    def encode(self, text: str, add_special: bool = True) -> List[int]:
        """Encode text to character IDs"""
        ids = [self.vocab.token_to_id(c) for c in text]
        if add_special:
            ids = [self.vocab.bos_id] + ids + [self.vocab.eos_id]
        return ids
    
    def decode(self, ids: List[int], skip_special: bool = True) -> str:
        """Decode character IDs back to text"""
        tokens = self.vocab.decode(ids, skip_special=skip_special)
        return "".join(tokens)
    
    @property
    def vocab_size(self) -> int:
        return self.vocab.size
    
    def save(self, dirpath: str):
        """Save tokenizer to directory"""
        os.makedirs(dirpath, exist_ok=True)
        self.vocab.save(os.path.join(dirpath, "vocab.json"))
        with open(os.path.join(dirpath, "config.json"), 'w') as f:
            json.dump({"type": "char", "max_vocab_size": self.vocab.max_size}, f)
    
    @classmethod
    def load(cls, dirpath: str) -> 'CharTokenizer':
        """Load tokenizer from directory"""
        with open(os.path.join(dirpath, "config.json"), 'r') as f:
            config = json.load(f)
        tok = cls(max_vocab_size=config.get("max_vocab_size", 500))
        tok.vocab = Vocabulary.load(os.path.join(dirpath, "vocab.json"))
        return tok


class WordTokenizer:
    """
    Word-level tokenizer with regex-based splitting.
    
    Usage:
        tok = WordTokenizer()
        tok.train(["Hello world! How are you?"])
        ids = tok.encode("Hello world")
        text = tok.decode(ids)
    """
    
    # Regex pattern: words, numbers, punctuation
    PATTERN = re.compile(r"[a-zA-Z]+|[0-9]+|[^\s\w]")
    
    def __init__(self, max_vocab_size: int = 50000, min_frequency: int = 1, lowercase: bool = True):
        self.vocab = Vocabulary(max_size=max_vocab_size, min_frequency=min_frequency)
        self.lowercase = lowercase
    
    def _tokenize(self, text: str) -> List[str]:
        """Split text into word tokens"""
        if self.lowercase:
            text = text.lower()
        return self.PATTERN.findall(text)
    
    def train(self, texts: List[str]):
        """Build word vocabulary from texts"""
        self.vocab.build_from_texts(texts, tokenize_fn=self._tokenize)
    
    def encode(self, text: str, add_special: bool = True) -> List[int]:
        """Encode text to word IDs"""
        tokens = self._tokenize(text)
        ids = self.vocab.encode(tokens)
        if add_special:
            ids = [self.vocab.bos_id] + ids + [self.vocab.eos_id]
        return ids
    
    def decode(self, ids: List[int], skip_special: bool = True) -> str:
        """Decode word IDs back to text"""
        tokens = self.vocab.decode(ids, skip_special=skip_special)
        return " ".join(tokens)
    
    @property
    def vocab_size(self) -> int:
        return self.vocab.size
    
    def save(self, dirpath: str):
        """Save tokenizer to directory"""
        os.makedirs(dirpath, exist_ok=True)
        self.vocab.save(os.path.join(dirpath, "vocab.json"))
        with open(os.path.join(dirpath, "config.json"), 'w') as f:
            json.dump({
                "type": "word",
                "max_vocab_size": self.vocab.max_size,
                "min_frequency": self.vocab.min_frequency,
                "lowercase": self.lowercase
            }, f)
    
    @classmethod
    def load(cls, dirpath: str) -> 'WordTokenizer':
        """Load tokenizer from directory"""
        with open(os.path.join(dirpath, "config.json"), 'r') as f:
            config = json.load(f)
        tok = cls(
            max_vocab_size=config.get("max_vocab_size", 50000),
            min_frequency=config.get("min_frequency", 1),
            lowercase=config.get("lowercase", True)
        )
        tok.vocab = Vocabulary.load(os.path.join(dirpath, "vocab.json"))
        return tok


class BPETokenizer:
    """
    Byte-Pair Encoding (BPE) tokenizer — pure Python implementation.
    
    This is the same algorithm used by GPT-2/GPT-3/GPT-4.
    Iteratively merges the most frequent adjacent token pairs.
    
    Usage:
        tok = BPETokenizer(vocab_size=1000)
        tok.train(["Hello world! This is a test."])
        ids = tok.encode("Hello world")
        text = tok.decode(ids)
    
    How BPE works:
        1. Start with character-level tokens
        2. Count frequency of all adjacent token pairs
        3. Merge the most frequent pair into a new token
        4. Repeat until desired vocabulary size
    """
    
    def __init__(self, vocab_size: int = 10000, min_frequency: int = 2):
        self.target_vocab_size = vocab_size
        self.min_frequency = min_frequency
        
        # BPE merge rules: list of (token_a, token_b) pairs in merge order
        self.merges: List[Tuple[str, str]] = []
        
        # Token vocabulary
        self.vocab = Vocabulary(max_size=vocab_size + 100, min_frequency=1)
        
        # Precompiled regex for pre-tokenization (GPT-2 style)
        self._pat = re.compile(
            r"""'s|'t|'re|'ve|'m|'ll|'d| ?\w+| ?\d+| ?[^\s\w\d]+|\s+(?!\S)|\s+""",
            re.IGNORECASE
        )
        
        self._trained = False
    
    def _pre_tokenize(self, text: str) -> List[str]:
        """
        Pre-tokenize text into word-level chunks (GPT-2 style).
        Each word becomes a sequence of characters for BPE processing.
        """
        return self._pat.findall(text)
    
    def _word_to_chars(self, word: str) -> Tuple[str, ...]:
        """Convert a word to a tuple of characters (BPE tokens)"""
        if not word:
            return ()
        # Add end-of-word marker
        chars = tuple(word[:-1]) + (word[-1] + "</w>",)
        return chars
    
    def _get_pair_counts(self, word_freqs: Dict[Tuple[str, ...], int]) -> Counter:
        """Count frequency of all adjacent token pairs across the corpus"""
        pairs = Counter()
        for word_tokens, freq in word_freqs.items():
            if len(word_tokens) < 2:
                continue
            for i in range(len(word_tokens) - 1):
                pair = (word_tokens[i], word_tokens[i + 1])
                pairs[pair] += freq
        return pairs
    
    def _merge_pair(self, word_tokens: Tuple[str, ...], pair: Tuple[str, str]) -> Tuple[str, ...]:
        """Merge all occurrences of a pair in a word"""
        new_tokens = []
        i = 0
        while i < len(word_tokens):
            if i < len(word_tokens) - 1 and word_tokens[i] == pair[0] and word_tokens[i + 1] == pair[1]:
                new_tokens.append(pair[0] + pair[1])
                i += 2
            else:
                new_tokens.append(word_tokens[i])
                i += 1
        return tuple(new_tokens)
    
    def train(self, texts: List[str], verbose: bool = False):
        """
        Train BPE tokenizer on a corpus of texts.
        
        Args:
            texts: List of training texts
            verbose: Print progress every 100 merges
        """
        # Step 1: Pre-tokenize and count word frequencies
        word_freqs: Dict[Tuple[str, ...], int] = defaultdict(int)
        for text in texts:
            words = self._pre_tokenize(text)
            for word in words:
                chars = self._word_to_chars(word)
                if chars:
                    word_freqs[chars] += 1
        
        # Step 2: Initialize vocabulary with all unique characters
        all_chars: Set[str] = set()
        for word_tokens in word_freqs:
            for token in word_tokens:
                all_chars.add(token)
        
        for char in sorted(all_chars):
            self.vocab.add_token(char)
        
        initial_size = self.vocab.size
        num_merges = self.target_vocab_size - initial_size
        
        if verbose:
            print(f"Initial vocab size: {initial_size}")
            print(f"Target merges: {max(0, num_merges)}")
        
        # Step 3: Iteratively merge most frequent pairs
        self.merges = []
        for merge_idx in range(max(0, num_merges)):
            # Count pairs
            pair_counts = self._get_pair_counts(word_freqs)
            if not pair_counts:
                break
            
            # Find most frequent pair
            best_pair = pair_counts.most_common(1)[0]
            pair, count = best_pair
            
            if count < self.min_frequency:
                break
            
            # Record merge
            self.merges.append(pair)
            
            # Merge in all words
            new_word_freqs = {}
            for word_tokens, freq in word_freqs.items():
                new_word = self._merge_pair(word_tokens, pair)
                new_word_freqs[new_word] = freq
            word_freqs = new_word_freqs
            
            # Add merged token to vocabulary
            merged_token = pair[0] + pair[1]
            self.vocab.add_token(merged_token)
            
            if verbose and (merge_idx + 1) % 100 == 0:
                print(f"Merge {merge_idx + 1}/{num_merges}: {pair} -> {merged_token} (count={count})")
        
        self._trained = True
        
        if verbose:
            print(f"Training complete. Final vocab size: {self.vocab.size}")
            print(f"Total merges: {len(self.merges)}")
    
    def _apply_bpe(self, word: str) -> List[str]:
        """Apply BPE merges to a single word"""
        if not word:
            return []
        
        tokens = list(self._word_to_chars(word))
        
        for pair in self.merges:
            i = 0
            new_tokens = []
            while i < len(tokens):
                if i < len(tokens) - 1 and tokens[i] == pair[0] and tokens[i + 1] == pair[1]:
                    new_tokens.append(pair[0] + pair[1])
                    i += 2
                else:
                    new_tokens.append(tokens[i])
                    i += 1
            tokens = new_tokens
            if len(tokens) == 1:
                break
        
        return tokens
    
    def tokenize(self, text: str) -> List[str]:
        """Tokenize text into BPE tokens"""
        words = self._pre_tokenize(text)
        tokens = []
        for word in words:
            tokens.extend(self._apply_bpe(word))
        return tokens
    
    def encode(self, text: str, add_special: bool = True) -> List[int]:
        """Encode text to token IDs"""
        tokens = self.tokenize(text)
        ids = self.vocab.encode(tokens)
        if add_special:
            ids = [self.vocab.bos_id] + ids + [self.vocab.eos_id]
        return ids
    
    def decode(self, ids: List[int], skip_special: bool = True) -> str:
        """Decode token IDs back to text"""
        tokens = self.vocab.decode(ids, skip_special=skip_special)
        # Reconstruct text by joining and removing end-of-word markers
        text = "".join(tokens)
        text = text.replace("</w>", " ")
        return text.strip()
    
    @property
    def vocab_size(self) -> int:
        return self.vocab.size
    
    def save(self, dirpath: str):
        """Save tokenizer (vocabulary + merge rules)"""
        os.makedirs(dirpath, exist_ok=True)
        self.vocab.save(os.path.join(dirpath, "vocab.json"))
        
        config = {
            "type": "bpe",
            "target_vocab_size": self.target_vocab_size,
            "min_frequency": self.min_frequency,
            "merges": [[a, b] for a, b in self.merges],
            "trained": self._trained,
        }
        with open(os.path.join(dirpath, "config.json"), 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
    
    @classmethod
    def load(cls, dirpath: str) -> 'BPETokenizer':
        """Load tokenizer from directory"""
        with open(os.path.join(dirpath, "config.json"), 'r', encoding='utf-8') as f:
            config = json.load(f)
        
        tok = cls(
            vocab_size=config.get("target_vocab_size", 10000),
            min_frequency=config.get("min_frequency", 2)
        )
        tok.vocab = Vocabulary.load(os.path.join(dirpath, "vocab.json"))
        tok.merges = [tuple(m) for m in config.get("merges", [])]
        tok._trained = config.get("trained", False)
        return tok
    
    def __repr__(self) -> str:
        return f"BPETokenizer(vocab_size={self.vocab_size}, merges={len(self.merges)}, trained={self._trained})"
