"""Dense single-stage spatial/temporal Set-BPE using AC2Sparse's merge trainer."""

from __future__ import annotations

import json
from collections.abc import Sequence
from dataclasses import dataclass
from numbers import Integral
from pathlib import Path

import numpy as np

from ._validation import active_dimensions
from .bpe import BPE_NAME, _dense_cells, _read_payload, _rules
from .cache import BoundedCache
from .grammar import _PrefixGrammar
from .sparse.training import fit_sparse_vocabulary
from .sparse.vocabulary import SparseVocabulary


@dataclass(frozen=True)
class DenseSetLayout:
    'Keep every active physical cell and append an end-of-frame atom.'

    bins: tuple[int, ...]
    dense = True

    @property
    def dims(self):
        return len(self.bins)

    @property
    def e_dim(self):
        return self.dims

    @property
    def zero(self):
        return (-2,) * self.dims

    @property
    def native_zero(self):
        return [-2] * self.dims

    @property
    def sparse_bins(self):
        return self.bins + (1,)

    def sparsify(self, cells):
        cells = np.asarray(cells, dtype=np.int64)
        if cells.ndim != 2 or cells.shape[1] != self.dims or not len(cells):
            raise ValueError("cells must be a nonempty dense trajectory")
        if np.any(cells < -1) or np.any(cells >= np.asarray(self.bins)):
            raise ValueError("cell outside the primitive range")
        if np.any((cells >= 0) != (cells[0] >= 0)) or not np.any(cells[0] >= 0):
            raise ValueError("active slots must remain fixed and nonempty")
        out = np.full((len(cells), self.dims + 1), -1, dtype=np.int64)
        out[:, : self.dims] = cells
        out[:-1, self.dims] = 0
        return out

    def densify(self, frames, active):
        cells = np.asarray(frames, dtype=np.int64)[:, : self.dims].copy()
        flags = np.asarray(active, dtype=bool)
        if np.any(cells[:, flags] < 0) or np.any(cells[:, ~flags] >= 0):
            raise ValueError("dense stream is missing an active cell")
        return cells


class UnifiedSetBPESerializer:
    'Single-stage Set-BPE over complete physical cells and frame markers.'

    scheme = "unified_set_bpe"

    def __init__(self, *, bins: Sequence[int], lookahead: int, merges):
        self.bins = tuple(int(v) for v in bins)
        self.lookahead = int(lookahead)
        self.merges = _rules(merges, 3, "merges")
        self.layout = DenseSetLayout(self.bins)
        self.vocabulary = SparseVocabulary(self.layout, (), self.merges, self.lookahead)
        self.vocab_size = self.vocabulary.vocab_size
        physical_mask = (1 << len(self.bins)) - 1
        self.token_cell_counts = np.asarray(
            [
                sum((int(mask) & physical_mask).bit_count() for mask in masks)
                for masks in self.vocabulary.tables.masks
            ],
            dtype=np.int64,
        )

    def configure_encoding(self, *, spatial_search=False, temporal_search=False, **kwargs):
        if spatial_search or temporal_search:
            raise ValueError("packing search is not available for unified Set-BPE")
        return self

    def encode_batch(self, cells, *, active_dims=None):
        dense, _ = _dense_cells(np.asarray(cells), self.bins, active_dims, ndim=3)
        return self.vocabulary.encode_batch(dense)

    def encode(self, cells, *, active_dims=None):
        dense, _ = _dense_cells(np.asarray(cells), self.bins, active_dims)
        return self.vocabulary.encode_batch(dense[None])[0]

    def decode_batch(self, streams, *, active_dims=None):
        dims = active_dimensions(active_dims, len(self.bins))
        rows, horizons = [], []
        for row in streams:
            checked = []
            for token in row:
                if (
                    isinstance(token, bool)
                    or not isinstance(token, Integral)
                    or not 0 <= int(token) < self.vocab_size
                ):
                    raise ValueError("token outside vocabulary")
                checked.append(int(token))
            count = sum(int(self.token_cell_counts[token]) for token in checked)
            horizon, remainder = divmod(count, len(dims))
            if not checked or not horizon or remainder:
                raise ValueError("incomplete unified Set-BPE stream")
            rows.append(checked)
            horizons.append(horizon)
        if not rows or len(set(horizons)) != 1:
            raise ValueError("streams must have the same nonempty horizon")
        flags = [d in dims for d in range(len(self.bins))]
        return self.vocabulary.decode_batch(rows, horizons[0], flags)

    def decode(self, tokens, *, active_dims=None):
        return self.decode_batch([tokens], active_dims=active_dims)[0]

    def to_dict(self):
        return {"merges": [list(rule) for rule in self.merges]}

    @classmethod
    def from_pretrained(cls, directory: str | Path, *, bins, lookahead):
        payload = _read_payload(directory)
        if set(payload) != {"merges"}:
            raise ValueError("unified bpe.json requires exactly merges")
        return cls(bins=bins, lookahead=lookahead, merges=payload["merges"])

    def save_pretrained(self, directory: str | Path):
        path = Path(directory) / BPE_NAME
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2) + "\n", encoding="utf-8")
        return path


