# Created by BVH (assisted), Jun 2026.
# action_format VERSION 2: YAMxdof joints -> EE pose via FK (first real divergence from v1).
#
# Selected via a4d_config data_defaults.action_format_version='action_format_v2' (default 'action_format_v1').
# v1 packs YAM raw joint angles into the 20D vector (JOINT-SPACE EXCEPTION); v2 runs forward
# kinematics through the official i2rt yam.urdf (see custom/dataset/yam_fk.py) so YAM emits the
# same xyz+rot6d+grip per-arm semantics as every other bimanual dataset. All other datasets
# dispatch to the exact v1 parsers (unchanged).

import numpy as np
from rich import print

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


def _parse_yam_ee(entry):
    '''
    YAMxdof with EE-pose semantics: FK(6 joint angles) -> flange (link_6) pose per arm.
    action = FK(arm_leader) = commanded EE pose; actual = FK(arm_proprio[:6]) = measured;
    desired = None. Grippers from eef_leader / eef_proprio[0] (~1=open).
    CAVEATS: poses live in PER-ARM base frames (no left/right rig calibration exists), and
    there is no TCP offset (XDOF gripper type unrecorded; fingertips ~13.5cm past flange).
    '''
    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), [right, left]
        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(v2): 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}


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