'''
calculate the velocity of the imu.

The measured acc is in imu body frame.

1. acc offset along each axis calculation.
    - gravity mapped to imu body frame, according to rpy;
    - subtract gravity items to get the pure acc along each axis;
    - keep imu static, average the pure acc values, as the offset;
    - calibrate the offset in following use.

2. vel calculation.
    - get calibrated acc in imu body frame;
    - convert this calibrate acc to world frame;
    - integtate this acc, get the vel in world frame;
    - convert the vel back to imu body frame.
'''


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],
    ]


class IMUVelocityEstimator:
    def __init__(self):
        # Velocity is stored in world frame, unit m/s
        self.vel_world = [0.0, 0.0, 0.0]

        # Acceleration offset in body frame, unit g
        self.acc_offset_body = [0.0, 0.0, 0.0]

    def reset_velocity(self):
        self.vel_world = [0.0, 0.0, 0.0]

    def set_acc_offset_body(self, acc_offset_body):
        self.acc_offset_body = list(acc_offset_body)

    def update(self, acc_body_g, rpy_deg, dt):
        """
        Parameters
        ----------
        acc_body_g:
            measured acceleration in IMU/body frame, unit g

        rpy_deg:
            [roll, pitch, yaw], unit degree

        dt:
            time step, unit second

        Returns
        -------
        acc_world_linear:
            gravity-compensated acceleration in world frame, unit m/s^2

        vel_world:
            velocity in world frame, unit m/s

        vel_body:
            velocity projected onto current IMU/body frame, unit m/s
        """

        R_bw = rotation_matrix_body_to_world(rpy_deg)
        R_wb = mat_transpose(R_bw)

        # 1. Remove body-frame acceleration offset
        acc_body_corrected_g = [
            acc_body_g[0] - self.acc_offset_body[0],
            acc_body_g[1] - self.acc_offset_body[1],
            acc_body_g[2] - self.acc_offset_body[2],
        ]

        # 2. Convert body acceleration from g to m/s^2
        acc_body_mps2 = [
            acc_body_corrected_g[0] * G,
            acc_body_corrected_g[1] * G,
            acc_body_corrected_g[2] * G,
        ]

        # 3. Convert acceleration body -> world
        acc_world_mps2 = mat_vec_mul(R_bw, acc_body_mps2)

        # 4. Subtract gravity in world frame
        # World Z is upward, and static flat IMU gives +1g on Z.
        acc_world_linear = [
            acc_world_mps2[0],
            acc_world_mps2[1],
            acc_world_mps2[2] - G,
        ]

        # 5. Integrate velocity in world frame
        self.vel_world[0] += acc_world_linear[0] * dt
        self.vel_world[1] += acc_world_linear[1] * dt
        self.vel_world[2] += acc_world_linear[2] * dt

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

        return acc_world_linear, list(self.vel_world), vel_body

    def calibrate_acc_offset(self, samples):
        """
        Calibrate acceleration offset.

        Parameters
        ----------
        samples:
            list of samples.
            Each sample should be:
                (acc_body_g, rpy_deg)

            acc_body_g: [Ax, Ay, Az], unit g
            rpy_deg:    [roll, pitch, yaw], degree

        Principle
        ---------
        During static calibration:

            measured_acc_body = gravity_body + offset_body

        So:

            offset_body = measured_acc_body - gravity_body

        We average offset_body over all samples.

        Returns
        -------
        acc_offset_body:
            acceleration offset in body frame, unit g
        """

        if len(samples) == 0:
            raise ValueError("samples is empty.")

        offset_sum = [0.0, 0.0, 0.0]

        for acc_body_g, rpy_deg in samples:
            R_bw = rotation_matrix_body_to_world(rpy_deg)
            R_wb = mat_transpose(R_bw)

            # Gravity in world frame.
            # Unit g.
            gravity_world_g = [0.0, 0.0, 1.0]

            # Convert gravity world -> body
            gravity_body_g = mat_vec_mul(R_wb, gravity_world_g)

            # Offset = measured acceleration - expected gravity
            offset_body_g = [
                acc_body_g[0] - gravity_body_g[0],
                acc_body_g[1] - gravity_body_g[1],
                acc_body_g[2] - gravity_body_g[2],
            ]

            offset_sum[0] += offset_body_g[0]
            offset_sum[1] += offset_body_g[1]
            offset_sum[2] += offset_body_g[2]

        n = len(samples)

        self.acc_offset_body = [
            offset_sum[0] / n,
            offset_sum[1] / n,
            offset_sum[2] / n,
        ]

        # Calibration means current velocity should be zero
        self.reset_velocity()

        return list(self.acc_offset_body)

if __name__ == "__main__":
    estimator = IMUVelocityEstimator()

    # Example calibration samples.
    # In real use, collect these while the IMU is static.
    calib_samples = [
        ([0.2, 0.73, 0.66], [46, -10, -8]),
        ([0.201, 0.731, 0.661], [46, -10, -8]),
        ([0.199, 0.729, 0.659], [46, -10, -8]),
    ]

    offset = estimator.calibrate_acc_offset(calib_samples)
    print("acc_offset_body =", offset)

    dt = 0.01

    acc = [0.22, 0.75, 0.67]
    rpy = [46, -10, -8]

    acc_world_linear, vel_world, vel_body = estimator.update(acc, rpy, dt)

    print("acc_world_linear =", acc_world_linear)
    print("vel_world        =", vel_world)
    print("vel_body         =", vel_body)