"""Policy transforms for the Unitree G1 + Inspire-hands red-ball dataset.

Adapted from libero_policy.py for the GR00T-format LeRobot dataset
``g1-inspire-red-ball-success`` (robot_type: unitree_g1).

Dataset layout (verified from meta/info.json + the parquet):
  observation.state : 63 dims = [29 proprio | 34 tactile]
      proprio 0:29 = left_arm 0:7, right_arm 7:14, left_hand 14:20,
                     right_hand 20:26, waist 26:29
  action            : 30 dims = left_arm 0:7, right_arm 7:14,
                     left_hand 14:20, right_hand 20:26, base_height 26:27,
                     navigate 27:30  (ABSOLUTE joint targets, verified:
                     mean|action-state| over arms = 0.047 rad)
  observation.images.ego_view : (240, 424, 3) single ego camera, no wrist

We feed pi0.5 the 29-dim proprio state (pads to the model's 32-dim slot) and
drop tactile (pi0.5 base has no tactile concept and the state slot caps at 32).
Arm actions are converted to deltas in the DataConfig; hands/base/nav stay
absolute. Output truncates back to the dataset's 30 action dims.
"""

import dataclasses

import einops
import numpy as np

from openpi import transforms
from openpi.models import model as _model

# proprio dims kept from the 63-dim observation.state (drop 29:63 tactile)
STATE_DIM = 29
# native action dims of this dataset
ACTION_DIM = 30


def make_g1_inspire_example() -> dict:
    """Random input example matching the inference observation schema."""
    return {
        "observation/state": np.random.rand(STATE_DIM),
        "observation/image": np.random.randint(256, size=(240, 424, 3), dtype=np.uint8),
        "prompt": "place the red ball in the box",
    }


def _parse_image(image) -> np.ndarray:
    image = np.asarray(image)
    if np.issubdtype(image.dtype, np.floating):
        image = (255 * image).astype(np.uint8)
    if image.shape[0] == 3:
        image = einops.rearrange(image, "c h w -> h w c")
    return image


@dataclasses.dataclass(frozen=True)
class G1InspireInputs(transforms.DataTransformFn):
    """Convert dataset/inference inputs into the pi0 model input schema."""

    model_type: _model.ModelType

    def __call__(self, data: dict) -> dict:
        # Keep only the proprio prefix; tactile (29:63) is dropped.
        state = np.asarray(data["observation/state"])[..., :STATE_DIM]

        base_image = _parse_image(data["observation/image"])

        inputs = {
            "state": state,
            "image": {
                # single ego camera -> third-person slot; no wrist cameras.
                "base_0_rgb": base_image,
                "left_wrist_0_rgb": np.zeros_like(base_image),
                "right_wrist_0_rgb": np.zeros_like(base_image),
            },
            "image_mask": {
                "base_0_rgb": np.True_,
                # pi0 (flow) masks padding images False; pi0-FAST expects True.
                "left_wrist_0_rgb": np.True_ if self.model_type == _model.ModelType.PI0_FAST else np.False_,
                "right_wrist_0_rgb": np.True_ if self.model_type == _model.ModelType.PI0_FAST else np.False_,
            },
        }

        if "actions" in data:
            inputs["actions"] = data["actions"]
        if "prompt" in data:
            inputs["prompt"] = data["prompt"]
        return inputs


@dataclasses.dataclass(frozen=True)
class G1InspireOutputs(transforms.DataTransformFn):
    """Truncate model actions back to the dataset's 30 action dims."""

    def __call__(self, data: dict) -> dict:
        return {"actions": np.asarray(data["actions"][..., :ACTION_DIM])}