"""Spatial Set-BPE fitting implementations.

Runtime vocabulary and serialization live in ``spatial``; this module is
loaded only by training entry points.
"""

from __future__ import annotations

import heapq
from collections.abc import Sequence

import numpy as np

from ..backend.native_backend import require_native, resolve_backend
from .setbpe_spatial import SpatialVocabulary

__all__ = ["train_spatial"]


class _SparseSpatialTrainer:
    """Columnar Set-BPE state with counters for observed pairs only."""

    def __init__(
        self,
        codes: np.ndarray,
        weights: np.ndarray,
        bins: Sequence[int],
    ) -> None:
        self.bins = tuple(int(value) for value in bins)
        self.dimension = len(self.bins)
        self.records = np.asarray(codes, dtype=np.int64)
        self.weights = np.asarray(weights, dtype=np.int64)
        if self.records.ndim != 2 or self.records.shape[1] != self.dimension:
            raise ValueError("codes must have shape [U, D] matching the bin counts")
        if self.weights.shape != (self.records.shape[0],):
            raise ValueError("weights must have one entry per record")
        counts = np.asarray(self.bins, dtype=np.int64)
        if np.any(self.records < -1) or np.any(self.records >= counts[None, :]):
            raise ValueError("codes must contain -1 or an in-range cell index")
        if np.any(np.all(self.records < 0, axis=1)):
            raise ValueError("every record must contain at least one active coordinate")
        self.count = int(self.records.shape[0])
        self.offsets = np.concatenate(
            ([0], np.cumsum(np.asarray(self.bins, dtype=np.int64)))
        )
        self.atom_count = int(self.offsets[-1])

        self.occurrence: list[np.ndarray] = []
        self.support: list[int] = []
        self.anchor: list[int] = []
        self.frequency: list[int] = []
        for index in range(self.dimension):
            column = self.records[:, index]
            for value in range(self.bins[index]):
                members = np.flatnonzero(column == value).astype(np.int64)
                self._append_token(members, 1 << index, index)
        self._scratch = np.zeros(self.count, dtype=bool)
        self.merges: list[tuple[int, int]] = []
        self.statistics = {
            "sparse_updates": 0,
            "sparse_heap_rebuilds": 0,
            "sparse_pair_types_peak": 0,
        }
        # Row-wise mirror used to initialize and incrementally update pair counts.
        self.member = np.where(
            self.records >= 0,
            self.offsets[:-1][None, :] + self.records,
            -1,
        )
        self._slot_pairs = tuple(
            (first, second)
            for first in range(self.dimension)
            for second in range(first + 1, self.dimension)
        )
        # The observed-pair backend is built on demand.  Keys use two fixed
        # 32-bit token fields, so appending vocabulary entries never invalidates
        # the counter map as a width-based packed index would.
        self._sparse_counts: dict[int, int] | None = None
        self._sparse_heap: list[tuple[int, int, int]] = []

    def _append_token(self, members: np.ndarray, support: int, anchor: int) -> int:
        token = len(self.occurrence)
        self.occurrence.append(members)
        self.support.append(support)
        self.anchor.append(anchor)
        self.frequency.append(int(self.weights[members].sum()) if members.size else 0)
        return token

    def _intersect(self, left: np.ndarray, right: np.ndarray) -> np.ndarray:
        if left.size > right.size:
            left, right = right, left
        if left.size == 0:
            return left
        self._scratch[left] = True
        result = right[self._scratch[right]]
        self._scratch[left] = False
        return result

    @staticmethod
    def _sparse_key(left: int, right: int) -> int:
        low, high = (left, right) if left < right else (right, left)
        return (int(low) << 32) | int(high)

    @staticmethod
    def _sparse_pair(key: int) -> tuple[int, int]:
        return key >> 32, key & 0xFFFFFFFF

    def build_sparse_index(self) -> None:
        """Materialize exact counters only for currently observed pairs once."""
        if self._sparse_counts is not None:
            return
        width = len(self.occurrence)
        counts = np.zeros(width * width, dtype=np.int64)
        for first, second in self._slot_pairs:
            left = self.member[:, first]
            right = self.member[:, second]
            valid = (left >= 0) & (right >= 0)
            if not valid.any():
                continue
            low = np.minimum(left[valid], right[valid])
            high = np.maximum(left[valid], right[valid])
            np.add.at(counts, low * width + high, self.weights[valid])
        observed = np.flatnonzero(counts)
        sparse: dict[int, int] = {}
        heap: list[tuple[int, int, int]] = []
        for packed in observed:
            left, right = divmod(int(packed), width)
            key = self._sparse_key(left, right)
            value = int(counts[packed])
            sparse[key] = value
            heap.append((-value, left, right))
        heapq.heapify(heap)
        self._sparse_counts = sparse
        self._sparse_heap = heap
        self.statistics["sparse_pair_types_peak"] = len(sparse)

    def sparse_best_pair(self) -> tuple[int, int, int] | None:
        """Return the exact winner from the observed-pair max heap."""
        if self._sparse_counts is None:
            self.build_sparse_index()
        assert self._sparse_counts is not None
        while self._sparse_heap:
            negative, left, right = heapq.heappop(self._sparse_heap)
            value = self._sparse_counts.get(self._sparse_key(left, right), 0)
            if value == -negative:
                return left, right, value
        return None

    def _update_sparse_index(
        self,
        left: int,
        right: int,
        child: int,
        rows: np.ndarray,
        members: np.ndarray,
        mass: int,
    ) -> None:
        if self._sparse_counts is None:
            return
        counts = self._sparse_counts
        heap = self._sparse_heap
        updates = 0

        def change(key: int, delta: int) -> None:
            nonlocal updates
            new = counts.get(key, 0) + delta
            if new < 0:
                pair = self._sparse_pair(key)
                raise RuntimeError(
                    f"sparse spatial pair counter became negative for {pair}"
                )
            if new:
                counts[key] = new
                pair = self._sparse_pair(key)
                heapq.heappush(heap, (-new, pair[0], pair[1]))
            else:
                counts.pop(key, None)
            updates += 1

        change(self._sparse_key(left, right), -mass)
        other_mass = np.zeros(len(self.occurrence), dtype=np.int64)
        weights = self.weights[members]
        for slot in range(rows.shape[1]):
            values = rows[:, slot]
            valid = (values >= 0) & (values != left) & (values != right)
            if valid.any():
                np.add.at(other_mass, values[valid], weights[valid])
        for raw_other in np.flatnonzero(other_mass):
            other = int(raw_other)
            value = int(other_mass[other])
            change(self._sparse_key(left, other), -value)
            change(self._sparse_key(right, other), -value)
            change(self._sparse_key(child, other), value)

        self.statistics["sparse_updates"] += updates
        self.statistics["sparse_pair_types_peak"] = max(
            self.statistics["sparse_pair_types_peak"], len(counts)
        )
        if len(heap) > 2 * len(counts) + 65_536:
            self._sparse_heap = [
                (-value, *self._sparse_pair(key)) for key, value in counts.items()
            ]
            heapq.heapify(self._sparse_heap)
            self.statistics["sparse_heap_rebuilds"] += 1

    def merge(self, left: int, right: int) -> int:
        """Apply one exact incremental merge update."""
        members = self._intersect(self.occurrence[left], self.occurrence[right])
        mass = int(self.weights[members].sum()) if members.size else 0
        rows = self.member[members] if members.size else self.member[:0].copy()
        self.occurrence[left] = np.setdiff1d(
            self.occurrence[left], members, assume_unique=True
        )
        self.occurrence[right] = np.setdiff1d(
            self.occurrence[right], members, assume_unique=True
        )
        self.frequency[left] -= mass
        self.frequency[right] -= mass
        token = self._append_token(
            members,
            self.support[left] | self.support[right],
            min(self.anchor[left], self.anchor[right]),
        )
        self._update_sparse_index(left, right, token, rows, members, mass)
        # Keep the row mirror in step.  Slot order inside a row carries no
        # meaning, so the child simply takes the left parent's slot.
        if members.size:
            rows[rows == right] = -1
            rows[rows == left] = token
            self.member[members] = rows
        self.merges.append((left, right))
        return token


