# - "body_transl_delta_pelv"
# - "body_orient_xy"
# - "z_orient_delta"
# - "body_pose"
# - "body_joints_local_wo_z_rot"
from transform3d import transform_body_pose, change_for, remove_z_rot, get_z_rot, rot_diff
from einops import rearrange
import torch

def to_tensor(array):
    if torch.is_tensor(array):
        return array
    else:
        return torch.tensor(array)

def _get_body_transl_delta_pelv(data):
    """
    get body pelvis tranlation delta relative to pelvis coord.frame
    v_i = t_i - t_{i-1} relative to R_{i-1}
    """
    trans = to_tensor(data['trans'])
    trans_vel = trans - trans.roll(1, 0)  # shift one right and subtract
    pelvis_orient = transform_body_pose(to_tensor(data['rots'][..., :3]), "aa->rot")
    trans_vel_pelv = change_for(trans_vel, pelvis_orient.roll(1, 0))
    trans_vel_pelv[0] = 0  # zero out velocity of first frame
    return trans_vel_pelv

def _get_body_orient_xy(data):
    """get body global orientation"""
    # default is axis-angle representation
    pelvis_orient = to_tensor(data['rots'][..., :3])
    # if rot_repr == "6d":
    # axis-angle to rotation matrix & drop last row
    pelvis_orient_xy = remove_z_rot(pelvis_orient, in_format="aa")
    return pelvis_orient_xy

def _get_body_pose(data):
    """get body pose"""
    # default is axis-angle representation: Frames x (Jx3) (J=21)
    pose = to_tensor(data['rots'][..., 3:3 + 21*3])  # drop pelvis orientation
    pose = transform_body_pose(pose, f"aa->6d")
    return pose

def _get_body_joints_local_wo_z_rot(data):
    """get body joint coordinates relative to the pelvis"""
    joints = to_tensor(data['joint_positions'][:, :22, :])
    pelvis_transl = to_tensor(joints[:, 0, :])
    joints_glob = to_tensor(joints[:, :22, :])
    pelvis_orient = to_tensor(data['rots'][..., :3])

    pelvis_orient_z = get_z_rot(pelvis_orient, in_format="aa")
    # pelvis_orient_z = transform_body_pose(pelvis_orient_z, "aa->rot").float()
    # relative_joints = R.T @ (p_global - pelvis_translation)
    rel_joints = torch.einsum('fdi,fjd->fji',
                                pelvis_orient_z,
                                joints_glob - pelvis_transl[:, None, :])

    return rearrange(rel_joints, '... j c -> ... (j c)')

def _get_z_orient_delta(data):
    """get global body orientation delta"""
    # default is axis-angle representation
    pelvis_orient = to_tensor(data['rots'][..., :3])
    pelvis_orient_z = get_z_rot(pelvis_orient, in_format="aa")
    pelvis_orient_z = transform_body_pose(pelvis_orient_z, "rot->aa")
    z_orient_delta = rot_diff(pelvis_orient_z, in_format="aa",
                                    out_format='6d')
    return z_orient_delta

FEAT_GET_METHODS = {
            "body_transl_delta_pelv": _get_body_transl_delta_pelv,
            "body_orient_xy": _get_body_orient_xy,
            "z_orient_delta": _get_z_orient_delta,
            "body_pose": _get_body_pose,
            "body_joints_local_wo_z_rot": _get_body_joints_local_wo_z_rot,
        }