#!/usr/bin/env python3
"""
Converts inertial data from the documented manner to the format
needed for an SCL spec file for the Barret WAM 7DOF arm. All
units are SI.

NOTE: This does not take into account motor inertias, which
      according to the documentation are approximately
      equal in magnitude to the link inertias.

NOTE: The hand is also not taken into account here, although the
      physical mass of the wrist motors in link 4 are included.
"""

import mathutils as mu
from math import radians

# -------------------------------------------------------
# Data from Barrett WAM Documentation and CAD Model
# -------------------------------------------------------

# Translations between the link input and output frames.
# For example, entry '2' corresponds to 'link2' and is
# thre translation from frame 1 to frame 2.
translations = {
    'B': (0.220000, 0.140000, 0.346000),  # Frame base to 0
    '1': (0.000000, 0.000000, 0.000000),  # Frame 0 to 1
    '2': (0.000000, 0.000000, 0.000000),
    '3': (0.04624199867248535, 0.0, 0.55024),
    '4': (-0.04624199867248535, 0.0, 0.0),
    '5': (0.0, 0.0, 0.3),
    '6': (0.0, 0.0, 0.0),
    '7': (0.0, 0.0, 0.0)
}

# Rotations between the link input and output frames. All links
# rotate about their input frame's Z coordinate.
rotations = {
    'B': (0.000, 0.000, 0.000),  # Frame base to 0
    '1': (-90.000, 0.000, 0.000),  # Frame 0 to 1
    '2': (+90.000, 0.000, 0.000),
    '3': (-90.000, 0.000, 0.000),
    '4': (+90.000, 0.000, 0.000),
    '5': (-90.000, 0.000, 0.000),
    '6': (+90.000, 0.000, 0.000),
    '7': (0, 0, 0)  # Frame 6 to 7
}

# Link masses
masses = {
    'B': 9.97059584,
    '1': 10.76768767,
    '2': 3.87493756,
    '3': 1.80228141,
    '4': 2.40016804,
    '5': 0.12376019,
    '6': 0.41797364,
    '7': 0.06864753
}

# Link center of masses from the OUTPUT frame
# For SCL, we need them from the INPUT frame
center_of_masses = {
    'B': (-0.02017671, -0.26604706, -0.14071720),
    '1': (-0.00443422, 0.12189039, -0.00066489),
    '2': (-0.00236983, 0.03105614, 0.01542114),
    '3': (-0.03825858, 0.20750770, 0.00003309),
    '4': (0.00498512, -0.00022942, 0.13271662),
    '5': (0.00008921, 0.00511217, 0.00435824),
    '6': (-0.00012262, -0.01703194, 0.02468336),
    '7': (-0.00007974, 0.00016313, -0.00323552)
}

# Link inertias about the center of mass, aligned with the OUTPUT
# coordinate frame. For SCL, we need to rotate them into the INPUT
# coordinate frame.
# Format: Ixx, Iyy, Izz, Ixy, Ixz, Iyz
inertias = {
    'B': (0.10916849, 0.18294303, 0.11760385, 0.00640270, 0.02557874, 0.00161433),
    '1': (0.13488033, 0.11328369, 0.09046330, -0.00213041, -0.00012485, 0.00068555),
    '2': (0.02140958, 0.01377875, 0.01558906, 0.00027172, 0.00002461, -0.00181920),
    '3': (0.05911077, 0.00324550, 0.05927043, -0.00249612, 0.00000738, -0.00001767),
    '4': (0.01491672, 0.01482922, 0.00294463, 0.00001741, -0.00150604, -0.00002109),
    '5': (0.00005029, 0.00007582, 0.00006270, 0.00000020, -0.00000005, -0.00000359),
    '6': (0.00055516, 0.00024367, 0.00045358, 0.00000061, -0.00000074, -0.00004590),
    '7': (0.00003773, 0.00003806, 0.00007408, -0.00000019, 0.00000000, 0.00000000)
}


def calculate_output(link):

    # Translation from input frame to output frame
    R_io = mu.Euler((radians(v) for v in rotations[link])).to_matrix()

    # Translation from output frame to input frame
    R_oi = R_io.transposed()

    # Indices to build inertia matrix
    I_inx = ((0, 3, 4), (3, 1, 5), (4, 5, 2))

    # Inertia matrix in the output frame
    I_o = mu.Matrix([[inertias[link][inx] for inx in row] for row in I_inx])

    # Inertia matrix in the input frame
    I_i = R_io * I_o * R_oi

    t = mu.Vector(translations[link])
    com_o = mu.Vector(center_of_masses[link])
    com_i = R_io * com_o + t

    q_oi = R_io.to_quaternion()
    q_oi_xyzw = (q_oi.x, q_oi.y, q_oi.z, q_oi.w)

    print('-------------------------------')
    print('Link {}'.format(link))
    print('-------------------------------')
    # Not printing translation and rotation, because for the SCL file they
    # actually correspond to the next link's entries, instead of the listed
    # link.
    # print('translation: {}'.format(' '.join('{:e}'.format(v) for v in t)))
    # print('rotation (quat): {}'.format(' '.join('{:e}'.format(v) for v in q_oi_xyzw)))
    print('mass: {:04f}'.format(masses[link]))
    print('center_of_mass: {}'.format(' '.join('{:e}'.format(v) for v in com_i)))
    print('inertia at com: {}'.format(format_inertia(I_i)))


def format_inertia(I):
    return ' '.join('{:e}'.format(v) for v in
                    [I[0][0], I[1][1], I[2][2], I[0][1], I[0][2], I[1][2]])


if __name__ == '__main__':

    for link in ['B', '1', '2', '3', '4', '5', '6', '7']:
        calculate_output(link)