"""Write the portable runtime and user guide for a saved ActionCodec2 artifact."""

import shutil
from pathlib import Path

# Files included in a portable runtime artifact.
MODULE_FILES: list[str]
SUPPORT_FILES: dict[str, str]
USER_GUIDE: str


class ArtifactExporter:
    """Copy portable code alongside numerical artifacts.

    Args:
        directory: Destination directory.
        source_directory: Existing artifact or package directory used as a
            fallback for runtime files and native resources.
    """

    def __init__(
        self, directory: str | Path, *, source_directory: str | Path | None = None
    ) -> None:
        self.output = Path(directory)
        self.package_root = Path(__file__).resolve().parent.parent
        self.source = Path(source_directory) if source_directory else self.package_root

    def export(self) -> list[str]:
        self.output.mkdir(parents=True, exist_ok=True)
        files = []
        runtime = self.output / "runtime"
        runtime.mkdir(parents=True, exist_ok=True)
        for name in MODULE_FILES:
            # Copy the current implementation so an artifact can be moved alone.
            source = self.package_root / name
            if not source.is_file():
                source = self.source / name
            target = runtime / name
            target.parent.mkdir(parents=True, exist_ok=True)
            if source.resolve() != target.resolve():
                shutil.copyfile(source, target)
            files.append(str(target))
        (runtime / "__init__.py").write_text(
            '"""Private inference runtime bundled with this ActionCodec2 artifact."""\n',
            encoding="utf-8",
        )
        entry = self.package_root / "hub_entry.py"
        if not entry.is_file():
            entry = self.source / "runtime" / "hub_entry.py"
        target = self.output / "processing_actioncodec2.py"
        if entry.resolve() != target.resolve():
            shutil.copyfile(entry, target)
        files.append(str(target))
        # A saved artifact needs runtime dependencies, but not training CLIs or
        # package-build metadata.
        for name in ("requirements.txt",):
            content = SUPPORT_FILES[name]
            target = self.output / name
            target.write_text(content, encoding="utf-8")
            files.append(str(target))
        if not (self.output / "README.md").exists():
            (self.output / "README.md").write_text(USER_GUIDE, encoding="utf-8")
        return files


