"""Chronos-style tokenizer: mean-scale a window, bin it, reserve two special ids."""

import torch

PAD_ID = 0
EOS_ID = 1
SPECIAL_TOKENS = 2

class Tokenizer:
    def __init__(self, n_bins: int= 4096, clip: float =15.0) -> None:
        self.n_bins = n_bins
        self.clip = clip
        self.bin_width = 2.0 * clip / n_bins

    @property
    def vocab_size(self) -> int:
        return self.n_bins + SPECIAL_TOKENS

    def centres(self, device: torch.device | None = None) -> torch.Tensor:
        half = self.bin_width / 2.0
        return torch.linspace(-self.clip + half, self.clip - half, self.n_bins, device=device)

    def mean_scale(self, x: torch.Tensor) -> torch.Tensor:
        # Gaps are NaN and must not poison the mean; a dead window scales by 1 so z stays finite.
        scale = torch.nanmean(x.abs(), dim=-1, keepdim=True)
        return torch.where(scale > 0.0, scale, torch.ones_like(scale))

    def encode(
            self, x: torch.Tensor, scale: torch.Tensor | None = None
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if scale is None:
            scale = self.mean_scale(x)
        gap = torch.isnan(x)
        z = torch.nan_to_num(x / scale, nan=0.0).clamp(-self.clip, self.clip)
        # z == clip lands on bin n_bins, one past the end, so the top bin absorbs it.
        bins = torch.floor((z + self.clip) / self.bin_width).long().clamp(max=self.n_bins - 1)
        ids = torch.where(gap, torch.full_like(bins, PAD_ID), bins + SPECIAL_TOKENS)
        return ids, scale

    def decode(self, ids: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
        centres = self.centres(device=ids.device)
        values = centres[(ids - SPECIAL_TOKENS).clamp(min=0)]* scale
        return torch.where(ids < SPECIAL_TOKENS, torch.full_like(values, float("nan")), values)