def _train_spatial_python(
    codes: np.ndarray,
    weights: np.ndarray,
    bins: Sequence[int],
    *,
    vocab_size: int,
    progress: bool = False,
) -> tuple[SpatialVocabulary, dict[str, object]]:
    """Train Set-BPE with exact counters for observed row-local pairs.

    ``vocab_size`` is the total budget including the ``sum(bins)`` primitive
    atoms.  Training stops early when no positive-frequency compatible pair
    remains, which happens once every record has collapsed to a single token.

    The one-time counter build costs one dense pass; later work is limited to
    records changed by the committed merge.
    """
    trainer = _SparseSpatialTrainer(codes, weights, bins)
    if vocab_size < trainer.atom_count:
        raise ValueError(
            f"vocab_size {vocab_size} is below the {trainer.atom_count} primitive atoms"
        )
    total_weight = int(trainer.weights.sum())
    active_width = np.sum(trainer.records >= 0, axis=1, dtype=np.int64)
    initial = int(np.dot(trainer.weights, active_width))
    length = initial
    trainer.build_sparse_index()
    for step in range(vocab_size - trainer.atom_count):
        winner = trainer.sparse_best_pair()
        if winner is None:
            break
        left, right, frequency = winner
        trainer.merge(left, right)
        length -= frequency
        if progress and step % 256 == 0:
            print(
                f"  merge {step:>5}  tokens/action {length / total_weight:.4f}",
                flush=True,
            )
    vocabulary = SpatialVocabulary(
        tuple(int(value) for value in bins), tuple(trainer.merges)
    )
    report = {
        "records": int(trainer.count),
        "weight": total_weight,
        "atoms": trainer.atom_count,
        "merges": len(trainer.merges),
        "vocab_size": vocabulary.vocab_size,
        "tokens_per_action_initial": initial / total_weight,
        "tokens_per_action": length / total_weight,
        "compression": initial / length,
        "sparse_updates": trainer.statistics["sparse_updates"],
        "sparse_heap_rebuilds": trainer.statistics["sparse_heap_rebuilds"],
        "sparse_pair_types_peak": trainer.statistics["sparse_pair_types_peak"],
    }
    return vocabulary, report