class UnifiedSetBPEGrammar(_PrefixGrammar):
    'Dense frame-coverage grammar for unified Set-BPE.'

    def __init__(self, serializer, *, horizon, active_dims=None):
        if isinstance(horizon, bool) or not isinstance(horizon, int) or horizon < 1:
            raise ValueError("horizon must be positive")
        self.horizon, self.vocab_size = horizon, serializer.vocab_size
        dims = active_dimensions(active_dims, len(serializer.bins))
        self._full = sum(1 << d for d in dims)
        self._e_bit = 1 << len(serializer.bins)
        self._e_atom = sum(serializer.bins)
        self._patterns = serializer.vocabulary.tables.masks
        self._lookahead = serializer.lookahead
        self._state_cache, self._mask_cache = BoundedCache(256), BoundedCache(64)
        self._empty_state = self._initial_state()
        self._native = None

    def _compiled(self):
        return None

    def _initial_state(self):
        return (0,) * self.horizon

    def _covered_steps(self, state):
        for index, mask in enumerate(state):
            needed = self._full | (self._e_bit if index < self.horizon - 1 else 0)
            if mask != needed:
                return index
        return self.horizon

    def _state_is_complete(self, state):
        return self._covered_steps(state) == self.horizon

    def _advance(self, state, token):
        if not 0 <= token < self.vocab_size:
            return None
        front = self._covered_steps(state)
        if front == self.horizon:
            return None
        masks = self._patterns[token]
        if front + len(masks) > self.horizon:
            return None
        updated = list(state)
        allowed = self._full | self._e_bit
        for offset, mask in enumerate(masks):
            mask = int(mask)
            if not mask:
                continue
            index = front + offset
            if mask & ~allowed or mask & updated[index] or updated[index] & self._e_bit:
                return None
            if index == self.horizon - 1 and mask & self._e_bit:
                return None
            updated[index] |= mask
            if mask & self._e_bit and updated[index] & self._full != self._full:
                return None
        return tuple(updated)

    def _python_mask(self, state):
        return np.asarray(
            [self._advance(state, token) is not None for token in range(self.vocab_size)]
        )

    def completion_token(self, state, zero_cells, bins):
        front = self._covered_steps(state)
        missing = self._full & ~state[front]
        if not missing:
            return self._e_atom
        dim = (missing & -missing).bit_length() - 1
        return sum(bins[:dim]) + int(zero_cells[dim])


def fit_unified_set_bpe(config, cells, *, threads, backend="native", progress=True):
    """Fit AC2Sparse's one-stage Set-BPE on complete dense-cell chunks."""
    from .training import build_primitives

    primitives = build_primitives(config)
    layout = DenseSetLayout(tuple(primitives.bins))
    vocab, report = fit_sparse_vocabulary(
        layout,
        cells,
        spatial_vocab_size=sum(layout.sparse_bins),
        vocab_size=config.train.bpe.vocab_size,
        lookahead=config.set_bpe.lookahead,
        backend=backend,
        threads=threads,
        progress=progress,
    )
    serializer = UnifiedSetBPESerializer(
        bins=primitives.bins, lookahead=config.set_bpe.lookahead, merges=vocab.temporal_merges
    )
    return serializer, report