import torch

TRANSFORMATION_AA_TO_SJA = torch.Tensor([
    [[1, 0, 0], [0, 0, 1], [0, -1, 0]],  # 00, 'left_hip',
    [[1, 0, 0], [0, 0, 1], [0, -1, 0]],  # 01, 'right_hip',
    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],  # 02, 'spine1',
    [[1, 0, 0], [0, 0, 1], [0, -1, 0]],  # 03, 'left_knee',
    [[1, 0, 0], [0, 0, 1], [0, -1, 0]],  # 04, 'right_knee',
    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],  # 05, 'spine2',
    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],  # 06, 'left_ankle',
    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],  # 07, 'right_ankle',
    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],  # 08, 'spine3',
    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],  # 09, 'left_foot',
    [[1, 0, 0], [0, 1, 0], [0, 0, 1]],  # 10, 'right_foot',
    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],  # 11, 'neck',
    [[0, 0, -1], [0, 1, 0], [1, 0, 0]],  # 12, 'left_collar',
    [[0, 0, 1], [0, 1, 0], [-1, 0, 0]],  # 13, 'right_collar',
    [[1, 0, 0], [0, 0, -1], [0, 1, 0]],  # 14, 'head',
    [[0, 0, -1], [0, 1, 0], [1, 0, 0]],  # 15, 'left_shoulder',
    [[0, 0, 1], [0, 1, 0], [-1, 0, 0]],  # 16, 'right_shoulder',
    [[0, 0, -1], [0, 1, 0], [1, 0, 0]],  # 17, 'left_elbow',
    [[0, 0, 1], [0, 1, 0], [-1, 0, 0]],  # 18, 'right_elbow',
    [[0, 0, -1], [0, 1, 0], [1, 0, 0]],  # 19, 'left_wrist',
    [[0, 0, 1], [0, 1, 0], [-1, 0, 0]],  # 20, 'right_wrist',
])

TRANSFORMATION_SJA_TO_AA = \
    torch.inverse(TRANSFORMATION_AA_TO_SJA)

# TODO: spines and shoulders may need further adjustment
STANDARD_JOINT_ANGLE_LIMITS = torch.deg2rad(
    torch.Tensor([
        [[-45, 155], [-88, 17], [-105, 85]],  # 00, 'left_hip',
        [[-45, 155], [-17, 88], [-85, 105]],  # 01, 'right_hip',
        [[-25, 15], [-20, 20], [-30, 30]],  # 02, 'spine1',
        [[0, 150], [0, 0], [0, 0]],  # 03, 'left_knee',
        [[0, 150], [0, 0], [0, 0]],  # 04, 'right_knee',
        [[-25, 15], [-15, 15], [-25, 25]],  # 05, 'spine2',
        [[-31, 63], [-26, 26], [-74, 15]],  # 06, 'left_ankle',
        [[-31, 63], [-26, 26], [-15, 74]],  # 07, 'right_ankle',
        [[-25, 15], [-15, 15], [-25, 25]],  # 08, 'spine3',
        [[-60, 45], [0, 0], [-45, 45]],  # 09, 'left_foot',
        [[-60, 45], [0, 0], [-45, 45]],  # 10, 'right_foot',
        [[-37, 22], [-30, 30], [-45, 45]],  # 11, 'neck',
        [[-30, 30], [-30, 10], [0, 0]],  # 12, 'left_collar',
        [[-30, 30], [-10, 30], [0, 0]],  # 13, 'right_collar',
        [[-37, 22], [-30, 30], [-45, 45]],  # 14, 'head',
        [[-90, 135], [-97, 91], [-90, 135]],  # 15, 'left_shoulder',
        [[-135, 90], [-91, 97], [-135, 90]],  # 16, 'right_shoulder',
        [[0, 0], [-150, 0], [0, 0]],  # 17, 'left_elbow',
        [[0, 0], [0, 150], [0, 0]],  # 18, 'right_elbow',
        [[-90, 90], [-45, 45], [-180, 60]],  # 19, 'left_wrist',
        [[-90, 90], [-45, 45], [-60, 180]],  # 20, 'right_wrist',
    ]))