def train_spatial(
    codes: np.ndarray,
    weights: np.ndarray,
    bins: Sequence[int],
    *,
    vocab_size: int,
    backend: str = "auto",
    progress: bool = False,
    **kwargs: object,
) -> tuple[SpatialVocabulary, dict[str, object]]:
    """Fit spatial Set-BPE, preferring the native kernel when it is built.

    ``codes`` is one row per distinct quantized action and ``weights`` its
    multiplicity (see :func:`aggregate_records`); collapsing duplicates first is
    lossless because identical actions follow identical merges forever.

    ``vocab_size`` is the *total* budget including the ``sum(bins)`` primitive
    atoms.  Training stops early once no positive-frequency compatible pair
    remains.  Both backends return the identical greedy merge sequence -- the
    test suite proves it -- so ``backend`` only trades speed for portability.
    """
    if resolve_backend(backend) == "python":
        # ``threads`` controls only the native OpenMP kernel. The profile fit API carries
        # it for either backend, but the Python reference has no such argument.
        kwargs.pop("threads", None)
        return _train_spatial_python(
            codes, weights, bins, vocab_size=vocab_size, progress=progress, **kwargs
        )
    values = np.ascontiguousarray(codes, dtype=np.int32)
    masses = np.ascontiguousarray(weights, dtype=np.int64)
    merges, report = require_native().train_spatial(
        values,
        masses,
        [int(count) for count in bins],
        int(vocab_size),
        int(kwargs.get("threads", 0)),
    )
    vocabulary = SpatialVocabulary(
        tuple(int(count) for count in bins),
        tuple((int(left), int(right)) for left, right in merges),
    )
    return vocabulary, dict(report)