# Created by BVH (assisted), Jun 2026.
'''
Dependency-free numpy forward kinematics for the I2RT YAM arm (YAMxdof dataset).

URDF: custom/dataset/assets/yam.urdf, verbatim (md5 f339ccff) from the official vendor repo
github.com/i2rt-robotics/i2rt @ main, i2rt/robot_models/arm/yam/yam.urdf (MIT). Same model as
MuJoCo Menagerie i2rt_yam. Cross-validated against pytorch_kinematics; see
_my-priv-util/a4d2/misc/claude/63_yamxdof_tools.py validate-fk.

FK output = flange (link_6) pose in the arm BASE frame. No TCP offset (XDOF gripper type is
unrecorded; linear_4310 fingertips would be ~13.5cm past the flange). Joint order = data
motor-ID order = URDF chain order joint1..joint6 (limits match per-dim on real data).
'''

import os
import xml.etree.ElementTree as ET

import numpy as np

YAM_URDF_PATH = os.path.join(os.path.dirname(__file__), 'assets', 'yam.urdf')


def _rpy_to_matrix(rpy):
    '''
    URDF fixed-axis XYZ rpy to rotation matrix: R = Rz(yaw) @ Ry(pitch) @ Rx(roll).
    '''
    r, p, y = float(rpy[0]), float(rpy[1]), float(rpy[2])
    cr, sr = np.cos(r), np.sin(r)
    cp, sp = np.cos(p), np.sin(p)
    cy, sy = np.cos(y), np.sin(y)
    Rx = np.array([[1, 0, 0], [0, cr, -sr], [0, sr, cr]])
    Ry = np.array([[cp, 0, sp], [0, 1, 0], [-sp, 0, cp]])
    Rz = np.array([[cy, -sy, 0], [sy, cy, 0], [0, 0, 1]])
    return Rz @ Ry @ Rx


def _axis_angle_to_matrix(axis, angle):
    '''
    Rodrigues formula, batched: axis (3,), angle (...,) -> R (..., 3, 3).
    '''
    axis = np.asarray(axis, dtype=np.float64)
    axis = axis / np.linalg.norm(axis)
    K = np.array([
        [0.0, -axis[2], axis[1]],
        [axis[2], 0.0, -axis[0]],
        [-axis[1], axis[0], 0.0],
    ])
    angle = np.asarray(angle, dtype=np.float64)[..., None, None]
    eye = np.eye(3)
    return eye + np.sin(angle) * K + (1.0 - np.cos(angle)) * (K @ K)


class SerialChainFK:
    '''
    Generic serial-chain FK from a URDF: T_tip = prod_i [Origin_i @ Rot(axis_i, q_i)] along the
    unique parent->child revolute chain from root to tip. Link visual/inertial origins are
    irrelevant for FK. Fixed joints along the chain are composed without a dof.
    '''

    def __init__(self, urdf_path, root='base_link', tip='link_6'):
        joints = {}
        for j in ET.parse(urdf_path).getroot().findall('joint'):
            child = j.find('child').attrib['link']
            origin = j.find('origin')
            xyz = [float(v) for v in origin.attrib.get('xyz', '0 0 0').split()]
            rpy = [float(v) for v in origin.attrib.get('rpy', '0 0 0').split()]
            T = np.eye(4)
            T[:3, :3] = _rpy_to_matrix(rpy)
            T[:3, 3] = xyz
            axis_el = j.find('axis')
            axis = [float(v) for v in axis_el.attrib['xyz'].split()] if axis_el is not None else [0, 0, 1]
            joints[child] = dict(
                name=j.attrib['name'], parent=j.find('parent').attrib['link'],
                type=j.attrib['type'], origin=T, axis=np.asarray(axis, dtype=np.float64))

        # Walk tip -> root via parent pointers, then reverse to chain order.
        chain = []
        link = tip
        while link != root:
            if link not in joints:
                raise ValueError(f'No joint chain from {root} to {tip} in {urdf_path} (stuck at {link})')
            chain.append(joints[link])
            link = joints[link]['parent']
        chain.reverse()
        self.chain = chain
        self.joint_names = [j['name'] for j in chain if j['type'] == 'revolute']
        self.num_dofs = len(self.joint_names)

    def fk(self, q):
        '''
        :param q: (..., num_dofs) joint angles [rad].
        :return T: (..., 4, 4) tip pose in the root frame.
        '''
        q = np.asarray(q, dtype=np.float64)
        assert q.shape[-1] == self.num_dofs, f'expected {self.num_dofs} dofs, got {q.shape}'
        T = np.broadcast_to(np.eye(4), q.shape[:-1] + (4, 4)).copy()
        qi = 0
        for joint in self.chain:
            T = T @ joint['origin']
            if joint['type'] == 'revolute':
                R = _axis_angle_to_matrix(joint['axis'], q[..., qi])
                T4 = np.broadcast_to(np.eye(4), R.shape[:-2] + (4, 4)).copy()
                T4[..., :3, :3] = R
                T = T @ T4
                qi += 1
        return T


_YAM_CHAIN = None


def yam_fk(q):
    '''
    YAM flange FK with a lazily-built singleton chain (safe for dataloader workers).
    :param q: (..., 6) joint angles [rad] in data motor-ID order.
    :return T: (..., 4, 4) link_6 pose in the arm base frame.
    '''
    global _YAM_CHAIN
    if _YAM_CHAIN is None:
        _YAM_CHAIN = SerialChainFK(YAM_URDF_PATH)
    return _YAM_CHAIN.fk(q)