USER_GUIDE = """# Loading this saved artifact

This directory is a complete, portable ActionCodec2 artifact. Install its runtime dependencies, then point Transformers at this directory. Keep every file and subdirectory together when copying or uploading it.

```bash
python -m pip install -r requirements.txt
```

```python
from transformers import AutoProcessor

codec = AutoProcessor.from_pretrained(
    "ZibinDong/ActionCodec2-2nd-order",
    trust_remote_code=True,
)
print(codec)
```

# ActionCodec2

ActionCodec2 converts continuous robot actions into integer token IDs and decodes
them to the same column layout. Physical quantization is lossy; token
serialization is exact for a fixed fitted codec. The Python `ActionCodec2` class
is the main interface. Its input is an action tensor, not a text prompt.

Python 3.10+ is required.

```bash
python -m pip install .
```

## Choose an action space before encoding

Action shape alone cannot identify units or column meanings. Confirm the action
order, units, rotation representation, absolute versus delta convention,
gripper convention, and sampling rate from your data source. `D=7` can match
several different layouts.

```python
from actioncodec2 import ActionCodec2

# No artifact is needed to inspect the built-in layout definitions.
print(ActionCodec2.builtin_action_spaces().keys())

ARTIFACT = "ZibinDong/ActionCodec2-2nd-order"  # or a Hub repository ID
print(ActionCodec2.describe_pretrained(ARTIFACT, action_dim=7))
codec = ActionCodec2.from_pretrained(ARTIFACT, action_space="single_eef_delta")
print(codec)  # All registered layouts, column meanings, units, and requirements
codec.print_action_spaces(action_dim=7)  # Candidates only; never auto-selects
```

A saved artifact can include its selected `action_space`; in that case
`ActionCodec2.from_pretrained(ARTIFACT)` restores it. For an unbound router,
`action_space` is required. You can inspect `codec.action_space` for the selected
contract and `codec.action_spaces` for every registered contract. Use
`codec.for_action_space("NAME")` to create another binding over the same fitted
profiles. The original object keeps its binding.

To define a layout with explicit semantics, register a preset mapping or YAML
file. Start from a known contract only when its column meanings really match:

```python
preset = codec.action_space.to_dict()
preset["name"] = "my_30hz_eef_delta"
preset["source_hz"] = 30.0
codec.register_action_space(preset)
codec = codec.for_action_space("my_30hz_eef_delta")
codec.save_pretrained("./my-bound-codec")  # Persists the custom layout and binding
```

Changing `source_hz` does not change the meaning of action columns. For a new
robot or a different physical tensor, define the complete preset instead of
renaming an existing layout.

## Encode and decode

```python
import numpy as np

actions = np.asarray(your_actions, dtype=np.float32)  # (T, D) or (B, T, D)
tokens = codec.encode(actions, fps=30)  # list[list[int]], one row per episode
reconstructed = codec.decode(tokens, fps=30)  # CPU torch.float32, (B, T_out, D)
features = codec(actions, fps=30)  # Transformers BatchFeature with input_ids
```

`T` can vary across calls. For unequal episode lengths, encode episodes one at a
time; a single tensor batch must have equal `T`. Pass the real sampling rate with
`fps`. If omitted, the codec warns once and assumes the selected preset's
`source_hz`. Decoding at a different rate resamples the reconstructed path, so
`T_out` may differ from the original `T`. Physical quantization and time
resampling prevent exact numeric round trips.

For independently generated rows with different decoded horizons, use
`codec.batch_decode(rows, fps=30)`. It returns a list of tensors shaped
`(1, T_i, D)`. With `errors="return"`, an invalid row becomes a `ValueError` in
its list position without aborting other rows. Ordinary `decode` expects a
regular batch with the same decoded length.

### Absolute actions and second-order profiles

An absolute layout needs the physical state immediately before the first action.
Pass the same `current_state` to encode and decode; it has shape `(B, D)` even
when the input is one `(T, D)` episode:

```python
tokens = codec.encode(actions, fps=30, current_state=current_state)
reconstructed = codec.decode(tokens, fps=30, current_state=current_state)
```

A second-order profile additionally needs the preceding reconstructed
first-order increment for each active component. Inspect
`codec.requires_previous_first_order`. The boundary mapping has the form
`{component_name: (B, D_active)}` at the codec's target rate. It is independent
of `current_state`.

```python
# Only use zeros when the preceding motion is known to be zero.
boundary = codec.first_order_template(batch_size=1)
tokens = codec.encode(actions, fps=30, previous_first_order=boundary)
reconstructed = codec.decode(tokens, fps=30, previous_first_order=boundary)
next_boundary = codec.final_first_order(
    tokens, previous_first_order=boundary, executed_steps=None
)
```

`final_first_order` returns the boundary for the next chunk. `executed_steps`
can select an executed prefix in codec-rate steps. For an absolute second-order
layout, pass both `current_state` and `previous_first_order` to encode/decode.
The library never invents a boundary from the first action.

For diagnostics, `codec.is_fitted`, `codec.fitted_profiles`, and
`codec.requires_previous_first_order` report whether the selected artifact is
ready for inference and whether the second-order boundary is mandatory.

## Fit and save

Fit on actions whose columns follow the selected layout. The same `ActionCodec2`
object is the training interface:

```python
codec = ActionCodec2(action_space="single_eef_delta")
codec.fit(episodes, fps=15, backend="auto")
print(codec.is_fitted, codec.fitted_profiles)
print(codec.training_report)
codec.save_pretrained("./my-codec")
```

`episodes` can be one `(T, D)` array, one regular `(B, T, D)` array, or an
iterable of variable-length episodes. For mixed layouts or sampling rates, pass
records with explicit metadata:

```python
records = [
    {"actions": episode_a, "action_space": "single_eef_delta", "fps": 15},
    {"actions": episode_b, "action_space": "single_joint_delta", "fps": 30},
]
codec.fit(records, backend="auto")
```

A record may also contain `current_state` and `previous_first_order` as required
by its layout. `fit` supports `primitive_order=2` codecs, an explicit
`bootstrap_context=True` option for two measured context frames, and
`profile_configs={"joint": ..., "eef": ...}` at construction for custom
physical profile settings. `backend` is `"auto"`, `"native"`, or `"python"`;
`threads`, `batch_size`, `progress`, `max_cell_bytes`, and `work_dir` control
resource use. Fitting replaces token assignments; re-encode datasets after
refitting. `save_pretrained` writes the complete loadable artifact, including
the layout binding and model card. `save_pretrained(path, push_to_hub=True,
repo_id="ZibinDong/my-actioncodec2")` also uploads it to a Hub repository you control.

The optional C++ extension accelerates fitting when available. `backend="auto"`
uses it if installed and otherwise uses the Python implementation.

### Reuse separately fitted physical profiles

If you already have fitted joint and EEF profile directories, assemble a router
without refitting. Budgets come from the recorded training configuration unless
you supply `token_budgets`:

```python
codec = ActionCodec2.from_profiles(
    "./joint-profile", "./eef-profile", action_space="single_eef_delta"
)
codec.save_pretrained("./combined-codec")
```

The equivalent command is:

```bash
actioncodec2-bundle --joint ./joint-profile --eef ./eef-profile \\
  --action-space single_eef_delta --out ./combined-codec
```

The command also accepts `--joint-budget` and `--eef-budget`. Omit
`--action-space` only when you intentionally want to select a layout at load
time.

## Constrain model generation

The codec supplies a grammar for an exact action horizon; it does not generate
actions by itself. `codec.grammar(horizon, fps=...)` has
`next_token_mask(prefix)` and `is_complete(tokens)`. A truncated or otherwise
invalid stream cannot be decoded as a complete action chunk.

```python
horizon = 12
grammar = codec.grammar(horizon, fps=15)
max_action_tokens = codec.max_tokens_for_horizon(horizon, fps=15)
assert grammar.is_complete(tokens)
```

For a Transformers model, use the adapter returned by
`codec.logits_processor(...)` in `generate`:

```python
from transformers import LogitsProcessorList

constraint = codec.logits_processor(
    horizon,
    fps=15,
    prompt_length=input_ids.shape[1],
    eos_token_id=model.config.eos_token_id,
    token_offset=action_token_offset,
)
output = model.generate(
    input_ids,
    logits_processor=LogitsProcessorList([constraint]),
)
```

The model must already have a contiguous range of action IDs beginning at
`action_token_offset`; the codec does not modify its vocabulary. EOS and any
padding ID must lie outside that range. `codec.vocab_size` is the number of
action IDs. `codec.new_action_tokens` provides their ordered names, and
`codec.to_token_strs(one_token_row)` renders a row for inspection. Remove the
model's token offset and EOS/padding IDs before passing generated action IDs to
`codec.decode`.

## Advanced profile commands

The recommended training path is `codec.fit`. The following commands operate
on physical profile tensors rather than raw action-space layouts. Use them when
you intentionally need to fit or audit individual joint/EEF profiles. Their
`.npy`/`.npz` inputs must already contain canonical `(T, D_profile)` or
`(B, T, D_profile)` trajectories; `actioncodec2-fit-profile` does not perform
layout routing or source-rate conversion.

```bash
actioncodec2-fit-profile --config ./eef-config.yaml --data ./eef.npy \\
  --output-dir ./eef-profile
actioncodec2-audit-profile stats --config ./eef-config.yaml \\
  --data ./eef.npy --output ./eef-stats.json
actioncodec2-audit-profile validate --artifact ./eef-profile \\
  --data ./eef.npy --output ./eef-validation.json
```

`--config` is an `ActionCodec2` physical profile YAML with a `train` section.
The fit command writes a physical profile and a JSON report. `stats` summarizes
source trajectories before fitting; `validate` checks a fitted profile against
data and writes a machine-readable report. Run each command with `--help` for
resource limits, metadata fields, and overwrite controls.

## Public interface at a glance

| Task | Entry point | Result |
| --- | --- | --- |
| Inspect layouts | `ActionCodec2.builtin_action_spaces()`, `describe_pretrained()`, `print(codec)`, `codec.print_action_spaces()` | Explicit physical contracts |
| Load and switch | `from_pretrained()`, `for_action_space()`, `register_action_space()` | Bound `ActionCodec2` |
| Fit and package | `ActionCodec2(...)`, `fit()`, `from_profiles()`, `save_pretrained()` | Loadable artifact |
| Encode and decode | `encode()`, `decode()`, `batch_decode()`, `codec(...)` | Token rows or action tensors |
| Stateful chunks | `first_order_template()`, `final_first_order()` | Explicit second-order boundaries |
| Generate | `grammar()`, `logits_processor()`, `max_tokens_for_horizon()` | Valid token constraints |
| Inspect tokens | `vocab_size`, `new_action_tokens`, `to_token_strs()` | ID-to-name mapping |

`AutoProcessor.from_pretrained(path, trust_remote_code=True)` loads an artifact
without importing this package first. If the package is installed and imported,
its processor registration also supports local `AutoProcessor` resolution.

## Tests

```bash
python -m pytest tests -q
```

## Artifact contents

A saved artifact contains the metadata, fitted profiles, private runtime, and
this guide. Keep the complete directory together when moving it or uploading it
to Hugging Face.
"""

