"""Pruned-vocab id remap for ALF-emb-micro."""

from __future__ import annotations

from pathlib import Path


class RemapTokenizer:
    def __init__(self, tok_dir: str | Path, keep_ids: list[int], unk_new: int | None = None):
        from transformers import AutoTokenizer
        import numpy as np
        import torch

        self.base = AutoTokenizer.from_pretrained(str(tok_dir))
        self.map = {old: new for new, old in enumerate(keep_ids)}
        unk_old = self.base.unk_token_id
        self.unk_new = unk_new if unk_new is not None else self.map.get(unk_old, 0)
        vocab = int(getattr(self.base, "vocab_size", 0) or 0)
        size = max(vocab, max(keep_ids) + 1 if keep_ids else 1)
        lut = np.full(size, self.unk_new, dtype=np.int64)
        for old, new in self.map.items():
            if 0 <= old < size:
                lut[old] = new
        self._lut_np = lut
        self._lut = torch.from_numpy(lut)

    def __call__(self, texts, **kwargs):
        enc = self.base(texts, **kwargs)
        ids = enc["input_ids"]
        if hasattr(ids, "clamp"):
            lut = self._lut.to(device=ids.device)
            safe = ids.clamp(0, lut.numel() - 1).long()
            enc["input_ids"] = lut[safe].to(dtype=ids.dtype)
        else:
            import numpy as np

            arr = np.asarray(ids, dtype=np.int64)
            mapped = self._lut_np[np.clip(arr, 0, len(self._lut_np) - 1)]
            tensors = kwargs.get("return_tensors")
            enc["input_ids"] = mapped if tensors in ("np", "pt") or hasattr(ids, "shape") else mapped.tolist()
        return enc