"""xtrainer2 command layout used between StarVLA and the real robot.

The converted LeRobot dataset stores each arm as ``xyz + quaternion + gripper``.
StarVLA / PyTorch3D read that quaternion as ``w, x, y, z``.  The robot HDF5
streams and the robot controller use ``x, y, z, w``.

The training transform already maps quaternion <-> rotation-6D.  After
un-normalization the server therefore has 16-D **wxyz** commands.  Sending
those 16 numbers to the robot without this reorder will execute the wrong
orientation.

Do not confuse these widths:

* 14  – joint angles in the source HDF5 (not used for this VLA)
* 16  – stored / robot EEF command (after this conversion, xyzw)
* 20  – model tensor after rotation-6D (horizon is a separate axis)
"""

from __future__ import annotations

import numpy as np

from starVLA.dataloader.relative_pose import compose_quaternion

QUATERNION_STARTS = (3, 11)
POSITION_SLICES = ((0, 3), (8, 11))
ACTION_DIM = 16
ROBOT_TYPE = "xtrainer2_bimanual_eef"
SOURCE_QUATERNION_FORMAT = "xyzw"
MODEL_QUATERNION_FORMAT = "wxyz"

#: Value of ``DataConfig.robot_command_layout`` that selects the conversions in
#: this module. The policy server reads that attribute instead of matching on a
#: class name.
COMMAND_LAYOUT = "xtrainer2_bimanual_eef_xyzw"


def _as_float_array(values: np.ndarray) -> np.ndarray:
    array = np.asarray(values, dtype=np.float32)
    if array.shape[-1] != ACTION_DIM:
        raise ValueError(f"expected last dimension {ACTION_DIM}, got {array.shape}")
    return array


def wxyz_to_xyzw_block(block: np.ndarray) -> np.ndarray:
    """``w, x, y, z`` -> ``x, y, z, w``."""
    return np.concatenate([block[..., 1:4], block[..., :1]], axis=-1)


def xyzw_to_wxyz_block(block: np.ndarray) -> np.ndarray:
    """``x, y, z, w`` -> ``w, x, y, z``."""
    return np.concatenate([block[..., 3:4], block[..., :3]], axis=-1)


def model_actions_to_robot(values: np.ndarray) -> np.ndarray:
    """Convert un-normalized StarVLA actions (wxyz) to robot commands (xyzw)."""
    converted = np.array(_as_float_array(values), copy=True)
    for start in QUATERNION_STARTS:
        converted[..., start : start + 4] = wxyz_to_xyzw_block(converted[..., start : start + 4])
    return converted


def robot_state_to_model(values: np.ndarray) -> np.ndarray:
    """Convert robot observations (xyzw) to the layout StarVLA was trained on."""
    converted = np.array(_as_float_array(values), copy=True)
    for start in QUATERNION_STARTS:
        converted[..., start : start + 4] = xyzw_to_wxyz_block(converted[..., start : start + 4])
    return converted


def compose_relative_eef_wxyz(relative_actions: np.ndarray, state_wxyz: np.ndarray) -> np.ndarray:
    """Turn rel-mode 16-D wxyz actions into absolute 16-D wxyz poses.

    Positions add the current EE; rotations left-multiply the current
    quaternion; grippers are already absolute and are left unchanged.
    """
    relative = np.array(_as_float_array(relative_actions), copy=True)
    state = np.asarray(state_wxyz, dtype=np.float32).reshape(-1, ACTION_DIM)[0]
    absolute = relative
    for start, end in POSITION_SLICES:
        absolute[..., start:end] = relative[..., start:end] + state[start:end]
    for start in QUATERNION_STARTS:
        absolute[..., start : start + 4] = compose_quaternion(
            state[start : start + 4],
            relative[..., start : start + 4],
        )
    return absolute


def is_relative_action_mode(action_mode: object | None) -> bool:
    if action_mode is None:
        return False
    return str(action_mode).lower() in {"rel", "relative"}