# Created by yams_any4d agent, 2026-07-12.
# action_format VERSION 3: YAMxdof joints -> grasp_site (TCP) pose via FK + a constant
# tool offset. v2 FKs to the flange (link_6); the yam_4310_linear gripper's toolpoint
# (`grasp_site` in the mujoco XML used by the DINO model + inference) is ~13.8cm past the
# flange (yam_fk docstring warns "fingertips ~13.5cm past flange"). Cameron's DINO dataset
# (yam_local_train_mirror/lib/dataset.py `_fk_ee_in_rbase`) uses that grasp_site as the
# target keypoint; this format matches it so v4head targets/keypoints land on the gripper.
#
# C = flange->grasp_site, a CONSTANT rigid transform (grasp_site is rigidly attached to
# link_6). Measured as inv(yam_fk(q)) @ mujoco_grasp_site(q) over 30 random configs:
# rotation constant to 1e-5 (clean 90deg about z), translation |t|=0.1372m (URDF-vs-MJCF
# model diff gives ~10mm spread — negligible vs the 138mm error this fixes). Applying C in
# numpy needs no runtime mujoco. Only the yam_4310_linear gripper uses this offset, so this
# is a NEW format (v2 unchanged) selected via data_defaults.action_format_version.

import numpy as np
from rich import print

import custom.dataset.action_format_v1 as v1
import custom.dataset.action_format_v2 as v2
from custom.dataset.yam_fk import yam_fk

# flange (link_6) -> grasp_site (yam_4310_linear TCP), constant. See module docstring.
_FLANGE_TO_GRASP = np.array([
    [0.0,  1.0, 0.0, 0.001215],
    [-1.0, 0.0, 0.0, 0.003868],
    [0.0,  0.0, 1.0, 0.137180],
    [0.0,  0.0, 0.0, 1.0],
], dtype=np.float64)


def _parse_yam_ee_grasp(entry):
    '''YAMxdof with grasp_site (TCP) EE-pose semantics: FK(6 joints) -> flange, then apply
    the constant flange->grasp_site tool offset. Poses in per-arm base frame; grippers from
    eef_leader / eef_proprio[0].'''
    def _fk_pack(d, right_key, left_key, right_grip, left_grip):
        q = np.stack([v1._to_np(d[right_key])[:6], v1._to_np(d[left_key])[:6]])  # (2, 6)
        T = yam_fk(q)  # (2, 4, 4) flange, [right, left]
        T = T @ _FLANGE_TO_GRASP  # (2, 4, 4) grasp_site (broadcast over arm dim)
        r_xyz, r_rot6d = v1._homogeneous_to_xyz_rot6d(T[0])
        l_xyz, l_rot6d = v1._homogeneous_to_xyz_rot6d(T[1])
        r_g = float(v1._to_np(right_grip)[0]) if right_grip is not None else 0.0
        l_g = float(v1._to_np(left_grip)[0]) if left_grip is not None else 0.0
        return v1._pack_bimanual(r_xyz, r_rot6d, l_xyz, l_rot6d, r_g, l_g)

    actions_dict = entry.get('actions', {})
    states_dict = entry.get('states', {})

    if 'right_arm_leader' in actions_dict and 'left_arm_leader' in actions_dict:
        action = _fk_pack(actions_dict, 'right_arm_leader', 'left_arm_leader',
                          actions_dict.get('right_eef_leader'), actions_dict.get('left_eef_leader'))
    else:
        print(f'[yellow]  YAMxdof(v3): no arm_leader keys in actions, keys={list(actions_dict.keys())}; '
              f'zero action[/yellow]')
        action = np.zeros(20, dtype=np.float32)

    actual = None
    if 'right_arm_proprio' in states_dict and 'left_arm_proprio' in states_dict:
        actual = _fk_pack(states_dict, 'right_arm_proprio', 'left_arm_proprio',
                          states_dict.get('right_eef_proprio'), states_dict.get('left_eef_proprio'))

    return {'action': action, 'desired': None, 'actual': actual}


_DATASET_PARSERS = {**v1._DATASET_PARSERS, 'yam': _parse_yam_ee_grasp}


def parse_anydata_raw_robot_action(sample, dset_name=''):
    '''v3 entry point: v1 dispatch with the YAM parser swapped to grasp_site EE-pose semantics.'''
    return v1.parse_anydata_raw_robot_action(sample, dset_name, parsers=_DATASET_PARSERS)