"""
Custom Byte-Pair Encoding (BPE) & Subword Tokenizer from Scratch.
Supports training on raw text / PDF corpora, vocabulary persistence, and document special tokens.
"""
import json
import re
from typing import List, Dict, Tuple, Optional, Set
from collections import defaultdict


SPECIAL_TOKENS = [
    "<pad>",    # ID: 0 (Padding)
    "<unk>",    # ID: 1 (Unknown)
    "<bos>",    # ID: 2 (Beginning of sequence)
    "<eos>",    # ID: 3 (End of sequence)
    "<doc>",    # ID: 4 (Document context start)
    "<query>",  # ID: 5 (Query start)
    "<ans>",    # ID: 6 (Answer start)
    "<sep>"     # ID: 7 (Separator)
]


class CustomBPETokenizer:
    """
    Byte-Pair Encoding Tokenizer implemented from first principles.
    Learns iterative frequent character/subword merges from training text.
    """
    def __init__(self, special_tokens: Optional[List[str]] = None):
        self.special_tokens = special_tokens or SPECIAL_TOKENS
        self.special_token_to_id = {tok: i for i, tok in enumerate(self.special_tokens)}
        self.id_to_special_token = {i: tok for i, tok in enumerate(self.special_tokens)}

        self.vocab: Dict[str, int] = {}
        self.id_to_token: Dict[int, str] = {}
        self.merges: List[Tuple[str, str]] = []

        # Special token IDs as convenient properties
        self.pad_id = self.special_token_to_id["<pad>"]
        self.unk_id = self.special_token_to_id["<unk>"]
        self.bos_id = self.special_token_to_id["<bos>"]
        self.eos_id = self.special_token_to_id["<eos>"]
        self.doc_id = self.special_token_to_id["<doc>"]
        self.query_id = self.special_token_to_id["<query>"]
        self.ans_id = self.special_token_to_id["<ans>"]
        self.sep_id = self.special_token_to_id["<sep>"]

    @property
    def vocab_size(self) -> int:
        return len(self.vocab)

    def _get_stats(self, words: Dict[Tuple[str, ...], int]) -> Dict[Tuple[str, str], int]:
        """Count frequencies of adjacent symbol pairs in word representations."""
        pairs = defaultdict(int)
        for word, freq in words.items():
            for i in range(len(word) - 1):
                pairs[(word[i], word[i + 1])] += freq
        return pairs

    def _merge_pair(self, pair: Tuple[str, str], words: Dict[Tuple[str, ...], int]) -> Dict[Tuple[str, ...], int]:
        """Merge all occurrences of a specific symbol pair across word representations."""
        new_words = {}
        bigram = pair
        first, second = bigram

        for word, freq in words.items():
            new_word = []
            i = 0
            while i < len(word):
                if i < len(word) - 1 and word[i] == first and word[i + 1] == second:
                    new_word.append(first + second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_words[tuple(new_word)] = freq
        return new_words

    def train(self, texts: List[str], target_vocab_size: int = 1024, min_frequency: int = 2):
        """
        Train the BPE tokenizer on a list of texts up to target_vocab_size.
        """
        # Step 1: Initialize vocabulary with special tokens
        self.vocab = {}
        self.id_to_token = {}

        for tok in self.special_tokens:
            idx = len(self.vocab)
            self.vocab[tok] = idx
            self.id_to_token[idx] = tok

        # Step 2: Extract base character frequencies from text
        # Regex to split on whitespace and punctuation while keeping words intact
        word_freqs = defaultdict(int)
        for text in texts:
            # Normalize whitespace
            clean_text = re.sub(r"\s+", " ", text).strip()
            # Split into words/tokens
            words = clean_text.split(" ")
            for w in words:
                if w:
                    # Append end-of-word marker </w> to distinguish word boundaries
                    chars = tuple(list(w) + ["</w>"])
                    word_freqs[chars] += 1

        # Add all unique single characters to initial vocabulary
        unique_chars = set()
        for word_tuple in word_freqs.keys():
            for char in word_tuple:
                unique_chars.add(char)

        for char in sorted(list(unique_chars)):
            if char not in self.vocab:
                idx = len(self.vocab)
                self.vocab[char] = idx
                self.id_to_token[idx] = char

        # Step 3: Iteratively find and merge the most frequent adjacent pair
        self.merges = []
        words = dict(word_freqs)

        num_merges = target_vocab_size - len(self.vocab)
        for _ in range(max(0, num_merges)):
            pairs = self._get_stats(words)
            if not pairs:
                break

            # Find most frequent pair
            best_pair = max(pairs, key=pairs.get)
            if pairs[best_pair] < min_frequency and len(self.vocab) >= min(target_vocab_size, 256):
                break

            words = self._merge_pair(best_pair, words)
            self.merges.append(best_pair)

            new_token = best_pair[0] + best_pair[1]
            if new_token not in self.vocab:
                idx = len(self.vocab)
                self.vocab[new_token] = idx
                self.id_to_token[idx] = new_token

            if len(self.vocab) >= target_vocab_size:
                break

    def _tokenize_word(self, word: str) -> List[str]:
        """Apply learned merges to segment a single word into subword tokens."""
        symbols = list(word) + ["</w>"]
        for first, second in self.merges:
            i = 0
            new_symbols = []
            while i < len(symbols):
                if i < len(symbols) - 1 and symbols[i] == first and symbols[i + 1] == second:
                    new_symbols.append(first + second)
                    i += 2
                else:
                    new_symbols.append(symbols[i])
                    i += 1
            symbols = new_symbols
        return symbols

    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False) -> List[int]:
        """
        Encode a string of text into a list of token IDs.
        Properly parses special tokens (like <doc>, <query>, <ans>, <eos>).
        """
        token_ids = []
        if add_bos:
            token_ids.append(self.bos_id)

        # Regex to split on special tokens while keeping them
        special_pattern = "(" + "|".join(re.escape(tok) for tok in self.special_tokens) + ")"
        parts = re.split(special_pattern, text)

        for part in parts:
            if not part:
                continue
            if part in self.special_token_to_id:
                token_ids.append(self.special_token_to_id[part])
            else:
                words = part.split()
                for w in words:
                    subwords = self._tokenize_word(w)
                    for sw in subwords:
                        if sw in self.vocab:
                            token_ids.append(self.vocab[sw])
                        else:
                            # Fallback to individual characters
                            for char in sw:
                                token_ids.append(self.vocab.get(char, self.unk_id))

        if add_eos:
            token_ids.append(self.eos_id)

        return token_ids

    def decode(self, token_ids: List[int], skip_special_tokens: bool = False) -> str:
        """
        Decode a list of token IDs back into human-readable text.
        """
        tokens = []
        for tid in token_ids:
            if tid in self.id_to_special_token:
                if not skip_special_tokens:
                    tokens.append(self.id_to_special_token[tid] + " ")
            elif tid in self.id_to_token:
                tok = self.id_to_token[tid]
                if tok.endswith("</w>"):
                    tokens.append(tok[:-4] + " ")
                else:
                    tokens.append(tok)
            else:
                if not skip_special_tokens:
                    tokens.append("<unk> ")

        text = "".join(tokens).strip()
        # Clean spacing around punctuation
        text = re.sub(r"\s+([.,!?;:])", r"\1", text)
        return text

    def save(self, filepath: str):
        """Save tokenizer vocabulary and merges to JSON."""
        data = {
            "special_tokens": self.special_tokens,
            "vocab": self.vocab,
            "merges": self.merges
        }
        with open(filepath, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    @classmethod
    def load(cls, filepath: str) -> "CustomBPETokenizer":
        """Load tokenizer vocabulary and merges from JSON."""
        with open(filepath, "r", encoding="utf-8") as f:
            data = json.load(f)

        tokenizer = cls(special_tokens=data.get("special_tokens", SPECIAL_TOKENS))
        tokenizer.vocab = data["vocab"]
        tokenizer.id_to_token = {int(v): k for k, v in tokenizer.vocab.items()}
        tokenizer.merges = [tuple(m) for m in data.get("merges", [])]
        return tokenizer