# createsonline/nlp/vocabulary.py
"""
CREATESONLINE Vocabulary Management

Pure Python vocabulary builder with special tokens, 
frequency tracking, and serialization.
"""

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


class Vocabulary:
    """
    Vocabulary manager for tokenization.
    
    Manages token-to-ID and ID-to-token mappings with special tokens,
    frequency tracking, and save/load capabilities.
    
    Usage:
        vocab = Vocabulary(max_size=10000)
        vocab.build_from_texts(["Hello world", "How are you"])
        ids = vocab.encode(["hello", "world"])
        tokens = vocab.decode(ids)
    """
    
    # Special tokens
    PAD_TOKEN = "[PAD]"
    UNK_TOKEN = "[UNK]"
    BOS_TOKEN = "[BOS]"
    EOS_TOKEN = "[EOS]"
    SEP_TOKEN = "[SEP]"
    CLS_TOKEN = "[CLS]"
    MASK_TOKEN = "[MASK]"
    
    SPECIAL_TOKENS = [PAD_TOKEN, UNK_TOKEN, BOS_TOKEN, EOS_TOKEN, SEP_TOKEN, CLS_TOKEN, MASK_TOKEN]
    
    def __init__(self, max_size: int = 50000, min_frequency: int = 1):
        self.max_size = max_size
        self.min_frequency = min_frequency
        
        # Core mappings
        self._token_to_id: Dict[str, int] = {}
        self._id_to_token: Dict[int, str] = {}
        self._token_freq: Counter = Counter()
        
        # Initialize special tokens
        self._init_special_tokens()
    
    def _init_special_tokens(self):
        """Initialize special tokens with reserved IDs"""
        for i, token in enumerate(self.SPECIAL_TOKENS):
            self._token_to_id[token] = i
            self._id_to_token[i] = token
    
    @property
    def size(self) -> int:
        """Current vocabulary size"""
        return len(self._token_to_id)
    
    @property
    def pad_id(self) -> int:
        return self._token_to_id[self.PAD_TOKEN]
    
    @property
    def unk_id(self) -> int:
        return self._token_to_id[self.UNK_TOKEN]
    
    @property
    def bos_id(self) -> int:
        return self._token_to_id[self.BOS_TOKEN]
    
    @property
    def eos_id(self) -> int:
        return self._token_to_id[self.EOS_TOKEN]
    
    def add_token(self, token: str, count: int = 1) -> int:
        """Add token to vocabulary. Returns token ID."""
        self._token_freq[token] += count
        if token not in self._token_to_id:
            if self.size >= self.max_size:
                return self.unk_id
            token_id = self.size
            self._token_to_id[token] = token_id
            self._id_to_token[token_id] = token
            return token_id
        return self._token_to_id[token]
    
    def add_tokens(self, tokens: List[str]):
        """Add multiple tokens"""
        for token in tokens:
            self.add_token(token)
    
    def build_from_texts(self, texts: List[str], tokenize_fn=None):
        """
        Build vocabulary from a list of texts.
        
        Args:
            texts: List of text strings
            tokenize_fn: Optional tokenization function, defaults to whitespace split
        """
        if tokenize_fn is None:
            tokenize_fn = lambda t: t.lower().split()
        
        # Count all tokens
        counter = Counter()
        for text in texts:
            tokens = tokenize_fn(text)
            counter.update(tokens)
        
        # Sort by frequency (descending), then alphabetically
        sorted_tokens = sorted(counter.items(), key=lambda x: (-x[1], x[0]))
        
        # Add tokens respecting min_frequency and max_size
        for token, count in sorted_tokens:
            if count < self.min_frequency:
                continue
            if self.size >= self.max_size:
                break
            self.add_token(token, count)
    
    def build_from_counter(self, counter: Counter):
        """Build vocabulary from a Counter of token frequencies"""
        sorted_tokens = sorted(counter.items(), key=lambda x: (-x[1], x[0]))
        for token, count in sorted_tokens:
            if count < self.min_frequency:
                continue
            if self.size >= self.max_size:
                break
            self.add_token(token, count)
    
    def token_to_id(self, token: str) -> int:
        """Convert token to ID. Returns UNK_ID for unknown tokens."""
        return self._token_to_id.get(token, self.unk_id)
    
    def id_to_token(self, token_id: int) -> str:
        """Convert ID to token. Returns UNK_TOKEN for unknown IDs."""
        return self._id_to_token.get(token_id, self.UNK_TOKEN)
    
    def encode(self, tokens: List[str]) -> List[int]:
        """Convert list of tokens to list of IDs"""
        return [self.token_to_id(t) for t in tokens]
    
    def decode(self, ids: List[int], skip_special: bool = True) -> List[str]:
        """Convert list of IDs to list of tokens"""
        special_ids = set(range(len(self.SPECIAL_TOKENS))) if skip_special else set()
        return [self.id_to_token(i) for i in ids if i not in special_ids]
    
    def encode_text(self, text: str, add_special: bool = True, tokenize_fn=None) -> List[int]:
        """Tokenize and encode a full text string"""
        if tokenize_fn is None:
            tokenize_fn = lambda t: t.lower().split()
        tokens = tokenize_fn(text)
        ids = self.encode(tokens)
        if add_special:
            ids = [self.bos_id] + ids + [self.eos_id]
        return ids
    
    def decode_text(self, ids: List[int], skip_special: bool = True) -> str:
        """Decode IDs back to text string"""
        tokens = self.decode(ids, skip_special=skip_special)
        return " ".join(tokens)
    
    def __contains__(self, token: str) -> bool:
        return token in self._token_to_id
    
    def __len__(self) -> int:
        return self.size
    
    def __getitem__(self, token: str) -> int:
        return self.token_to_id(token)
    
    def get_frequency(self, token: str) -> int:
        """Get token frequency"""
        return self._token_freq.get(token, 0)
    
    def most_common(self, n: int = 10) -> List[Tuple[str, int]]:
        """Get most common tokens"""
        return self._token_freq.most_common(n)
    
    def save(self, filepath: str):
        """Save vocabulary to JSON file"""
        data = {
            "version": "1.0",
            "max_size": self.max_size,
            "min_frequency": self.min_frequency,
            "token_to_id": self._token_to_id,
            "token_freq": dict(self._token_freq),
        }
        os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else ".", exist_ok=True)
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    
    @classmethod
    def load(cls, filepath: str) -> 'Vocabulary':
        """Load vocabulary from JSON file"""
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        vocab = cls(
            max_size=data.get("max_size", 50000),
            min_frequency=data.get("min_frequency", 1)
        )
        vocab._token_to_id = {k: int(v) for k, v in data["token_to_id"].items()}
        vocab._id_to_token = {int(v): k for k, v in data["token_to_id"].items()}
        vocab._token_freq = Counter(data.get("token_freq", {}))
        return vocab
    
    def __repr__(self) -> str:
        return f"Vocabulary(size={self.size}, max_size={self.max_size})"
