# Based on togethercomputer/evo-1-131k-base/tokenizer.py (Apache 2.0).
# Adapted: minor fix (typing), Dict -> dict for forward-compatibility.
"""ByteTokenizer for Evo1.

Maps text to/from raw UTF-8 byte values, vocab_size = 512 (no real vocab,
just the byte range padded out).
"""
from __future__ import annotations

from os import PathLike
from typing import List, Tuple

import numpy as np

from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import BatchEncoding, TruncationStrategy
from transformers.utils.generic import PaddingStrategy


EMPTY: str = ""


class ByteTokenizer(PreTrainedTokenizer):
    """UTF-8 byte-level tokenizer (vocab_size = 512).

    Special tokens follow the original ``CharLevelTokenizer`` convention:
      * ``\x00`` (chr(0)) is end-of-document / end-of-sequence.
      * ``\x01`` (chr(1)) is the padding token.

    These are wired up so ``HF`` tokenization helpers (``padding=True``,
    ``model_max_length``, etc.) work with this tokenizer.
    """

    def __init__(self, byte_level: bool = True, **kwargs):
        # Defaults; ``tokenizer_config.json`` may override via kwargs.
        # Padding token: byte 0x01 (matches the original CharLevelTokenizer
        # pad convention). We deliberately do NOT set eos_token / cls_token:
        # Evo1 is a pure byte-level model with no special tokens added at
        # encoding time, so downstream pooling logic (e.g. mRNABench) should
        # treat every non-pad position as a real token.
        kwargs.setdefault("pad_token", chr(1))
        super().__init__(byte_level=byte_level, **kwargs)
        # The model only consumes input_ids and attention_mask (no segment ids).
        self.model_input_names = ["input_ids", "attention_mask"]

    @property
    def vocab_size(self) -> int:
        return 512

    @property
    def byte_level(self) -> bool:
        return self.init_kwargs.get("byte_level", True)

    def get_vocab(self) -> dict:
        return {chr(i): i for i in range(self.vocab_size)}

    def __len__(self) -> int:
        return self.vocab_size

    def clamp(self, n: int) -> int:
        return max(32, min(n, self.vocab_size))

    def _tokenize(self, text: str, **kwargs) -> List[str]:
        return list(text)

    def byte_tokenize(self, text: str) -> np.ndarray:
        return np.frombuffer(text.encode("utf-8"), dtype=np.uint8)

    def _convert_token_to_id(self, token: str) -> int:
        return self.clamp(ord(token))

    def _convert_id_to_token(self, index: int) -> str:
        return chr(self.clamp(index))

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        return EMPTY.join(tokens)

    def _decode(self, token_ids: List[int], **kwargs) -> str:
        # numpy >= 2 disallows clipping a uint8 array to a value outside [0, 255]
        # (vocab_size=512). Convert to int16 first, clip, then back to uint8.
        indices = np.asarray(token_ids, dtype=np.int16)
        indices = indices.clip(min=32, max=255).astype(np.uint8)
        return indices.tobytes().decode("utf-8", errors="replace")

    def _encode_plus(self, text: str, **kwargs) -> BatchEncoding:
        first_ids = self.byte_tokenize(text).tolist()
        return self.prepare_for_model(
            first_ids,
            pair_ids=None,
            add_special_tokens=kwargs.get("add_special_tokens", False),
            padding=kwargs.get("padding_strategy", PaddingStrategy.DO_NOT_PAD).value,
            truncation=kwargs.get("truncation_strategy", TruncationStrategy.DO_NOT_TRUNCATE).value,
            max_length=kwargs.get("max_length"),
            stride=kwargs.get("stride", 0),
            pad_to_multiple_of=kwargs.get("pad_to_multiple_of"),
            return_tensors=kwargs.get("return_tensors"),
            prepend_batch_axis=True,
            return_attention_mask=kwargs.get("return_attention_mask"),
            return_token_type_ids=kwargs.get("return_token_type_ids"),
            return_overflowing_tokens=kwargs.get("return_overflowing_tokens", False),
            return_special_tokens_mask=kwargs.get("return_special_tokens_mask", False),
            return_length=kwargs.get("return_length", False),
            verbose=kwargs.get("verbose", True),
        )

    def _batch_encode_plus(self, batch_text_or_text_pairs, **kwargs) -> BatchEncoding:
        input_ids = [(self.byte_tokenize(t).tolist(), None) for t in batch_text_or_text_pairs]
        return self._batch_prepare_for_model(
            input_ids,
            add_special_tokens=kwargs.get("add_special_tokens", False),
            padding_strategy=kwargs.get("padding_strategy", PaddingStrategy.DO_NOT_PAD),
            truncation_strategy=kwargs.get("truncation_strategy", TruncationStrategy.DO_NOT_TRUNCATE),
            max_length=kwargs.get("max_length"),
            stride=kwargs.get("stride", 0),
            pad_to_multiple_of=kwargs.get("pad_to_multiple_of"),
            return_attention_mask=kwargs.get("return_attention_mask"),
            return_token_type_ids=kwargs.get("return_token_type_ids"),
            return_overflowing_tokens=kwargs.get("return_overflowing_tokens", False),
            return_special_tokens_mask=kwargs.get("return_special_tokens_mask", False),
            return_length=kwargs.get("return_length", False),
            return_tensors=kwargs.get("return_tensors"),
            verbose=kwargs.get("verbose", True),
        )

    def _save_pretrained(
        self, save_directory: str | PathLike, file_names: Tuple[str], **kwargs
    ) -> Tuple[str]:
        return file_names