import pickle
from collections.abc import Iterable, Iterator
import math
import regex as re



class Tokenizer:
    def __init__(self, vocab: dict[int, bytes], merges: list[tuple[bytes,bytes]], special_tokens: list[str] | None = None): 
        """Construct a tokenizer from a given vocabulary, list of merges, and (optionally) a list of special tokens.
        This function should accept the following parameters:
        vocab: dict[int, bytes]  
        merges: list[tuple[bytes, bytes]]  
        special_tokens: list[str] | None = None  
        """
        self.vocab = vocab
        self.byte_to_id: dict[bytes, int] = {
            token: idx for idx, token in vocab.items()
        }
        self.merge_ranks: dict[tuple[bytes, bytes], int] = {
            merge: rank for rank, merge in enumerate(merges)
        }
        self.special_tokens = special_tokens or []
        if self.special_tokens:
            self.special_token_to_id: dict[str, int] = {
                token: self.byte_to_id[token.encode("utf-8")] for token in self.special_tokens 
            }
        self.PAT = re.compile(
            r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""  
        )

    @classmethod
    def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None):
        """Class method that constructs and returns a Tokenizer from a serialized vocabulary and list of merges (in the 
        same format that your BPE training code output) and (optionally) a list of special tokens. 
        This method should accept the following additional parameters:
        vocab_filepath: str  
        merges_filepath: str  
        special_tokens: list[str] | None = None  
        """
        raise NotImplementedError

    @classmethod      
    def from_file_with_pickle(cls, filepath, special_tokens=None) -> "Tokenizer":
        with open(filepath, "rb") as f:
            vocab, merges = pickle.load(f)
        return cls(vocab, merges, special_tokens)

    def encode(self, text: str) -> list[int]:
        """Encode an input text into a sequence of token IDs."""
        res: list[int] = []
        parts: list[str] = []

        if self.special_tokens:
            special_tokens = sorted(self.special_tokens, key = len, reverse= True)
            escaped_special_token = [re.escape(tok) for tok in special_tokens]
            pattern = "(" + "|".join(escaped_special_token) + ")"
            parts = re.split(pattern, text)
        else:
            parts = [text]

        for part in parts:
            if part == "":
                continue
            if part in self.special_tokens:
                res.append(self.special_token_to_id[part])
                continue
            for match in self.PAT.finditer(part):
                tokens: list[bytes] = []
                pre_token = match.group().encode("utf-8")
                for ch in pre_token:
                    tokens.append(bytes([ch]))

                while True:
                    merge, exists = self.find_small_rank(tokens)
                    if not exists:
                        break
                    new_tokens = []
                    i = 0
                    while i < len(tokens):
                        if i < len(tokens) - 1 and (tokens[i],tokens[i + 1]) == merge:
                            new_tokens.append(tokens[i] + tokens[i + 1])
                            i += 2
                        else:
                            new_tokens.append(tokens[i])
                            i += 1
                    tokens = new_tokens

                for i in range(0, len(tokens)):
                    res.append(self.byte_to_id[tokens[i]])
        return res
          

        
    def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
        """Given an iterable of strings (e.g., a Python file handle), return a generator that lazily yields token IDs.
        This is required for memory-efficient tokenization of large files that we cannot directly load into memory.
        """
        for text in iterable:
            yield from self.encode(text)

    
    def decode(self, ids: list[int]) -> str:
        """ Decode a sequence of token IDs into text.
        To test your Tokenizer against our provided tests, you will first need to implement the test 
        adapter at [adapters.get_tokenizer] . 
        """
        output_bytes: list[bytes] = []

        for id in ids:
            if id not in self.vocab:
                continue
            output_bytes.append(self.vocab[id])
        return b"".join(output_bytes).decode("utf-8",errors = "replace")

    def find_small_rank(self, tokens: list[bytes]) -> tuple[tuple[bytes, bytes],bool]:
        smallest_pair = (b"", b"")
        find = False
        rank = math.inf
        i = 0
        for i in range(0, len(tokens) - 1):
            pair = (tokens[i],tokens[i + 1])
            if pair not in self.merge_ranks:
                continue
            r = self.merge_ranks[pair]
            if r < rank:
                smallest_pair = pair
                rank = r
                find = True
        return (smallest_pair,find)