from transformers import PreTrainedTokenizer
from typing import List, Optional, Dict

# Number of symbolic graph-node ids "0".."{NUM_NODES-1}". Node id == token id, so the
# LM head is a direct classifier over node ids (the node-dictionary probe). A 2-arm star
# of depth L needs 2 + 4*L distinct nodes, so 100 nodes supports up to L=24.
NUM_NODES = 100


class STokenizer(PreTrainedTokenizer):
    def __init__(self, num_nodes=NUM_NODES):
        # Create vocabulary: "0" -> 0, "1" -> 1, ..., "{num_nodes-1}" -> num_nodes-1.
        # `num_nodes` is configurable so checkpoints trained with a different node
        # vocab (e.g. the 31-node L6 model) can be loaded/probed with a matching
        # tokenizer, while run.py / L20 default to the module-level NUM_NODES.
        self.vocab = {str(i): i for i in range(0, num_nodes)}
        n = num_nodes
        self.vocab['<|start-latent|>'] = n
        self.vocab['<|end-latent|>'] = n + 1
        self.vocab['<|latent|>'] = n + 2
        self.vocab['|'] = n + 3
        self.vocab['[Q]'] = n + 4
        self.vocab['[R]'] = n + 5
        self.vocab['[A]'] = n + 6
        # Add special tokens
        self.vocab['<eos>'] = n + 7
        self.vocab['<|no-answer|>'] = n + 8
        # Model configs often use vocab_size=128 (> tokenizer size 109). Generation can
        # emit those unused ids; decode must not crash (KeyError) on them.
        self.vocab['<unk>'] = n + 9
        
        # Create inverse vocabulary (id to token mapping)
        self.ids_to_tokens = {v: k for k, v in self.vocab.items()}
        
        # Set special token attributes
        super().__init__(
            pad_token="<eos>", eos_token="<eos>", bos_token="<eos>", unk_token="<unk>"
        )

    def get_vocab(self) -> Dict[str, int]:
        """Returns the vocabulary as a dict"""
        return self.vocab.copy()

    @property
    def vocab_size(self) -> int:
        return len(self.vocab)
        
    def _tokenize(self, text: str) -> List[str]:
        # Split on whitespace and validate each token is a number in range
        tokens = []
        for token in text.replace("\n", " ").strip().split():
            if token in self.vocab:
                tokens.append(token)
            else:
                raise ValueError(f"Token {token} not in vocabulary")
        return tokens
    
    def _convert_token_to_id(self, token: str) -> int:
        # Convert token to id, return unk_token_id if token not in vocab
        return self.vocab[token]
    
    def _convert_id_to_token(self, index: int) -> str:
        # Convert id back to token. Unknown ids (e.g. model LM head slots beyond
        # the symbolic vocab) map to <unk> instead of raising KeyError.
        return self.ids_to_tokens.get(int(index), "<unk>")
    
    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        # Join tokens with spaces
        return ' '.join(tokens)
    
    def build_inputs_with_special_tokens(self, token_ids_0: List[int], 
                                       token_ids_1: Optional[List[int]] = None) -> List[int]:
        # Add special tokens around sequence(s)
        if token_ids_1 is None:
            return token_ids_0 + [self.vocab['<eos>']]
        return token_ids_0 + [self.vocab['<eos>']] + token_ids_1 + [self.vocab['<eos>']]

    def get_special_tokens_mask(self, token_ids_0: List[int], 
                              token_ids_1: Optional[List[int]] = None,
                              already_has_special_tokens: bool = False) -> List[int]:
        # Return a mask indicating special tokens
        if already_has_special_tokens:
            return [1 if token_id in [self.vocab['<eos>'], self.vocab['<pad>'], self.vocab['<unk>']]
                   else 0 for token_id in token_ids_0]
        if token_ids_1 is None:
            return [0] * len(token_ids_0) + [1]
        return [0] * len(token_ids_0) + [1] + [0] * len(token_ids_1) + [1]