import math


G = 9.80665


def mat_vec_mul(R, v):
    return [
        R[0][0] * v[0] + R[0][1] * v[1] + R[0][2] * v[2],
        R[1][0] * v[0] + R[1][1] * v[1] + R[1][2] * v[2],
        R[2][0] * v[0] + R[2][1] * v[1] + R[2][2] * v[2],
    ]


def mat_transpose(R):
    return [
        [R[0][0], R[1][0], R[2][0]],
        [R[0][1], R[1][1], R[2][1]],
        [R[0][2], R[1][2], R[2][2]],
    ]


def rotation_matrix_body_to_world(rpy_deg):
    """
    rpy_deg = [roll, pitch, yaw], degree

    Convention:
        roll  around X
        pitch around Y
        yaw   around Z

    R_body_to_world = Rz(yaw) @ Ry(pitch) @ Rx(roll)
    """

    roll = math.radians(rpy_deg[0])
    pitch = math.radians(rpy_deg[1])
    yaw = math.radians(rpy_deg[2])

    cr, sr = math.cos(roll), math.sin(roll)
    cp, sp = math.cos(pitch), math.sin(pitch)
    cy, sy = math.cos(yaw), math.sin(yaw)

    return [
        [cy * cp, cy * sp * sr - sy * cr, cy * sp * cr + sy * sr],
        [sy * cp, sy * sp * sr + cy * cr, sy * sp * cr - cy * sr],
        [-sp,     cp * sr,                  cp * cr],
    ]

#####################################################################

dt = 0.02

rpy_deg = [10, 20, 30]

linear_acc_ms2 = [3.2, 5.5 , 7.7]           #   linear acc in imu body frame

vel_world = [0.0, 0.0, 0.0]                 #   linear vel in world frame


#   get direction cosine matrix
R_bw = rotation_matrix_body_to_world(rpy_deg)
R_wb = mat_transpose(R_bw)


#   Convert acceleration body -> world
acc_world_mps2 = mat_vec_mul(R_bw, linear_acc_ms2)

#   Integrate velocity in world frame
vel_world[0] += acc_world_mps2[0] * dt
vel_world[1] += acc_world_mps2[1] * dt
vel_world[2] += acc_world_mps2[2] * dt

#   Convert velocity world -> current body
vel_body = mat_vec_mul(R_wb, vel_world)