# src/util/urdf_to_dh/generate_dh.py
# This module generates Denavit-Hartenberg parameters from a URDF file.

from anytree import AnyNode, LevelOrderIter
import numpy as np
import math

import reforge_core.util.urdf_to_dh.kinematics_helpers as kh
import reforge_core.util.urdf_to_dh.geometry_helpers as gh
import reforge_core.util.urdf_to_dh.urdf_helpers as uh
from typing import Any


class GenerateDhParams:
    """Generate Denavit–Hartenberg parameters from a URDF file.

    Args:
        None.

    Side Effects:
        None.

    Raises:
        None.

    Preconditions:
        None.
    """

    def __init__(self, urdf_file: str) -> None:
        """Initialize URDF-to-DH conversion state.

        Args:
            urdf_file: Path to the URDF file.
        Returns:
            `None`.
        Side Effects:
            Initializes internal dictionaries and transform placeholders.
        Raises:
            None.
        Preconditions:
            `urdf_file` points to a valid URDF that will be parsed later.
        """
        print("Initializing DH parameter generator...")
        # All joints in the URDF (including fixed)
        self.urdf_joints_all: dict[str, dict[str, Any]] = {}
        # Only moving joints (revolute/prismatic) used for DH and dynamics
        self.urdf_joints: dict[str, dict[str, Any]] = {}
        self.moving_joints: list[str] = []
        self.urdf_links: dict[str, dict[str, Any]] = {}
        self.urdf_file = urdf_file
        self.urdf_tree_nodes: list[AnyNode] = []
        self.root_link = None
        self.dh_dict: dict[str, dict[int, Any]] = {}
        self.inertia_dict: dict[str, dict[int, Any]] = {}
        self.verbose = False
        self.base_prefix_tf = np.eye(4)
        self.tool_offset_tf = np.eye(4)

        print("URDF file = %s" % self.urdf_file)

    def parse_urdf(self) -> None:
        """Parse the URDF and build the link/joint tree representation.

        Returns:
            `None`.

        Side Effects:
            Populates `urdf_links`, `urdf_joints_all`, `urdf_tree_nodes`, and
            `root_link`.

        Raises:
            RuntimeError: If the URDF has multiple root links.

        Preconditions:
            `self.urdf_file` points to a readable URDF file.
        """
        # Get the root of the URDF and extract all of the joints
        urdf_root = uh.get_urdf_root(self.urdf_file)

        # Parse all links first and add to tree
        for child in urdf_root:
            if child.tag == "link":
                _, link_data = uh.process_link(child)
                link_name = child.get("name") or ""
                self.urdf_links[link_name] = {
                    "rel_tf": np.eye(4),
                    "abs_tf": np.eye(4),
                    "dh_tf": np.eye(4),
                    "abs_dh_tf": np.eye(4),
                    "dh_found": False,
                    "mass": link_data["mass"],
                    "center_of_mass": link_data["center_of_mass"],
                    "inertia_tensor": link_data["inertia_tensor"],
                }
                node = AnyNode(id=link_name, parent=None, children=None, type="link")
                self.urdf_tree_nodes.append(node)

        # Parse all joints and add to tree
        for child in urdf_root:
            if child.tag == "joint":
                joint_name, joint_data = uh.process_joint(child)
                self.urdf_joints_all[joint_name] = joint_data
                node = AnyNode(id=joint_name, parent=None, children=None, type="joint")

                # Find parent and child link
                for n in self.urdf_tree_nodes:
                    if n.id == joint_data["parent"]:
                        node.parent = n
                    if n.id == joint_data["child"]:
                        n.parent = node
                self.urdf_tree_nodes.append(node)

        # Find root link
        num_nodes_no_parent = 0
        for n in self.urdf_tree_nodes:
            if n.parent is None:
                num_nodes_no_parent += 1
                self.root_link = n

        if num_nodes_no_parent == 1:
            # Root link DH will be identity, set dh_found = True
            # TO DO: Probably not needed since order iter is used
            if self.root_link is None:
                raise RuntimeError("Root link could not be determined from URDF tree")
            self.urdf_links[self.root_link.id]["dh_found"] = True
            # print("URDF Tree:")
            # for pre, _, node in RenderTree(self.root_link):
            #     print('%s%s' % (pre, node.id))

            # print("Joint Info:")
            # pprint.pprint(self.urdf_joints)
        else:
            raise RuntimeError("Error: Should only be one root link")

    def calculate_tfs_in_world_frame(self) -> None:
        """Compute absolute transforms for each link in the world frame.

        Returns:
            `None`.

        Side Effects:
            Populates each link's `rel_tf` and `abs_tf` in `urdf_links`.

        Raises:
            None.

        Preconditions:
            `parse_urdf()` has been called and `root_link` is set.
        """
        # print("Calculate world tfs:")
        if self.root_link is None:
            raise RuntimeError("Root link is not set. Call parse_urdf() first.")
        for n in LevelOrderIter(self.root_link):
            if n.type == "link" and n.parent is not None:
                # print("\nget tf from ", n.parent.parent.id, " to ", n.id)

                parent_tf_world = self.urdf_links[n.parent.parent.id]["abs_tf"]
                xyz = self.urdf_joints_all[n.parent.id]["xyz"]
                rpy = self.urdf_joints_all[n.parent.id]["rpy"]
                tf = np.eye(4)
                tf[0:3, 0:3] = kh.get_extrinsic_rotation(rpy)
                tf[0:3, 3] = xyz
                self.urdf_links[n.id]["rel_tf"] = tf

                abs_tf = np.eye(4)
                abs_tf = np.matmul(parent_tf_world, tf)
                self.urdf_links[n.id]["abs_tf"] = abs_tf

    def calculate_params(self) -> None:
        """Compute Denavit–Hartenberg and inertia parameters from the URDF tree.

        Returns:
            `None`.

        Side Effects:
            Populates `dh_dict`, `inertia_dict`, and moving joint metadata.

        Raises:
            None.

        Preconditions:
            `parse_urdf()` and `calculate_tfs_in_world_frame()` have been called.
        """
        robot_dh_params = []
        robot_mass_params = []

        moving_joint_names = []

        for urdf_node in LevelOrderIter(self.root_link):
            if (
                urdf_node.type == "link"
                and self.urdf_links[urdf_node.id]["dh_found"] is False
            ):
                # TF from current link frame to world frame
                link_to_world = self.urdf_links[urdf_node.id]["abs_tf"]

                # DH frame from parent link frame to world frame
                parent_to_world_dh = self.urdf_links[urdf_node.parent.parent.id][
                    "abs_dh_tf"
                ]

                # TF from link frame to parent dh frame
                link_to_parent_dh = np.matmul(
                    kh.inv_tf(parent_to_world_dh), link_to_world
                )

                # Skip fixed joints but propagate their transform downstream
                joint_name = urdf_node.parent.id
                joint_data = self.urdf_joints_all[joint_name]
                if joint_data.get("type", "").lower() == "fixed":
                    # Fold the fixed joint transform into the child's DH frame
                    self.urdf_links[urdf_node.id]["dh_tf"] = link_to_parent_dh
                    self.urdf_links[urdf_node.id]["abs_dh_tf"] = np.matmul(
                        parent_to_world_dh, link_to_parent_dh
                    )
                    self.urdf_links[urdf_node.id]["dh_found"] = True
                    continue

                # Find DH parameters
                axis = np.matmul(link_to_parent_dh[0:3, 0:3], joint_data["axis"])
                dh_params = self.__get_joint_dh_params(link_to_parent_dh, axis)

                # Preserve your original rounding in radians (before deg conversion below)
                dh_params = np.round(dh_params.astype(float), 5)

                dh_frame = kh.get_dh_frame(dh_params)
                abs_dh_frame = np.matmul(parent_to_world_dh, dh_frame)

                self.urdf_links[urdf_node.id]["dh_tf"] = dh_frame
                self.urdf_links[urdf_node.id]["abs_dh_tf"] = abs_dh_frame

                # Order: [joint, parent, child, d, theta(rad, rounded), r, alpha(rad, rounded)]
                robot_dh_params.append(
                    [
                        joint_name,
                        urdf_node.parent.parent.id,
                        urdf_node.id,
                        float(dh_params[0]),
                        float(dh_params[1]),
                        float(dh_params[2]),
                        float(dh_params[3]),
                    ]
                )
                moving_joint_names.append(joint_name)

                robot_mass_params.append(
                    [
                        urdf_node.id,
                        self.urdf_links[urdf_node.id]["mass"],
                        self.urdf_links[urdf_node.id]["center_of_mass"],
                        self.urdf_links[urdf_node.id]["inertia_tensor"],
                    ]
                )

        # Build dictionaries with the same structure as pandas.DataFrame(...).to_dict()
        # i.e., {column_name: {row_index: value, ...}, ...}
        dh_columns = ["joint", "parent", "child", "d", "theta", "r", "alpha"]
        dh_dict: dict[str, dict[int, Any]] = {col: {} for col in dh_columns}

        for idx, row in enumerate(robot_dh_params):
            d_val, theta_rad, r_val, alpha_rad = row[3], row[4], row[5], row[6]
            theta_deg = theta_rad * 180.0 / math.pi
            alpha_deg = alpha_rad * 180.0 / math.pi

            values = [row[0], row[1], row[2], d_val, theta_deg, r_val, alpha_deg]
            for col, val in zip(dh_columns, values):
                dh_dict[col][idx] = val

        inertia_columns = ["link", "mass", "com", "inertia_tensor"]
        inertia_dict: dict[str, dict[int, Any]] = {col: {} for col in inertia_columns}
        for idx, row in enumerate(robot_mass_params):
            for col, val in zip(inertia_columns, row):
                inertia_dict[col][idx] = val

        self.dh_dict = dh_dict
        self.inertia_dict = inertia_dict
        self.moving_joints = moving_joint_names
        # Preserve moving joints in the public urdf_joints field
        self.urdf_joints = {
            name: self.urdf_joints_all[name] for name in moving_joint_names
        }

        # Capture any upstream fixed transform (world -> parent of first moving joint) using URDF abs_tf
        if len(robot_dh_params) > 0:
            first_parent = dh_dict["parent"][0]
            if first_parent in self.urdf_links:
                self.base_prefix_tf = self.urdf_links[first_parent]["abs_tf"]

            # Capture downstream fixed transform from last moving link to tool (if present) using URDF abs_tf
            last_child = dh_dict["child"][len(robot_dh_params) - 1]
            if last_child in self.urdf_links:
                last_child_tf = self.urdf_links[last_child]["abs_tf"]
                tool_abs_tf = None
                if "tool0" in self.urdf_links:
                    tool_abs_tf = self.urdf_links["tool0"].get("abs_tf", None)
                # Fallback: if no explicit tool, leave as identity
                if tool_abs_tf is not None:
                    self.tool_offset_tf = np.matmul(
                        kh.inv_tf(last_child_tf), tool_abs_tf
                    )

    def __get_joint_dh_params(
        self,
        rel_link_frame: np.ndarray,
        axis: np.ndarray,
    ) -> np.ndarray:
        """Compute DH parameters for a joint given its relative frame and axis.

        Args:
            rel_link_frame: Relative transform from parent link to joint frame.
            axis: Joint axis vector.

        Returns:
            `np.ndarray` of length 4 containing DH parameters.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `rel_link_frame` is a valid 4x4 transform and `axis` is length 3.
        """
        dh_params = np.zeros(4)
        origin_xyz = rel_link_frame[0:3, 3]
        z_axis = np.array([0, 0, 1])

        # Collinear case
        if gh.are_collinear(np.zeros(3), z_axis, origin_xyz, axis):
            dh_params = self.__process_collinear_case(
                origin_xyz, rel_link_frame[0:3, 0]
            )

        # Parallel case
        elif gh.are_parallel(z_axis, axis):
            dh_params = self.__process_parallel_case(origin_xyz)

        # Intersect case
        elif gh.lines_intersect(np.zeros(3), z_axis, origin_xyz, axis)[0]:
            dh_params = self.__process_intersection_case(origin_xyz, axis)

        # Skew case
        else:
            dh_params = self.__process_skew_case(origin_xyz, axis)

        for i in range(len(dh_params)):
            if np.isnan(dh_params[i]):
                dh_params[i] = 0
        return dh_params

    def __process_collinear_case(
        self,
        origin: np.ndarray,
        xaxis: np.ndarray,
    ) -> np.ndarray:
        """Compute DH parameters when joint axis is collinear with z-axis.

        Args:
            origin: Joint origin in parent frame.
            xaxis: X-axis of the joint frame (unused).

        Returns:
            `np.ndarray` of length 4 containing DH parameters.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `origin` is length 3.
        """
        dh_params = np.zeros(4)
        dh_params[0] = origin[2]
        return dh_params

    def __process_parallel_case(self, origin: np.ndarray) -> np.ndarray:
        """Compute DH parameters when joint axis is parallel to z-axis.

        Args:
            origin: Joint origin in parent frame.

        Returns:
            `np.ndarray` of length 4 containing DH parameters.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `origin` is length 3.
        """
        dh_params = np.zeros(4)
        dh_params[0] = origin[2]
        dh_params[1] = math.atan2(origin[1], origin[0])
        dh_params[2] = math.sqrt(origin[0] ** 2 + origin[1] ** 2)
        return dh_params

    def __process_intersection_case(
        self,
        origin: np.ndarray,
        axis: np.ndarray,
    ) -> np.ndarray:
        """Compute DH parameters when joint axis intersects the z-axis.

        Args:
            origin: Joint origin in parent frame.
            axis: Joint axis vector.

        Returns:
            `np.ndarray` of length 4 containing DH parameters.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `origin` and `axis` are length 3.
        """
        dh_params = np.zeros(4)
        _, x = gh.lines_intersect(np.zeros(3), np.array([0, 0, 1]), origin, axis)
        # `x` contains the parameters along the parent z-axis (x[0]) and the
        # child joint axis (x[1]) at the point of intersection. For DH we want
        # the offset along the *parent* axis, so use x[0] as d.
        dh_params[0] = float(x[0, 0])

        zaxis = np.array([0.0, 0.0, 1.0])
        for i in range(0, 3):
            if abs(axis[i]) < 1.0e-5:
                axis[i] = 0

        cn = np.cross(zaxis, axis)
        for i in range(0, 3):
            if abs(cn[i]) < 1.0e-6:
                cn[i] = 0
        if cn[0] < 0:
            cn = cn * -1
        dh_params[1] = math.atan2(cn[1], cn[0])
        # print(math.atan2(np.dot(np.cross(xaxis, cn), zaxis), np.dot(xaxis, cn)))

        dh_params[2] = 0

        vn = cn / np.linalg.norm(cn)
        dh_params[3] = math.atan2(
            np.dot(np.cross(zaxis, axis), vn), np.dot(zaxis, axis)
        )

        return dh_params

    def __process_skew_case(
        self,
        origin: np.ndarray,
        direction: np.ndarray,
    ) -> np.ndarray:
        """Compute DH parameters when joint axis is skew to z-axis.

        Args:
            origin: Joint origin in parent frame.
            direction: Joint axis vector.

        Returns:
            `np.ndarray` of length 4 containing DH parameters.

        Side Effects:
            None.

        Raises:
            None.

        Preconditions:
            `origin` and `direction` are length 3.
        """
        pointA = np.zeros(3)
        pointB = np.zeros(3)
        dh_params = np.zeros(4)

        # Avoid exploding values when the axis is nearly parallel to the parent
        # z-axis (tiny x/y components make the projection ill-conditioned).
        denom = direction[0] ** 2 + direction[1] ** 2
        if denom < 1e-9:
            return self.__process_parallel_case(origin)

        # Find closest points along parent z-axis (pointA) and joint axis (pointB)
        t = -1.0 * (origin[0] * direction[0] + origin[1] * direction[1]) / denom
        pointB = origin + t * direction
        pointA[2] = pointB[2]

        # 'd' is offset along parent z axis
        dh_params[0] = pointA[2]

        # 'r' is the length of the common normal
        dh_params[2] = np.linalg.norm(pointB - pointA)

        # 'theta' is the angle between the x-axis and the common normal
        dh_params[1] = math.atan2(pointB[1], pointB[0])

        # 'alpha' is the angle between the current z-axis and the joint axis
        cn = pointB - pointA
        vn = cn / np.linalg.norm(cn)
        zaxis = np.array([0, 0, 1])
        dh_params[3] = math.atan2(
            np.dot(np.cross(zaxis, direction), vn), np.dot(zaxis, direction)
        )

        return dh_params
