"""Runtime spatial Set-BPE vocabulary and lossless row serialization."""

from __future__ import annotations

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

import numpy as np

from ..common._validation import active_dimensions

__all__ = ["SpatialVocabulary", "aggregate_records"]


def aggregate_records(
    codes: np.ndarray, bins: Sequence[int] | None = None
) -> tuple[np.ndarray, np.ndarray]:
    """Lossless micro-aggregation of identical quantized actions.

    Replacing ``w`` identical records by one record of weight ``w`` leaves every
    pair frequency and therefore every greedy decision unchanged, because
    identical actions follow identical merges forever.
    """
    values = np.asarray(codes, dtype=np.int64)
    if values.ndim != 2:
        raise ValueError("codes must have shape [U, D]")
    if bins is not None:
        dimensions = tuple(int(count) for count in bins)
        if len(dimensions) != values.shape[1]:
            raise ValueError("bins must have one entry per code coordinate")
        counts = np.asarray(dimensions, dtype=np.int64)
        valid = values >= 0
        if np.any(values < -1) or np.any(values >= counts[None, :]):
            raise ValueError("codes must contain -1 or an in-range cell index")
        cardinality = int(np.prod(np.asarray(dimensions, dtype=object)))
        if np.all(valid) and cardinality <= np.iinfo(np.int64).max:
            keys = np.ravel_multi_index(values.T, dimensions)
            unique_keys, counts = np.unique(keys, return_counts=True)
            unique = np.stack(np.unravel_index(unique_keys, dimensions), axis=1)
            return unique.astype(np.int64, copy=False), counts.astype(np.int64)
    unique, counts = np.unique(values, axis=0, return_counts=True)
    return unique, counts.astype(np.int64)


@dataclass
class SpatialVocabulary:
    """Atoms, merge rules and the decoder metadata they induce."""

    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)
        dimension = len(self.bins)
        support = np.zeros(self.vocab_size, dtype=np.int64)
        expansion = np.full((self.vocab_size, dimension), -1, dtype=np.int64)
        for index, count in enumerate(self.bins):
            for value in range(count):
                token = int(self.offsets[index]) + value
                support[token] = 1 << index
                expansion[token, index] = value
        for rank, (left, right) in enumerate(self.merges):
            token = self.atom_count + rank
            if left >= token or right >= token:
                raise ValueError("a merge rule must reference already-created tokens")
            if support[left] & support[right]:
                raise ValueError("a merge rule joins two overlapping supports")
            support[token] = support[left] | support[right]
            merged = np.maximum(expansion[left], expansion[right])
            expansion[token] = merged
        self.support = support
        self.expansion = expansion
        self.anchor = np.asarray(
            [int(mask & -mask).bit_length() - 1 for mask in support], dtype=np.int64
        )
        self.full_support = (1 << dimension) - 1
        self._rules: dict[tuple[int, int], tuple[int, int]] = {}
        for rank, (left, right) in enumerate(self.merges):
            key = (min(left, right), max(left, right))
            self._rules[key] = (self.atom_count + rank, rank)

    @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 encode(
        self,
        code: Sequence[int],
        *,
        active_dims: Sequence[int] | None = None,
    ) -> list[int]:
        """Deterministic online Set-BPE encoding (Algorithm 8).

        Merge rules are applied in creation-rank order, which is what makes the
        encoder reproducible: the same quantized action always yields the same
        token set regardless of how the search that learned the rules ran.
        ``-1`` in a dense row, or a missing coordinate in ``active_dims``, means
        semantic absence: it emits no atom and can never participate in a merge.
        """
        compact = [int(value) for value in code]
        if active_dims is None:
            if len(compact) != self.dimension:
                raise ValueError(
                    f"a dense quantized action must have {self.dimension} slots"
                )
            pairs = enumerate(compact)
        else:
            dimensions = active_dimensions(active_dims, self.dimension)
            if len(dimensions) != len(compact):
                raise ValueError(
                    "active_dims must have one entry per compact coordinate"
                )
            pairs = zip(dimensions, compact, strict=True)
        tokens = set()
        for index, value in pairs:
            if value == -1:
                if active_dims is not None:
                    raise ValueError("compact coordinates must not contain -1")
                continue
            if not 0 <= value < self.bins[index]:
                raise ValueError(f"bin {value} is out of range for dimension {index}")
            tokens.add(self.atom(index, value))
        if not tokens:
            raise ValueError(
                "a quantized action must contain at least one active coordinate"
            )
        heap: list[tuple[int, int, int, int]] = []
        ordered = sorted(tokens)
        for position, left in enumerate(ordered):
            for right in ordered[position + 1 :]:
                rule = self._rules.get((left, right))
                if rule is not None:
                    heapq.heappush(heap, (rule[1], rule[0], left, right))
        while heap:
            _, child, left, right = heapq.heappop(heap)
            if left not in tokens or right not in tokens:
                continue
            tokens.discard(left)
            tokens.discard(right)
            for other in tokens:
                key = (min(child, other), max(child, other))
                rule = self._rules.get(key)
                if rule is not None:
                    heapq.heappush(heap, (rule[1], rule[0], key[0], key[1]))
            tokens.add(child)
        return self.serialize(tokens)

    def serialize(self, tokens: Iterable[int]) -> list[int]:
        """Canonical order for autoregressive supervision.

        Set semantics carry no order, but a policy must be trained on one fixed
        target.  Tokens within an action have pairwise disjoint supports, so
        their anchor dimensions are distinct and sorting by anchor is a total
        order.  Anchor order is preferred over raw token id because it is the
        one that keeps constrained decoding feasible: the next token's anchor is
        always the lowest uncovered dimension, and a primitive atom for that
        dimension always exists, so no prefix can paint itself into a corner.
        """
        return sorted(tokens, key=lambda token: (int(self.anchor[token]), int(token)))

    def decode(
        self,
        tokens: Sequence[int],
        *,
        active_dims: Sequence[int] | None = None,
    ) -> np.ndarray:
        """Expand tokens to a dense row, using ``-1`` for absent coordinates."""
        code = np.full(self.dimension, -1, dtype=np.int64)
        covered = 0
        for token in tokens:
            value = int(token)
            if not 0 <= value < self.vocab_size:
                raise ValueError(f"token {value} is outside the vocabulary")
            if covered & int(self.support[value]):
                raise ValueError("token supports overlap; this is not a valid action")
            covered |= int(self.support[value])
            local = self.expansion[value]
            code = np.where(local >= 0, local, code)
        if active_dims is None:
            expected = self.full_support
        else:
            dimensions = active_dimensions(active_dims, self.dimension)
            expected = sum(1 << dimension for dimension in dimensions)
        if covered != expected:
            raise ValueError(
                "token set does not exactly cover the requested active dimensions"
            )
        return code