SUPPORT_FILES = {
    "requirements.txt": "numpy\nscipy\ntorch>=2.3\ntransformers>=4.57,<6\nhuggingface-hub\npyyaml\n",
}

MODULE_FILES = [
    "backend/__init__.py",
    "backend/native_backend.py",
    "common/__init__.py",
    "common/_frontend_base.py",
    "common/_types.py",
    "common/_validation.py",
    "frontend/__init__.py",
    "frontend/tokenizer.py",
    "integration/__init__.py",
    "integration/exporting.py",
    "integration/hf.py",
    "physical/__init__.py",
    "physical/_so3.py",
    "physical/boundary.py",
    "physical/configuration.py",
    "physical/physical_codec.py",
    "physical/physical_processor.py",
    "physical/primitives.py",
    "physical/semantic.py",
    "physical/trajectory.py",
    "routing/__init__.py",
    "routing/cache.py",
    "routing/defaults.py",
    "routing/presets.py",
    "routing/routing.py",
    "tokenization/__init__.py",
    "tokenization/bpe.py",
    "tokenization/bpe_training.py",
    "tokenization/bpe_vocabulary.py",
    "tokenization/generation.py",
    "tokenization/grammar.py",
    "tokenization/setbpe_spatial.py",
    "tokenization/setbpe_spatial_training.py",
    "tokenization/setbpe_temporal.py",
    "tokenization/setbpe_temporal_training.py",
    "training/__init__.py",
    "training/fitting.py",
    "training/training.py",
    "processing_actioncodec2.py",
]