"""One ActionCodec2Sparse physical profile: unchanged physical stage plus sparse token layer."""

from __future__ import annotations

import hashlib
import json
from collections.abc import Sequence
from numbers import Integral
from pathlib import Path

import numpy as np

from .._validation import active_dimensions
from ..cache import BoundedCache
from ..configuration import CONFIG_NAME, ActionCodec2Config, ConfigError
from ..hf import save_huggingface_metadata
from ..physical_codec import PhysicalCodec
from ..primitives import PRIMITIVES_NAME, PrimitiveStore
from .grammar import SparseGrammar
from .layout import SparseLayout
from .vocabulary import SparseVocabulary

PROFILE_PROCESSOR_CLASS = "ActionCodec2SparseProcessor"


class ActionCodec2SparseProcessor:
    """Fitted sparse profile over canonical first- or second-order profile actions."""

    def __init__(
        self,
        config: ActionCodec2Config,
        primitives: PrimitiveStore,
        vocabulary: SparseVocabulary,
    ) -> None:
        if config.set_bpe.scheme != "sparse_set_bpe":
            raise ConfigError("ActionCodec2Sparse profiles require set_bpe.scheme=sparse_set_bpe")
        if vocabulary.layout != SparseLayout.from_config(config, primitives):
            raise ValueError("sparse layout does not match the physical primitives")
        if vocabulary.lookahead != config.set_bpe.lookahead:
            raise ValueError("sparse vocabulary lookahead does not match config.yaml")
        self.config, self.primitives, self.vocabulary = config, primitives, vocabulary
        self.physical = PhysicalCodec(config, primitives)
        self.requires_previous_first_order = self.physical.requires_previous_first_order
        self.action_dim = config.physical.input_dim
        self.bins = primitives.bins
        self._grammar_cache = BoundedCache(8)

    @property
    def vocab_size(self) -> int:
        return self.vocabulary.vocab_size

    @property
    def layout(self) -> SparseLayout:
        return self.vocabulary.layout

    @classmethod
    def from_pretrained(cls, directory: str | Path) -> ActionCodec2SparseProcessor:
        """Load ``config.yaml``, ``primitives.npz`` and ``sparse_bpe.json``."""
        root = Path(directory)
        config = ActionCodec2Config.from_pretrained(root)
        primitives = PrimitiveStore.from_pretrained(root, input_dim=config.physical.input_dim)
        return cls(config, primitives, SparseVocabulary.from_pretrained(root))

    def save_pretrained(self, directory: str | Path) -> list[str]:
        root = Path(directory)
        self.config.save_pretrained(root)
        self.primitives.save_pretrained(root)
        files = [root / CONFIG_NAME, root / PRIMITIVES_NAME, self.vocabulary.save_pretrained(root)]
        files.extend(
            save_huggingface_metadata(
                root, processor_class=PROFILE_PROCESSOR_CLASS, vocab_size=self.vocab_size
            )
        )
        return [str(path) for path in files]

    @property
    def artifact_fingerprint(self) -> str:
        """Content hash over config, primitive axes and the ordered sparse rules."""
        digest = hashlib.sha256()
        digest.update(
            json.dumps(self.config.to_dict(), sort_keys=True, separators=(",", ":")).encode()
        )
        for axis in self.primitives.values:
            values = np.ascontiguousarray(axis)
            digest.update(str(values.dtype).encode())
            digest.update(np.asarray(values.shape, dtype=np.int64).tobytes())
            digest.update(values.tobytes())
        digest.update(
            json.dumps(self.vocabulary.to_dict(), sort_keys=True, separators=(",", ":")).encode()
        )
        return digest.hexdigest()

    def _active(self, dims: tuple[int, ...]) -> np.ndarray:
        flags = np.zeros(self.action_dim, dtype=bool)
        flags[list(dims)] = True
        return flags

    def _dense_previous_first_order(
        self, previous_first_order: np.ndarray | None, dims: tuple[int, ...], batch: int
    ) -> np.ndarray | None:
        active = set(dims)
        for component in self.config.physical.components.values():
            present = active.intersection(component.dims)
            if (
                component.primitive_order == 2
                and component.space == "so3"
                and present
                and present != set(component.dims)
            ):
                raise ValueError(
                    "second-order SO(3) requires all three active component dimensions"
                )
        if previous_first_order is None:
            if any(
                c.primitive_order == 2 and active.intersection(c.dims)
                for c in self.config.physical.components.values()
            ):
                raise ValueError(
                    f"second-order tokenization requires previous_first_order with shape {(batch, len(dims))} in canonical input units"
                )
            return (
                np.zeros((batch, self.action_dim), dtype=np.float32)
                if self.requires_previous_first_order
                else None
            )
        values = np.asarray(previous_first_order, dtype=np.float32)
        if values.shape != (batch, len(dims)) or not np.all(np.isfinite(values)):
            raise ValueError(f"previous_first_order must be finite with shape {(batch, len(dims))}")
        dense = np.zeros((batch, self.action_dim), dtype=np.float32)
        dense[:, dims] = values
        return dense

    def quantize_batch(
        self, actions: np.ndarray, *, active_dims=None, previous_first_order=None
    ) -> np.ndarray:
        """Compact ``(B,T,D_active)`` actions to dense cells ``(B,T,D)``, ``-1`` inactive."""
        values = np.asarray(actions, dtype=np.float32)
        if values.ndim != 3 or len(values) < 1 or values.shape[1] < 1:
            raise ValueError("actions must have non-empty shape [B, T, D_active]")
        dims = active_dimensions(active_dims, self.action_dim)
        if len(dims) != values.shape[2]:
            raise ValueError("active_dims must have one entry per action coordinate")
        dense = np.zeros((len(values), values.shape[1], self.action_dim), dtype=np.float32)
        dense[:, :, dims] = values
        previous = self._dense_previous_first_order(previous_first_order, dims, len(values))
        cells = np.asarray(
            self.physical.encode_batch(dense, previous_first_order=previous), dtype=np.int64
        )
        cells[:, :, ~self._active(dims)] = -1
        return cells

    def encode_batch(
        self,
        actions: np.ndarray,
        *,
        active_dims: Sequence[int] | None = None,
        previous_first_order: np.ndarray | None = None,
        threads: int = 0,
    ) -> list[list[int]]:
        """Encode independent ``(B,T,D_active)`` chunks into ragged local token rows."""
        cells = self.quantize_batch(
            actions, active_dims=active_dims, previous_first_order=previous_first_order
        )
        return self.vocabulary.encode_batch(cells, threads=threads)

    def cells_from_tokens(
        self, rows: Sequence[Sequence[int]], *, horizon: int, active_dims=None
    ) -> np.ndarray:
        """Losslessly recover dense cells ``(B,horizon,D)`` with ``-1`` for inactive slots."""
        dims = active_dimensions(active_dims, self.action_dim)
        return self.vocabulary.decode_batch(rows, horizon, self._active(dims))

    def _filled(self, cells: np.ndarray) -> np.ndarray:
        zero = [0 if z is None else z for z in self.layout.zero]
        return np.where(cells >= 0, cells, np.asarray(zero, dtype=np.int64))

    def decode_batch(
        self,
        rows: Sequence[Sequence[int]],
        *,
        horizon: int,
        active_dims=None,
        previous_first_order: np.ndarray | None = None,
    ) -> np.ndarray:
        """Decode rows to ``(B,horizon,D_active)`` reconstructed profile actions."""
        dims = active_dimensions(active_dims, self.action_dim)
        cells = self.cells_from_tokens(rows, horizon=horizon, active_dims=dims)
        previous = self._dense_previous_first_order(previous_first_order, dims, len(cells))
        decoded = self.physical.decode_batch(self._filled(cells), previous_first_order=previous)
        return np.take(decoded, dims, axis=2)

    def final_first_order(
        self,
        rows: Sequence[Sequence[int]],
        *,
        horizon: int,
        previous_first_order: np.ndarray,
        active_dims=None,
        executed_steps: int | None = None,
    ) -> np.ndarray:
        """Next compact ``(B,D_active)`` boundary after ``executed_steps`` codec frames."""
        dims = active_dimensions(active_dims, self.action_dim)
        cells = self.cells_from_tokens(rows, horizon=horizon, active_dims=dims)
        if executed_steps is not None:
            if (
                isinstance(executed_steps, bool)
                or not isinstance(executed_steps, Integral)
                or not 1 <= executed_steps <= cells.shape[1]
            ):
                raise ValueError("executed_steps must be within the chunk horizon")
            cells = cells[:, :executed_steps]
        previous = self._dense_previous_first_order(previous_first_order, dims, len(cells))
        result = self.physical.final_first_order(self._filled(cells), previous_first_order=previous)
        return np.take(result, dims, axis=1)

    def grammar(self, horizon: int, *, active_dims=None) -> SparseGrammar:
        """Prefix grammar of one ``horizon``-frame chunk over the active slots."""
        dims = active_dimensions(active_dims, self.action_dim)
        key = (int(horizon), dims)
        grammar = self._grammar_cache.get(key)
        if grammar is None:
            grammar = SparseGrammar(self.vocabulary, horizon, self._active(dims))
            self._grammar_cache[key] = grammar
        return grammar