'Runtime plain-BPE vocabulary over the row-major flattened cell stream.'

from __future__ import annotations

import heapq
from collections.abc import Sequence
from dataclasses import dataclass

import numpy as np

from ._validation import active_dimensions
from .native_serialization import NativeCacheMixin, decode_native, encode_native

__all__ = ["BPEVocabulary"]


@dataclass
class BPEVocabulary(NativeCacheMixin):
    """Primitive atoms plus ordered, adjacency-based merge rules."""

    bins: tuple[int, ...]
    merges: tuple[tuple[int, int], ...]

    def __post_init__(self) -> None:
        if not self.bins or any(count < 1 for count in self.bins):
            raise ValueError("every dimension needs at least one bin")
        self.offsets = np.concatenate(([0], np.cumsum(np.asarray(self.bins, dtype=np.int64))))
        self.atom_count = int(self.offsets[-1])
        self.vocab_size = self.atom_count + len(self.merges)
        atom_dims = np.repeat(
            np.arange(len(self.bins), dtype=np.int64), np.asarray(self.bins, dtype=np.int64)
        )


        expansions: list[tuple[int, ...]] = [(token,) for token in range(self.atom_count)]
        for rank, (left, right) in enumerate(self.merges):
            token = self.atom_count + rank
            if not 0 <= left < token or not 0 <= right < token:
                raise ValueError("a merge rule must reference already-created tokens")
            expansions.append(expansions[left] + expansions[right])
        self.expansions = tuple(expansions)
        self.token_dims = tuple(
            tuple(int(atom_dims[atom]) for atom in expansion) for expansion in expansions
        )
        self.token_cells = np.asarray([len(expansion) for expansion in expansions], dtype=np.int64)
        self._rules = {
            (int(left), int(right)): (self.atom_count + rank, rank)
            for rank, (left, right) in enumerate(self.merges)
        }

    @property
    def dimension(self) -> int:
        return len(self.bins)

    def atom(self, dimension: int, value: int) -> int:
        return int(self.offsets[dimension]) + int(value)

    def _flatten(self, values: np.ndarray, dimensions: tuple[int, ...]) -> list[int]:
        
        counts = np.asarray(self.bins, dtype=np.int64)
        if np.any(values < -1) or np.any(values >= counts[None, :]):
            raise ValueError("cells contain an invalid active or absent index")
        support = values >= 0
        if np.any(support != support[:1]):
            raise ValueError("active dimensions must remain constant within a trajectory")
        present = tuple(int(dim) for dim in np.flatnonzero(support[0]))
        if present != dimensions:
            raise ValueError("cells do not activate exactly the requested dimensions")
        if not present:
            raise ValueError("every control step needs at least one active coordinate")
        active = values[:, list(present)]
        base = np.asarray([int(self.offsets[dim]) for dim in present], dtype=np.int64)
        return (active + base[None, :]).reshape(-1).tolist()

    def encode(self, cells: np.ndarray) -> list[int]:
        'Apply the merges greedily in creation-rank order to one trajectory.'
        values = np.asarray(cells, dtype=np.int64)
        if values.ndim != 2 or len(values) < 1 or values.shape[1] < 1:
            raise ValueError("cells must have non-empty shape [T, D]")
        dimensions = tuple(int(dim) for dim in np.flatnonzero(values[0] >= 0))
        stream = self._flatten(values, dimensions)
        result = encode_native(self, values[None])
        if result is not None:
            return result[0]
        count = len(stream)
        previous = list(range(-1, count - 1))
        following = list(range(1, count + 1))
        following[-1] = -1
        alive = [True] * count
        rules = self._rules
        heap: list[tuple[int, int]] = []
        for index in range(count - 1):
            rule = rules.get((stream[index], stream[index + 1]))
            if rule is not None:
                heap.append((rule[1], index))
        heapq.heapify(heap)
        while heap:
            rank, index = heapq.heappop(heap)
            if not alive[index]:
                continue
            follower = following[index]
            if follower < 0 or not alive[follower]:
                continue
            rule = rules.get((stream[index], stream[follower]))
            if rule is None or rule[1] != rank:
                continue
            stream[index] = rule[0]
            alive[follower] = False
            after = following[follower]
            following[index] = after
            if after >= 0:
                previous[after] = index
                successor = rules.get((stream[index], stream[after]))
                if successor is not None:
                    heapq.heappush(heap, (successor[1], index))
            before = previous[index]
            if before >= 0:
                predecessor = rules.get((stream[before], stream[index]))
                if predecessor is not None:
                    heapq.heappush(heap, (predecessor[1], before))
        return [stream[index] for index in range(count) if alive[index]]

    def decode(
        self,
        tokens: Sequence[int],
        *,
        active_dims: Sequence[int] | None = None,
    ) -> np.ndarray:
        """Expand tokens to dense cells, writing ``-1`` in absent coordinates."""
        dimensions = active_dimensions(active_dims, self.dimension)
        result = decode_native(self, [[int(token) for token in tokens]], dimensions)
        if result is not None:
            return result[0]
        atoms: list[int] = []
        for index, raw_token in enumerate(tokens):
            token = int(raw_token)
            if not 0 <= token < self.vocab_size:
                raise ValueError(f"tokens[{index}] is outside the vocabulary")
            atoms.extend(self.expansions[token])
        width = len(dimensions)
        if not atoms or len(atoms) % width:
            raise ValueError("token stream does not cover a whole number of routed control steps")
        values = np.asarray(atoms, dtype=np.int64).reshape(-1, width)
        starts = np.asarray([int(self.offsets[dim]) for dim in dimensions], dtype=np.int64)
        stops = np.asarray([int(self.offsets[dim + 1]) for dim in dimensions], dtype=np.int64)
        if np.any(values < starts[None, :]) or np.any(values >= stops[None, :]):
            raise ValueError("token stream does not follow the routed dimension order")
        cells = np.full((len(values), self.dimension), -1, dtype=np.int64)
        cells[:, list(dimensions)] = values - starts[None, :]
        return cells