runtime/processing_actioncodec2.py

"""Hugging Face processor over the shared ActionCodec2 frontend implementation."""

from __future__ import annotations

import json
from collections.abc import Mapping
from pathlib import Path
from typing import Any

import numpy as np
import torch
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin

from .routing.defaults import CodecBuilder, ProfileBundleBuilder, builtin_action_spaces
from .integration.hf import (
    PROCESSOR_CONFIG_NAME,
    ROUTER_ALLOW_PATTERNS,
    ActionCodec2FrontendHFConfig,
    resolve_pretrained_directory,
    save_huggingface_metadata,
    save_model_card,
    upload_pretrained_directory,
)
from .routing.presets import PresetConfig
from .routing.routing import ActionCodec2Router
from .frontend.tokenizer import ActionCodec2Tokenizer


def _saved_metadata(directory: Path) -> dict[str, Any]:
    metadata = directory / PROCESSOR_CONFIG_NAME
    return (
        json.loads(metadata.read_text(encoding="utf-8")) if metadata.is_file() else {}
    )


def _fitted_status(profiles) -> str:
    status = ", ".join(sorted(profiles)) or "none; call codec.fit(episodes)"
    return f"Fitted profiles: {status}\n"


class ActionCodec2(ActionCodec2Tokenizer, ProcessorMixin):
    """Action-space-bound processor with lossless token serialization.

    ``encode`` returns ragged token-ID rows; ``decode`` reconstructs a CPU
    ``torch.float32`` tensor in the original ``(B,T,D)`` column layout.
    Calling the processor wraps those same rows in a ``BatchFeature`` under
    ``input_ids``. Physical quantization remains lossy as configured.

    Released first- and second-order bundles both register 13 layouts: single
    and dual EEF arms, single and dual arms with 6 or 7 joint coordinates, each
    with absolute/delta variants and binary grippers, plus a six-continuous-
    coordinate absolute layout without a separate binary gripper column.
    The registered names and column meanings come from the loaded artifact.
    ``print(codec)`` lists all of them; ``print(cls.describe_pretrained(path))``
    inspects them before binding. ``print_action_spaces(action_dim=D)`` only
    filters candidates, because shape alone cannot identify physical meaning.

    Inputs accept ``(T,D)`` or ``(B,T,D)`` with any positive T; unequal-length
    episodes must be encoded separately. Absolute layouts require the same
    ``current_state(B,D)`` for encode/decode. Second-order physical profiles
    additionally require ``previous_first_order: {component: (B,D_active)}``
    at the action space's configured codec rate, independently of the
    absolute/delta choice.
    Boundary arrays retain B=1 for a single episode. Physical order does not
    change the caller's action layout into an acceleration tensor.

    Args:
        backend: Optional fitted router. Omit to construct a fresh codec for fit().
        action_space: Registered name defining the input tensor's semantics.
        primitive_order: Physical primitive order, 1 or 2, for fresh construction.
        bpe_scheme: Optional 'set_bpe' or 'bpe' override for fresh construction.
        profile_configs: Optional joint/EEF physical configs, mappings or YAML paths.
            Defaults ship inside the package; no repository checkout is required.
    """

    model_type = ActionCodec2FrontendHFConfig.model_type
    model_input_names = ["input_ids"]
    # Numerical artifacts are owned by the router, not modality sub-processors.
    attributes = []

    def __init__(
        self,
        backend: ActionCodec2Router | None = None,
        *,
        action_space: str,
        primitive_order: int = 1,
        bpe_scheme: str | None = None,
        profile_configs: Mapping[str, object] | None = None,
    ) -> None:
        fresh = backend is None
        if fresh:
            backend = CodecBuilder(
                primitive_order=primitive_order,
                bpe_scheme=bpe_scheme,
                profile_configs=profile_configs,
            ).build()
        elif (
            primitive_order != 1
            or bpe_scheme is not None
            or profile_configs is not None
        ):
            raise ValueError(
                "profile construction options cannot accompany an existing backend"
            )
        super().__init__(backend, action_space=action_space)
        ProcessorMixin.__init__(self)
        self._fitted_profiles = frozenset() if fresh else frozenset(backend.profiles)
        self.training_report = None
        self._source_directory = Path(__file__).parent

    def __dir__(self):
        inapplicable = {
            "action_tokenizer",
            "attributes",
            "audio_tokenizer",
            "chat_template",
            "check_argument_for_proper_class",
            "encode_kwargs",
            "feature_extractor_class",
            "from_args_and_dict",
            "get_possibly_dynamic_module",
            "get_processor_dict",
            "model_input_names",
            "model_type",
            "optional_attributes",
            "optional_call_args",
            "apply_chat_template",
            "post_process_image_text_to_text",
            "pad_action_id",
            "bos_blk_id",
            "parts_meta",
            "BAR_SENTINEL_TOKENS",
            "push_to_hub",
            "register_for_auto_class",
            "tokenizer_class",
            "to_dict",
            "to_json_file",
            "to_json_string",
            "valid_processor_kwargs",
            "validate_init_kwargs",
        }
        return [name for name in super().__dir__() if name not in inapplicable]

    def __repr__(self) -> str:
        return (
            f"ActionCodec2(action_space={self.action_space.name!r}, "
            f"D={self.action_dim}, vocab={self.action_tokenizer.vocab_size})"
        )

    @staticmethod
    def builtin_action_spaces():
        """Return the immutable named layout definitions shipped in the package."""
        return builtin_action_spaces()

    @classmethod
    def from_profiles(cls, joint, eef, *, action_space: str, token_budgets=None):
        """Bind an artifact assembled from existing joint/EEF profile directories.

        Args:
            joint: Fitted joint profile directory.
            eef: Fitted EEF profile directory.
            action_space: Explicit built-in tensor layout to bind.
            token_budgets: Optional joint/EEF overrides; defaults to recorded
                training budgets, or fitted vocabulary sizes if unavailable.
        """
        router = ProfileBundleBuilder(joint, eef, token_budgets=token_budgets).build()
        return cls(router, action_space=action_space)

    def logits_processor(
        self,
        horizon: int,
        *,
        fps: float,
        prompt_length: int,
        eos_token_id: int,
        token_offset: int = 0,
        pad_token_id=None,
    ):
        """Create a Transformers LogitsProcessor for a fixed action horizon.

        ``prompt_length`` excludes leading model prompt IDs from the grammar;
        ``token_offset`` explicitly maps codec IDs to a contiguous model range.
        EOS/padding IDs must be outside that range. No model vocabulary is changed.
        """
        from .tokenization.generation import ActionCodec2LogitsProcessor

        return ActionCodec2LogitsProcessor(
            self,
            horizon,
            fps=fps,
            prompt_length=prompt_length,
            eos_token_id=eos_token_id,
            token_offset=token_offset,
            pad_token_id=pad_token_id,
        )

    def batch_decode(
        self,
        token_rows,
        *,
        fps=None,
        current_state=None,
        previous_first_order=None,
        errors="raise",
    ):
        """Decode independent token rows, including different valid horizons.

        Args:
            token_rows: Iterable of token-ID rows. Each must cover a complete
                trajectory; intended generation horizon must be checked separately
                with grammar.is_complete before decode.
            fps: Output sampling rate, shared by all rows.
            current_state: Optional (B,D) preceding states in original layout.
            previous_first_order: Optional {component: (B,D_active)} boundaries.
            errors: 'raise' stops at the first invalid row; 'return' returns the
                ValueError at that position so other valid rows remain usable.

        Returns:
            List of CPU float32 tensors shaped (1,T_i,D), or ValueError entries
            when errors='return'. No truncated row is padded or executed partially.
        """
        self._require_fitted()
        if errors not in ("raise", "return"):
            raise ValueError("errors must be 'raise' or 'return'")
        rows = list(token_rows)
        current = (
            None
            if current_state is None
            else np.asarray(
                current_state.detach().cpu()
                if isinstance(current_state, torch.Tensor)
                else current_state
            )
        )
        if current is not None and current.shape != (len(rows), self.action_dim):
            raise ValueError(
                f"current_state must have shape {(len(rows), self.action_dim)}"
            )
        previous = self._numpy_boundary(previous_first_order)
        if previous is not None and any(
            np.ndim(v) != 2 or len(v) != len(rows) for v in previous.values()
        ):
            raise ValueError(
                "previous_first_order arrays must have batch dimension equal to token rows"
            )
        result = []
        for index, row in enumerate(rows):
            try:
                decoded = self.decode(
                    [row],
                    fps=fps,
                    current_state=None
                    if current is None
                    else current[index : index + 1],
                    previous_first_order=None
                    if previous is None
                    else {k: v[index : index + 1] for k, v in previous.items()},
                )
            except ValueError as error:
                error = ValueError(f"token row {index}: {error}")
                if errors == "raise":
                    raise error from None
                decoded = error
            result.append(decoded)
        return result

    @property
    def fitted_profiles(self) -> frozenset[str]:
        """Names of profiles fitted on data; persisted by save_pretrained()."""
        return self._fitted_profiles

    @property
    def is_fitted(self) -> bool:
        """Whether every profile used by the selected action space is fitted."""
        return all(
            c.space in self.fitted_profiles for c in self.action_space.components
        )

    def _require_fitted(self):
        missing = {c.space for c in self.action_space.components} - self.fitted_profiles
        if missing:
            raise RuntimeError(
                f"Unfitted profiles {sorted(missing)}; call codec.fit(episodes) first"
            )

    def _describe_action_spaces(self, action_dim: int | None = None) -> str:
        return _fitted_status(self.fitted_profiles) + super()._describe_action_spaces(
            action_dim
        )

    def for_action_space(self, name: str) -> ActionCodec2:
        """Share current weights with a new binding; subsequent fit replaces only that instance."""
        codec = super().for_action_space(name)
        codec._fitted_profiles = self.fitted_profiles
        codec.training_report = self.training_report
        codec._source_directory = self._source_directory
        return codec

    def fit(
        self,
        episodes,
        *,
        fps: float | None = None,
        current_state=None,
        previous_first_order=None,
        bootstrap_context: bool = False,
        backend: str = "auto",
        threads: int = 0,
        batch_size: int = 256,
        progress: bool = False,
        max_cell_bytes: int | None = None,
        work_dir: str | Path | None = None,
    ) -> ActionCodec2:
        """Fit physical vocabularies from actions in declared source layouts.

        Args:
            episodes: Array/tensor (T,D), (B,T,D), or iterable of variable-length
                episodes. Records may contain ``actions``, ``action_space`` (a
                registered name), ``fps``, ``current_state(B,D)``, and
                ``previous_first_order: {component: (B,D_active)}``. A record's
                metadata must not duplicate the corresponding keyword argument.
            fps: Source frequency override shared by all input episodes.
            current_state: Preceding source state (B,D), required for absolute
                inputs. Use per-record values when episodes have different states.
            previous_first_order: Preceding canonical codec-rate increments by component,
                with shape (B,D_active), required for second-order profiles.
            bootstrap_context: Use two measured canonical frames as second-order
                context, training on the remaining suffix. Mutually exclusive with
                previous_first_order. Inference still requires explicit boundaries.
            backend: Exact BPE implementation: 'auto', 'native', or 'python'.
            threads: Native worker count; 0 uses the backend default.
            batch_size: Maximum canonical trajectories per quantization batch.
            progress: Display BPE training progress.
            max_cell_bytes: Hard bound on retained quantized corpus bytes per profile.
            work_dir: Parent scratch directory for sequential source staging files,
                automatically removed on success or failure. Defaults to system temp.

        Returns:
            This instance, with per-profile results in ``training_report``. Only
            profiles present in the data are fitted; other weights are preserved.
            Vocabulary replacement is atomic on success. Refit changes token IDs:
            regenerate previously encoded datasets before using the new vocabulary.
        """
        # Training is intentionally lazy: a loaded Hub artifact only needs the
        # inference runtime. The full fitting pipeline remains available from
        # the source package and is imported when the user calls fit().
        from .training.fitting import ActionCodec2Fitter

        router, report = ActionCodec2Fitter(
            self, bootstrap_context=bootstrap_context
        ).fit(
            episodes,
            work_dir=work_dir,
            metadata=dict(
                fps=fps,
                current_state=current_state,
                previous_first_order=previous_first_order,
            ),
            backend=backend,
            threads=threads,
            batch_size=batch_size,
            progress=progress,
            max_cell_bytes=max_cell_bytes,
        )
        fitted = self.fitted_profiles | report["profiles"].keys()
        replacement = type(self)(router, action_space=self.action_space.name)
        replacement._fitted_profiles = frozenset(fitted)
        replacement.training_report = report
        replacement._source_directory = self._source_directory
        replacement._warned_default_fps = self._warned_default_fps
        self.__dict__.update(replacement.__dict__)
        return self

    def _encode_action_indices(self, action, encode_kwargs=None):
        self._require_fitted()
        return super()._encode_action_indices(action, encode_kwargs)

    def _decode_action_indices(self, action_indices, **kwargs):
        self._require_fitted()
        return super()._decode_action_indices(action_indices, **kwargs)

    def grammar(self, horizon: int, *, fps: float | None = None):
        self._require_fitted()
        return super().grammar(horizon, fps=fps)

    def final_first_order(
        self, action_token_ids, *, previous_first_order, executed_steps=None
    ):
        self._require_fitted()
        return super().final_first_order(
            action_token_ids,
            previous_first_order=previous_first_order,
            executed_steps=executed_steps,
        )

    @classmethod
    def describe_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *,
        action_dim: int | None = None,
        **kwargs: Any,
    ) -> str:
        """Inspect a local/Hub artifact before choosing an action space.

        Args:
            pretrained_model_name_or_path: Saved processor/router directory or
                Hugging Face Hub repository ID. Loads its fitted router for inspection.
            action_dim: Optional D from ``actions.shape[-1]``. Only filters the
                report; tensor shape cannot establish physical semantics.
            **kwargs: Hub options such as ``revision``, ``subfolder``,
                ``cache_dir``, ``token``, and ``local_files_only``.

        Returns:
            Human-readable layout guide, suitable for ``print(...)``. No binding
            is required or inferred, and no action data is encoded.
        """
        directory = resolve_pretrained_directory(
            pretrained_model_name_or_path,
            allow_patterns=ROUTER_ALLOW_PATTERNS,
            **kwargs,
        )
        router = ActionCodec2Router.from_pretrained(directory)
        metadata = _saved_metadata(directory)
        return _fitted_status(
            metadata.get("fitted_profiles", router.profiles)
        ) + router.describe_action_spaces(
            selected=metadata.get("action_space"), action_dim=action_dim
        )

    @classmethod
    def from_pretrained(
        cls,
        pretrained_model_name_or_path: str | Path,
        *,
        action_space: str | Path | PresetConfig | Mapping[str, object] | None = None,
        **kwargs: Any,
    ) -> ActionCodec2:
        """Load a local/Hub artifact and restore or override its action space.

        Args:
            pretrained_model_name_or_path: Saved processor or router directory,
                or a Hugging Face Hub repository ID.
            action_space: Registered name, preset YAML path, config, or mapping.
                Required for an unbound router; otherwise defaults to the saved
                binding. Custom definitions are saved in the router manifest.
            **kwargs: Hub options including ``revision``, ``subfolder``,
                ``cache_dir``, ``token``, and ``local_files_only``.
        """
        directory = resolve_pretrained_directory(
            pretrained_model_name_or_path,
            allow_patterns=ROUTER_ALLOW_PATTERNS,
            **kwargs,
        )
        if action_space is None:
            action_space = _saved_metadata(directory).get("action_space")
        if action_space is None:
            raise ValueError(
                "action_space is required when loading an unbound router artifact. "
                "First inspect its layouts with print(ActionCodec2.describe_pretrained(path)), "
                "then pass action_space='NAME'; shape alone cannot identify physical meaning."
            )
        codec = super().from_pretrained(directory, action_space=action_space)
        metadata = _saved_metadata(directory)
        fitted = metadata.get("fitted_profiles", list(codec.action_tokenizer.profiles))
        if not isinstance(fitted, list) or any(
            not isinstance(name, str) or name not in codec.action_tokenizer.profiles
            for name in fitted
        ):
            raise ValueError("invalid fitted_profiles in processor_config.json")
        codec._fitted_profiles = frozenset(fitted)
        codec._source_directory = directory
        report_path = directory / "fit_report.json"
        if report_path.is_file():
            codec.training_report = json.loads(report_path.read_text(encoding="utf-8"))
        return codec

    def to_dict(self, legacy_serialization: bool = True) -> dict[str, object]:
        """Return lightweight discovery metadata.

        ``legacy_serialization`` is accepted for ProcessorMixin signature
        compatibility. ActionCodec2 has one serialization format for both values.
        """
        del legacy_serialization
        return {
            "processor_class": type(self).__name__,
            "model_type": self.model_type,
            "action_space": self.action_space.name,
            "fitted_profiles": sorted(self.fitted_profiles),
        }

    def save_pretrained(
        self, save_directory: str | Path, push_to_hub: bool = False, **kwargs: Any
    ) -> list[str]:
        """Save the complete router and the selected action-space binding."""
        if not push_to_hub and kwargs:
            raise TypeError(
                f"unused save_pretrained keyword(s): {', '.join(sorted(kwargs))}"
            )
        from .integration.exporting import ArtifactExporter

        directory = Path(save_directory)
        save_model_card(directory, str(self))
        files = self.action_tokenizer.save_pretrained(directory, include_code=False)
        files.extend(
            ArtifactExporter(
                directory, source_directory=self._source_directory
            ).export()
        )
        save_huggingface_metadata(
            directory,
            **self.to_dict(),
            auto_map={"AutoProcessor": "processing_actioncodec2.ActionCodec2"},
        )
        report_path = directory / "fit_report.json"
        if self.training_report is not None:
            report_path.write_text(
                json.dumps(self.training_report, indent=2) + "\n", encoding="utf-8"
            )
            files.append(str(report_path))
        elif report_path.exists():
            report_path.unlink()
        if push_to_hub:
            upload_pretrained_directory(directory, **kwargs)
        return files

    def __call__(
        self,
        action: torch.Tensor | np.ndarray,
        *,
        fps: float | None = None,
        current_state: torch.Tensor | np.ndarray | None = None,
        previous_first_order: Mapping[str, torch.Tensor | np.ndarray] | None = None,
        encode_kwargs: dict | None = None,
        return_tensors: str | None = None,
    ) -> BatchFeature:
        """Prepare ``input_ids`` from ``(B,T,D)`` or a single ``(T,D)`` action.

        ``current_state(B,D)`` is required for absolute actions. Second-order
        artifacts require ``previous_first_order`` mapping component names to
        ``(B,D_active)`` canonical increments, as in :meth:`encode`.
        ``return_tensors`` may be ``"pt"`` or ``"np"`` for equal-length rows;
        omitted, rows remain ragged lists without adding padding tokens.
        """
        rows = self.encode(
            action,
            fps=fps,
            current_state=current_state,
            previous_first_order=previous_first_order,
            encode_kwargs=encode_kwargs,
        )
        return BatchFeature({"input_ids": rows}, tensor_type=return_tensors)