# urdf_helpers.py
# Author: Nosa Edoimioya
# Description: General 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 xml.etree.ElementTree as ET
import numpy as np
from typing import Tuple, Dict


# Helper functions for parsing the URDF
def get_urdf_root(urdf_file: str) -> ET.Element:
    """Parse a URDF file and return the root element.

    Args:
        urdf_file: Absolute path to the URDF file.

    Returns:
        `xml.etree.ElementTree.Element` root node of the URDF.

    Side Effects:
        Reads the URDF file from disk.

    Raises:
        ET.ParseError: If the URDF cannot be parsed.
        OSError: If the file cannot be read.

    Preconditions:
        `urdf_file` exists and is readable.
    """
    try:
        tree = ET.parse(urdf_file)
    except ET.ParseError:
        print("ERROR: Could not parse urdf file.")

    return tree.getroot()


def process_joint(joint: ET.Element) -> Tuple[str, Dict]:
    """Extract joint attributes into a dictionary.

    Args:
        joint: Joint XML element from the URDF.

    Returns:
        `tuple[str, dict]` containing the joint name and joint data.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        `joint` is a valid URDF joint element.
    """
    axis = np.array([1, 0, 0])
    xyz = np.zeros(3)
    rpy = np.zeros(3)
    parent_link = ""
    child_link = ""
    joint_type = joint.get("type", "").lower()

    joint_name = joint.get("name") or ""

    for child in joint:
        if child.tag == "axis":
            axis = np.array((child.get("xyz") or "1 0 0").split(), dtype=float)
        elif child.tag == "origin":
            xyz = np.array((child.get("xyz") or "0 0 0").split(), dtype=float)
            rpy = np.array((child.get("rpy") or "0 0 0").split(), dtype=float)
        elif child.tag == "parent":
            parent_link = child.get("link") or ""
        elif child.tag == "child":
            child_link = child.get("link") or ""
    return joint_name, {
        "axis": axis,
        "xyz": xyz,
        "rpy": rpy,
        "parent": parent_link,
        "child": child_link,
        "type": joint_type,
        "dh": np.zeros(4),
    }


def process_link(link: ET.Element) -> Tuple[str, Dict]:
    """Extract link inertial properties into a dictionary.

    Args:
        link: Link XML element from the URDF.

    Returns:
        `tuple[str, dict]` containing the link name and properties.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        `link` is a valid URDF link element.
    """
    mass = 0.0
    center_of_mass = np.zeros(3)
    inertia_tensor = np.zeros((3, 3))

    link_name = link.get("name") or ""

    for child in link:
        if child.tag == "inertial":
            for grand_child in child:
                if grand_child.tag == "mass":
                    mass = float(grand_child.get("value") or 0.0)
                elif grand_child.tag == "inertia":
                    inertia_tensor = np.array(
                        [
                            [
                                grand_child.get("ixx"),
                                grand_child.get("ixy"),
                                grand_child.get("ixz"),
                            ],
                            [
                                grand_child.get("ixy"),
                                grand_child.get("iyy"),
                                grand_child.get("iyz"),
                            ],
                            [
                                grand_child.get("ixz"),
                                grand_child.get("iyz"),
                                grand_child.get("izz"),
                            ],
                        ],
                        dtype=float,
                    )
                elif grand_child.tag == "origin":
                    center_of_mass = np.array(
                        (grand_child.get("xyz") or "0 0 0").split(), dtype=float
                    )

            # Break loop after finding the 'inertial' tag
            break

    return link_name, {
        "mass": mass,
        "center_of_mass": center_of_mass,
        "inertia_tensor": inertia_tensor,
    }
