import numpy as np
import torch
from SMPLX.rotation_conversions import rotation_6d_to_matrix, matrix_to_axis_angle

def qinv(q):
    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
    mask = torch.ones_like(q)
    mask[..., 1:] = -mask[..., 1:]
    return q * mask

def qrot(q, v):
    """
    Rotate vector(s) v about the rotation described by quaternion(s) q.
    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
    where * denotes any number of dimensions.
    Returns a tensor of shape (*, 3).
    """
    assert q.shape[-1] == 4
    assert v.shape[-1] == 3
    assert q.shape[:-1] == v.shape[:-1]

    original_shape = list(v.shape)
    # print(q.shape)
    q = q.contiguous().view(-1, 4)
    v = v.contiguous().view(-1, 3)

    qvec = q[:, 1:]
    uv = torch.cross(qvec, v, dim=1)
    uuv = torch.cross(qvec, uv, dim=1)
    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)


def recover_root_rot_pos(data):
    rot_vel = data[..., 0]
    r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
    '''Get Y-axis rotation from rotation velocity'''
    r_rot_ang[..., 1:] = rot_vel[..., :-1]
    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)

    r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)
    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
    r_rot_quat[..., 2] = torch.sin(r_rot_ang)

    r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)
    r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
    '''Add Y-axis rotation to root position'''
    r_pos = qrot(qinv(r_rot_quat), r_pos)

    r_pos = torch.cumsum(r_pos, dim=-2)

    r_pos[..., 1] = data[..., 3]
    return r_rot_quat, r_pos

def recover_from_ric(data, joints_num):
    if isinstance(data, np.ndarray):
        data = torch.from_numpy(data).float()
        dtype = "numpy"
    else:
        data = data.float()
        dtype = "tensor"
    r_rot_quat, r_pos = recover_root_rot_pos(data)
    positions = data[..., 4:(joints_num - 1) * 3 + 4]
    positions = positions.view(positions.shape[:-1] + (-1, 3))

    '''Add Y-axis rotation to local joints'''
    positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)

    '''Add root XZ to joints'''
    positions[..., 0] += r_pos[..., 0:1]
    positions[..., 2] += r_pos[..., 2:3]

    '''Concate root and joints'''
    positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)

    if dtype == "numpy":
        positions = positions.numpy()

    return positions

def recover_pose_from_smr(data, njoints=22):
    joints = recover_from_ric(data, njoints)    
    trans = joints[:, 0, :] - joints[0:1, 0, :]

    pose = data[:, 4 + (njoints - 1) * 3:10 + (njoints - 1) * 9]
    pose = pose.reshape(pose.shape[0], njoints, 6)
    ptype = type(pose)
    if ptype == np.ndarray:
        pose = torch.from_numpy(pose).float()
        pose = rotation_6d_to_matrix(pose)
        pose = matrix_to_axis_angle(pose)
        pose = pose.numpy()
    elif ptype == torch.Tensor:
        pose = rotation_6d_to_matrix(pose)
        pose = matrix_to_axis_angle(pose)
    
    pose = pose.reshape(pose.shape[0], -1)

    if njoints < 24:
        if ptype == np.ndarray:
            addition = np.zeros([pose.shape[0], 72-njoints*3])
            pose = np.concatenate([pose, addition], axis=1)
        elif ptype == torch.Tensor:
            addition = torch.zeros([pose.shape[0], 72-njoints*3], dtype=pose.dtype, device=pose.device)
            pose = torch.cat([pose, addition], dim=1)

    if ptype == np.ndarray:
        pose = np.concatenate([pose, trans], axis=1)
    elif ptype == torch.Tensor:
        pose = torch.cat([pose, trans], dim=1)

    return pose