src/sata/utils/joint2humanml.py

"""
Joint to HumanML3D Representation Converter

This module provides a class to convert joint positions (shape: [T, 22, 3])
to HumanML3D representation format (shape: [T, 263]).
"""

import numpy as np
import torch
from os.path import join as pjoin
from .humanml_skeleton import Skeleton
from .quaternion import *
from .humanml_paramUtil import t2m_raw_offsets, t2m_kinematic_chain


class JointToHumanML3D:
    """
    Convert joint positions to HumanML3D representation format.
    
    The output format is a 263-dimensional feature vector containing:
    - Root rotation velocity (1D): rotation velocity along y-axis
    - Root linear velocity (2D): linear velocity on xz plane
    - Root height (1D): y-coordinate of root joint
    - Joint rotation invariant positions (63D): local positions of 21 joints
    - Joint rotations (126D): continuous 6D rotation representation for 21 joints
    - Joint velocities (66D): local velocities of 22 joints
    - Foot contacts (4D): contact labels for left/right foot
    
    Total: 1 + 2 + 1 + 63 + 126 + 66 + 4 = 263
    """
    
    def __init__(self, example_id="000021", data_dir='./joints/'):
        """
        Initialize the converter with pre-loaded target skeleton.
        
        Args:
            example_id (str): The example motion ID to extract target skeleton offsets
            data_dir (str): Directory containing the example motion data
        """
        # Lower legs indices for scale calculation
        self.l_idx1, self.l_idx2 = 5, 8
        
        # Right/Left foot indices for foot contact detection
        self.fid_r, self.fid_l = [8, 11], [7, 10]
        
        # Face direction joints: r_hip, l_hip, sdr_r, sdr_l
        self.face_joint_indx = [2, 1, 17, 16]
        
        # Hip indices
        self.r_hip, self.l_hip = 2, 1
        
        # Number of joints
        self.joints_num = 22
        
        # Load kinematic chain and raw offsets
        self.n_raw_offsets = torch.from_numpy(t2m_raw_offsets)
        self.kinematic_chain = t2m_kinematic_chain
        
        # Get target skeleton offsets from example
        example_data = np.load(pjoin(data_dir, example_id + '.npy'))
        example_data = example_data.reshape(len(example_data), -1, 3)
        example_data = torch.from_numpy(example_data)
        
        tgt_skel = Skeleton(self.n_raw_offsets, self.kinematic_chain, 'cpu')
        self.tgt_offsets = tgt_skel.get_offsets_joints(example_data[0])
        
        print(f"JointToHumanML3D initialized with target skeleton from {example_id}")
    
    def uniform_skeleton(self, positions, target_offset):
        """
        Normalize skeleton to target proportions using leg length scaling.
        
        Args:
            positions (np.ndarray): Joint positions, shape [T, joints_num, 3]
            target_offset (torch.Tensor): Target skeleton offsets
            
        Returns:
            np.ndarray: Normalized joint positions
        """
        src_skel = Skeleton(self.n_raw_offsets, self.kinematic_chain, 'cpu')
        src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0]))
        src_offset = src_offset.numpy()
        tgt_offset = target_offset.numpy()
        
        # Calculate scale ratio based on leg lengths
        src_leg_len = np.abs(src_offset[self.l_idx1]).max() + np.abs(src_offset[self.l_idx2]).max()
        tgt_leg_len = np.abs(tgt_offset[self.l_idx1]).max() + np.abs(tgt_offset[self.l_idx2]).max()
        
        scale_rt = tgt_leg_len / src_leg_len
        
        src_root_pos = positions[:, 0]
        tgt_root_pos = src_root_pos * scale_rt
        
        # Inverse Kinematics
        quat_params = src_skel.inverse_kinematics_np(positions, self.face_joint_indx)
        
        # Forward Kinematics with target skeleton
        src_skel.set_offset(target_offset)
        new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos)
        
        return new_joints
    
    def process_joints(self, positions, feet_thre=0.002):
        """
        Process joint positions and convert to HumanML3D representation.
        
        Args:
            positions (np.ndarray): Joint positions, shape [T, joints_num, 3]
            feet_thre (float): Threshold for foot contact detection
            
        Returns:
            np.ndarray: HumanML3D representation, shape [T-1, 263]
        """
        # Uniform skeleton normalization
        positions = self.uniform_skeleton(positions, self.tgt_offsets)
        
        # Put on floor
        floor_height = positions.min(axis=0).min(axis=0)[1]
        positions[:, :, 1] -= floor_height
        
        # Center XZ at origin
        root_pos_init = positions[0]
        root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1])
        positions = positions - root_pose_init_xz
        
        # Align all poses to initially face Z+
        r_hip, l_hip, sdr_r, sdr_l = self.face_joint_indx
        across1 = root_pos_init[r_hip] - root_pos_init[l_hip]
        across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l]
        across = across1 + across2
        across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis]
        
        # Calculate forward direction (rotate around y-axis)
        forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
        forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis]
        
        # Rotate to face Z+
        target = np.array([[0, 0, 1]])
        root_quat_init = qbetween_np(forward_init, target)
        root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init
        
        positions = qrot_np(root_quat_init, positions)
        
        # Store global positions
        global_positions = positions.copy()
        
        # Detect foot contacts
        feet_l, feet_r = self._foot_detect(positions, feet_thre)
        
        # Get continuous 6D representation
        cont_6d_params, r_velocity, velocity, r_rot = self._get_cont6d_params(positions)
        
        # Get rotation invariant position representation
        positions = self._get_rifke(positions, r_rot)
        
        # Root height
        root_y = positions[:, 0, 1:2]
        
        # Root rotation and linear velocity
        r_velocity = np.arcsin(r_velocity[:, 2:3])  # (T-1, 1)
        l_velocity = velocity[:, [0, 2]]  # (T-1, 2)
        root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)  # (T-1, 4)
        
        # Joint rotation representation
        rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)  # (T, 126)
        
        # Joint rotation invariant position representation
        ric_data = positions[:, 1:].reshape(len(positions), -1)  # (T, 63)
        
        # Joint velocity representation
        local_vel = qrot_np(
            np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),
            global_positions[1:] - global_positions[:-1]
        )
        local_vel = local_vel.reshape(len(local_vel), -1)  # (T-1, 66)
        
        # Concatenate all features
        data = root_data  # (T-1, 4)
        data = np.concatenate([data, ric_data[:-1]], axis=-1)  # (T-1, 4+63)
        data = np.concatenate([data, rot_data[:-1]], axis=-1)  # (T-1, 4+63+126)
        data = np.concatenate([data, local_vel], axis=-1)  # (T-1, 4+63+126+66)
        data = np.concatenate([data, feet_l, feet_r], axis=-1)  # (T-1, 263)
        
        return data
    
    def _foot_detect(self, positions, thres):
        """
        Detect foot contacts based on velocity threshold.
        
        Args:
            positions (np.ndarray): Joint positions
            thres (float): Velocity threshold
            
        Returns:
            tuple: (feet_l, feet_r) contact labels for left and right feet
        """
        velfactor = np.array([thres, thres])
        
        # Left foot
        feet_l_x = (positions[1:, self.fid_l, 0] - positions[:-1, self.fid_l, 0]) ** 2
        feet_l_y = (positions[1:, self.fid_l, 1] - positions[:-1, self.fid_l, 1]) ** 2
        feet_l_z = (positions[1:, self.fid_l, 2] - positions[:-1, self.fid_l, 2]) ** 2
        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float32)
        
        # Right foot
        feet_r_x = (positions[1:, self.fid_r, 0] - positions[:-1, self.fid_r, 0]) ** 2
        feet_r_y = (positions[1:, self.fid_r, 1] - positions[:-1, self.fid_r, 1]) ** 2
        feet_r_z = (positions[1:, self.fid_r, 2] - positions[:-1, self.fid_r, 2]) ** 2
        feet_r = ((feet_r_x + feet_r_y + feet_r_z) < velfactor).astype(np.float32)
        
        return feet_l, feet_r
    
    def _get_cont6d_params(self, positions):
        """
        Get continuous 6D rotation parameters.
        
        Args:
            positions (np.ndarray): Joint positions
            
        Returns:
            tuple: (cont_6d_params, r_velocity, velocity, r_rot)
        """
        skel = Skeleton(self.n_raw_offsets, self.kinematic_chain, "cpu")
        quat_params = skel.inverse_kinematics_np(positions, self.face_joint_indx, smooth_forward=True)
        
        # Quaternion to continuous 6D
        cont_6d_params = quaternion_to_cont6d_np(quat_params)
        
        # Root rotation
        r_rot = quat_params[:, 0].copy()
        
        # Root linear velocity
        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
        velocity = qrot_np(r_rot[1:], velocity)
        
        # Root angular velocity
        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
        
        return cont_6d_params, r_velocity, velocity, r_rot
    
    def _get_rifke(self, positions, r_rot):
        """
        Get rotation invariant position representation.
        
        Args:
            positions (np.ndarray): Joint positions
            r_rot (np.ndarray): Root rotations
            
        Returns:
            np.ndarray: Rotation invariant positions
        """
        # Local pose (relative to root XZ)
        positions[..., 0] -= positions[:, 0:1, 0]
        positions[..., 2] -= positions[:, 0:1, 2]
        
        # All poses face Z+
        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)
        
        return positions
    
    def convert(self, joint_positions, valid_frames):
        """
        Convert joint positions to HumanML3D representation.
        
        Args:
            joint_positions (np.ndarray or torch.Tensor): Joint positions with shape [T, 22, 3]
            valid_frames (int): Number of valid frames to use. If provided and less than T,
                                         only the first valid_frames will be processed, and the output
                                         will be padded to match the expected length [T-1, 263].
                                         If None, all frames are processed.
            
        Returns:
            np.ndarray: HumanML3D representation with shape [T-1, 263]
        """
        # Convert to numpy if input is torch tensor
        if isinstance(joint_positions, torch.Tensor):
            joint_positions = joint_positions.cpu().numpy()
        
        # Validate input shape
        if len(joint_positions.shape) != 3:
            raise ValueError(f"Expected 3D input with shape [T, 22, 3], got shape {joint_positions.shape}")
        
        if joint_positions.shape[1] != 22 or joint_positions.shape[2] != 3:
            raise ValueError(f"Expected shape [T, 22, 3], got {joint_positions.shape}")
        
        original_length = joint_positions.shape[0]
        
        # Truncate to valid frames if specified
        if valid_frames is not None and valid_frames < original_length:
            joint_positions = joint_positions[:valid_frames, ...]
        else:
            valid_frames = original_length

        # Process joints (this will reduce length by 1 due to velocity calculation)
        humanml3d_repr = self.process_joints(joint_positions, feet_thre=0.002)
        # humanml3d_repr shape: [valid_frames-1, 263]

        # Padding to match expected output length [original_length-1, 263]
        # if frames were truncated
        if valid_frames < original_length:
            target_length = original_length - 1
            current_length = humanml3d_repr.shape[0]  # valid_frames - 1
            
            if current_length < target_length:
                # Pad with the last valid frame
                pad_length = target_length - current_length
                last_frame = humanml3d_repr[-1:, :]  # [1, 263]
                padding = np.repeat(last_frame, pad_length, axis=0)  # [pad_length, 263]
                humanml3d_repr = np.concatenate([humanml3d_repr, padding], axis=0)
                
        return humanml3d_repr


if __name__ == "__main__":
    # Example usage
    converter = JointToHumanML3D(example_id="000021", data_dir='./joints/')
    
    # Load example joint data
    example_joints = np.load('./joints/000021.npy')[:, :22, :]
    example_joints = example_joints.reshape(-1, 22, 3)
    
    print(f"Input shape: {example_joints.shape}")
    
    # Test 1: Convert without truncation
    humanml3d_data = converter.convert(example_joints, valid_frames=None)
    print(f"Output shape (no truncation): {humanml3d_data.shape}")
    print(f"Expected output shape: [{example_joints.shape[0]-1}, 263]")
    
    # Test 2: Convert with truncation and padding
    if example_joints.shape[0] > 50:
        valid_frames = 50
        humanml3d_data_truncated = converter.convert(example_joints, valid_frames=valid_frames)
        print(f"\nOutput shape (truncated to {valid_frames} frames): {humanml3d_data_truncated.shape}")
        print(f"Expected shape: [{example_joints.shape[0]-1}, 263]")
        print(f"Frames processed: {valid_frames-1}, Frames padded: {example_joints.shape[0]-1-(valid_frames-1)}")