import numpy as np
from typing import List, Dict

class CharTokenizer:
    """A simple character-level tokenizer for PotatoPLM."""
    
    def __init__(self, chars: List[str] = None):
        # We include a space and some common symbols by default if chars is empty
        if chars:
            self.chars = sorted(list(set(chars)))
        else:
            self.chars = []
        
        # Add special UNK token if not present
        if "[UNK]" not in self.chars:
            self.chars.insert(0, "[UNK]")
            
        self.vocab_size = len(self.chars)
        self.stoi = {ch: i for i, ch in enumerate(self.chars)}
        self.itos = {i: ch for i, ch in enumerate(self.chars)}

    def encode(self, s: str) -> List[int]:
        """Convert a string into a list of integers."""
        return [self.stoi.get(c, self.stoi["[UNK]"]) for c in s]

    def decode(self, l: List[int]) -> str:
        """Convert a list of integers into a string."""
        return ''.join([self.itos.get(i, "[UNK]") for i in l])

    @classmethod
    def from_text(cls, text: str):
        """Build character vocabulary from a given text."""
        chars = "".join(sorted(list(set(text))))
        return cls(chars)

    def save_vocab(self, path: str):
        """Save vocabulary to a file."""
        with open(path, 'w', encoding='utf-8') as f:
            f.write("".join(self.chars))

    @classmethod
    def load_vocab(cls, path: str):
        """Load vocabulary from a file."""
        with open(path, 'r', encoding='utf-8') as f:
            chars = f.read()
        return cls(chars)
