import math
from typing import List, Tuple


G = 9.80665  # m/s^2


def mat_vec_mul(R: List[List[float]], v: List[float]) -> List[float]:
    """3x3 matrix times 3x1 vector."""
    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 transpose(R: List[List[float]]) -> List[List[float]]:
    """Transpose of 3x3 matrix."""
    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: List[float]) -> List[List[float]]:
    """
    Convert roll, pitch, yaw to rotation matrix.

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

    Convention:
        roll  : rotation around X axis
        pitch : rotation around Y axis
        yaw   : rotation around Z axis

    Rotation order:
        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 = math.cos(roll)
    sr = math.sin(roll)

    cp = math.cos(pitch)
    sp = math.sin(pitch)

    cy = math.cos(yaw)
    sy = math.sin(yaw)

    # R = Rz(yaw) @ Ry(pitch) @ Rx(roll)
    R = [
        [
            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,
        ],
    ]

    return R


def gravity_body_from_rpy(rpy_deg: List[float]) -> List[float]:
    """
    Calculate gravity vector projected onto body frame.

    Unit: g

    This matches your original formula:
        g_x = -sin(pitch)
        g_y = cos(pitch) * sin(roll)
        g_z = cos(pitch) * cos(roll)

    Yaw does not affect gravity projection.
    """

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

    g_x = -math.sin(pitch)
    g_y = math.cos(pitch) * math.sin(roll)
    g_z = math.cos(pitch) * math.cos(roll)

    return [g_x, g_y, g_z]


def remove_gravity_in_body(acc_body_g: List[float], rpy_deg: List[float]) -> List[float]:
    """
    Remove gravity directly in body frame.

    This is equivalent to your original acc_calib().
    Unit input : g
    Unit output: g
    """

    g_body = gravity_body_from_rpy(rpy_deg)

    return [
        acc_body_g[0] - g_body[0],
        acc_body_g[1] - g_body[1],
        acc_body_g[2] - g_body[2],
    ]


class IMUVelocityEstimator:
    """
    Estimate velocity from IMU acceleration.

    Internal velocity state is stored in world frame:
        self.v_world = [vx, vy, vz], unit m/s

    Each update returns:
        a_world_linear_mps2 : gravity-compensated linear acceleration in world frame
        v_world             : velocity in world frame
        v_body              : velocity projected onto current IMU/body XYZ axes
    """

    def __init__(self):
        self.v_world = [0.0, 0.0, 0.0]

    def reset_velocity(self, v_world: List[float] = None):
        if v_world is None:
            self.v_world = [0.0, 0.0, 0.0]
        else:
            self.v_world = list(v_world)

    def update(
        self,
        acc_body_g: List[float],
        rpy_deg: List[float],
        dt: float,
    ) -> Tuple[List[float], List[float], List[float]]:
        """
        One IMU update step.

        Parameters
        ----------
        acc_body_g:
            IMU measured acceleration in current body frame.
            Unit: g
            Example: [Ax, Ay, Az]

        rpy_deg:
            IMU attitude angle.
            Unit: degree
            Format: [roll, pitch, yaw]

        dt:
            Sampling interval.
            Unit: second

        Returns
        -------
        a_world_linear_mps2:
            Linear acceleration in world frame after gravity removal.
            Unit: m/s^2

        v_world:
            Integrated velocity in world frame.
            Unit: m/s

        v_body:
            Current velocity projected onto current IMU/body XYZ axes.
            Unit: m/s
        """

        # 1. Rotation matrix: body frame -> world frame
        R_bw = rotation_matrix_body_to_world(rpy_deg)

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

        # 3. Transform acceleration from body frame to world frame
        acc_world_mps2 = mat_vec_mul(R_bw, acc_body_mps2)

        # 4. Remove gravity in world frame
        #
        # Since your IMU flat on table gives Az ≈ +1g,
        # and world Z is upward, gravity component in this convention is:
        #
        #     g_world = [0, 0, +G]
        #
        # Therefore linear acceleration is:
        #
        #     a_linear_world = acc_world - g_world
        #
        a_world_linear_mps2 = [
            acc_world_mps2[0],
            acc_world_mps2[1],
            acc_world_mps2[2] - G,
        ]

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

        # 6. Transform world velocity back to current body frame
        R_wb = transpose(R_bw)
        v_body = mat_vec_mul(R_wb, self.v_world)

        return a_world_linear_mps2, list(self.v_world), v_body


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

    acc = [0.2, 0.73, 0.66]      # unit g
    rpy = [46, -10, -8]          # [roll, pitch, yaw], unit degree
    dt = 0.01                    # 100 Hz sample rate

    a_world, v_world, v_body = estimator.update(acc, rpy, dt)

    print("a_world_linear_mps2 =", a_world)
    print("v_world_mps         =", v_world)
    print("v_body_mps          =", v_body)

    # For comparison with your original direct body-frame gravity compensation:
    print("a_body_linear_g     =", remove_gravity_in_body(acc, rpy))