# kinematic_helpers.py
# Author: Nosa Edoimioya
# Description: Kinematic helper functions to generate DH parameters based on robot urdf file.
# Version: 0.1
# Date: 03-27-2024
# Based on https://github.com/mcevoyandy/urdf_to_dh

import numpy as np
import math


# Kinematics helper functions
def x_rotation(theta: float) -> np.ndarray:
    """Return rotation matrix for a rotation about the x-axis.

    Args:
        theta: Rotation angle in radians.

    Returns:
        `np.ndarray` 3x3 rotation matrix.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        None.
    """
    return np.array(
        [
            [1, 0, 0],
            [0, math.cos(theta), -math.sin(theta)],
            [0, math.sin(theta), math.cos(theta)],
        ]
    )


def y_rotation(theta: float) -> np.ndarray:
    """Return rotation matrix for a rotation about the y-axis.

    Args:
        theta: Rotation angle in radians.

    Returns:
        `np.ndarray` 3x3 rotation matrix.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        None.
    """
    return np.array(
        [
            [math.cos(theta), 0, math.sin(theta)],
            [0, 1, 0],
            [-math.sin(theta), 0, math.cos(theta)],
        ]
    )


def z_rotation(theta: float) -> np.ndarray:
    """Return rotation matrix for a rotation about the z-axis.

    Args:
        theta: Rotation angle in radians.

    Returns:
        `np.ndarray` 3x3 rotation matrix.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        None.
    """
    return np.array(
        [
            [math.cos(theta), -math.sin(theta), 0],
            [math.sin(theta), math.cos(theta), 0],
            [0, 0, 1],
        ]
    )


def get_extrinsic_rotation(rpy: np.ndarray) -> np.ndarray:
    """Return extrinsic rotation matrix from roll, pitch, yaw.

    Args:
        rpy: Roll, pitch, yaw angles in radians.

    Returns:
        `np.ndarray` 3x3 rotation matrix.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        `rpy` has length 3.
    """

    x_rot = x_rotation(rpy[0])
    y_rot = y_rotation(rpy[1])
    z_rot = z_rotation(rpy[2])
    return np.matmul(z_rot, np.matmul(y_rot, x_rot))


def inv_tf(tf: np.ndarray) -> np.ndarray:
    """Return the inverse of a homogeneous transform.

    Args:
        tf: 4x4 homogeneous transform.

    Returns:
        `np.ndarray` 4x4 inverse transform.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        `tf` is a valid 4x4 transform matrix.
    """
    inv_tf = np.eye(4)
    inv_tf[0:3, 0:3] = np.transpose(tf[0:3, 0:3])
    inv_tf[0:3, 3] = -1.0 * np.matmul(np.transpose(tf[0:3, 0:3]), tf[0:3, 3])
    return inv_tf


def get_dh_frame(dh_params: np.ndarray) -> np.ndarray:
    """Return the transform for the given DH parameters.

    Args:
        dh_params: DH parameters `[d, theta, r, alpha]`.

    Returns:
        `np.ndarray` 4x4 transform matrix.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        `dh_params` has length 4.
    """
    d = dh_params[0]
    theta = dh_params[1]
    r = dh_params[2]
    alpha = dh_params[3]

    dh_frame = np.eye(4)

    dh_frame[0, 0] = math.cos(theta)
    dh_frame[0, 1] = -math.sin(theta) * math.cos(alpha)
    dh_frame[0, 2] = math.sin(theta) * math.sin(alpha)
    dh_frame[0, 3] = r * math.cos(theta)

    dh_frame[1, 0] = math.sin(theta)
    dh_frame[1, 1] = math.cos(theta) * math.cos(alpha)
    dh_frame[1, 2] = -math.cos(theta) * math.sin(alpha)
    dh_frame[1, 3] = r * math.sin(theta)

    dh_frame[2, 1] = math.sin(alpha)
    dh_frame[2, 2] = math.cos(alpha)
    dh_frame[2, 3] = d
    return dh_frame
