re_importer 0.32.0

use std::path::{Path, PathBuf};
use std::sync::Arc;

use ahash::{HashMap, HashMapExt as _, HashSet, HashSetExt as _};
use anyhow::{Context as _, bail};
use arrow::array::{Array as _, Float64Array, ListArray, StringArray, StructArray};
use arrow::buffer::OffsetBuffer;
use arrow::compute::cast;
use arrow::datatypes::{DataType, Field, Fields};
use itertools::Itertools as _;
use re_arrow_util::ArrowArrayDowncastRef as _;
use re_chunk::EntityPath;
use re_log_types::EntityPathPart;
use re_sdk_types::Loggable as _;
use re_sdk_types::archetypes::Transform3D;
use re_sdk_types::datatypes::{Quaternion, Vec3D};
use urdf_rs::{Geometry, Joint, Link, Material, Robot};

use super::joint_transform;

const DEFAULT_TF_STATIC_ENTITY_PATH: &str = "tf_static";

/// Returns a short name for the geometry type, used as a path segment.
fn geometry_type_name(geometry: &Geometry) -> &'static str {
    match geometry {
        Geometry::Mesh { .. } => "mesh",
        Geometry::Box { .. } => "box",
        Geometry::Cylinder { .. } => "cylinder",
        Geometry::Capsule { .. } => "capsule",
        Geometry::Sphere { .. } => "sphere",
    }
}

/// Helper struct containing the (root) entity paths where the different parts of the URDF model are logged.
pub(crate) struct UrdfLogPaths {
    /// The root entity path of the robot.
    pub root: EntityPath,

    /// We separate visual and collision geometries under different paths below `root_path`.
    /// This makes it for example easier to toggle the visibility of all visual or collision geometries at once.
    pub visual_root: EntityPath,
    pub collision_root: EntityPath,

    // We log all default transforms to the same entity path because we use Transform3D with frame names.
    pub transforms: EntityPath,
}

impl UrdfLogPaths {
    pub fn new(robot_name: &str, entity_path_prefix: Option<EntityPath>) -> Self {
        let root = match entity_path_prefix {
            Some(prefix) => prefix / EntityPath::from_single_string(robot_name),
            None => EntityPath::from_single_string(robot_name),
        };
        let visual_root = root.clone() / EntityPathPart::new("visual_geometries");
        let collision_root = root.clone() / EntityPathPart::new("collision_geometries");
        let transforms = EntityPath::from_single_string(DEFAULT_TF_STATIC_ENTITY_PATH);

        Self {
            root,
            visual_root,
            collision_root,
            transforms,
        }
    }
}

/// A `.urdf` file loaded into memory (excluding any mesh files).
///
/// Can be used to inspect any link or joint in the URDF.
pub struct UrdfTree {
    /// The dir containing the .urdf file.
    ///
    /// Used to find mesh files (.stl etc) relative to the URDF file.
    pub(crate) urdf_dir: Option<PathBuf>,
    pub(crate) log_paths: UrdfLogPaths,

    name: String,
    root: Link,
    joints: Vec<Joint>,
    links: HashMap<String, Link>,
    children: HashMap<String, Vec<Joint>>,
    materials: HashMap<String, Material>,
    frame_prefix: Option<String>,
}

impl UrdfTree {
    /// Given a path to an `.urdf` file, load it.
    pub fn from_file_path<P: AsRef<Path>>(
        path: P,
        entity_path_prefix: Option<EntityPath>,
    ) -> anyhow::Result<Self> {
        let path = path.as_ref();
        let robot = urdf_rs::read_file(path)?;
        let urdf_dir = path.parent().map(|p| p.to_path_buf());
        Self::new(robot, urdf_dir, entity_path_prefix)
    }

    /// The `urdf_dir` is the directory containing the `.urdf` file,
    /// which can later be used to resolve relative paths to mesh files.
    pub fn new(
        robot: Robot,
        urdf_dir: Option<PathBuf>,
        entity_path_prefix: Option<EntityPath>,
    ) -> anyhow::Result<Self> {
        let urdf_rs::Robot {
            name,
            links,
            joints,
            materials,
        } = robot;

        let materials = materials
            .into_iter()
            .map(|material| (material.name.clone(), material))
            .collect::<HashMap<_, _>>();

        let links: HashMap<String, Link> = links
            .into_iter()
            .map(|link| (link.name.clone(), link))
            .collect();

        let mut children = HashMap::<String, Vec<Joint>>::new();
        let mut child_links = HashSet::<String>::new();

        for joint in &joints {
            children
                .entry(joint.parent.link.clone())
                .or_default()
                .push(joint.clone());

            child_links.insert(joint.child.link.clone());
        }

        let roots = links
            .iter()
            .filter(|(name, _)| !child_links.contains(*name))
            .map(|(_, link)| link)
            .collect_vec();

        let root = match roots.len() {
            0 => {
                bail!("No root link found in URDF");
            }
            1 => roots[0].clone(),
            _ => {
                bail!("Multiple roots in URDF");
            }
        };

        for joint in &joints {
            if !links.contains_key(&joint.child.link) {
                bail!(
                    "Joint '{}' references unknown child link '{}'",
                    joint.name,
                    joint.child.link
                );
            }
            if !links.contains_key(&joint.parent.link) {
                bail!(
                    "Joint '{}' references unknown parent link '{}'",
                    joint.name,
                    joint.parent.link
                );
            }
        }

        let log_paths = UrdfLogPaths::new(&name, entity_path_prefix);

        Ok(Self {
            urdf_dir,
            name,
            root: root.clone(),
            joints,
            links,
            children,
            materials,
            log_paths,
            frame_prefix: None,
        })
    }

    /// Set the frame prefix applied to all frame IDs.
    pub fn with_frame_prefix(mut self, prefix: impl Into<String>) -> Self {
        self.frame_prefix = Some(prefix.into());
        self
    }

    /// Set the entity path used for static transforms emitted by this URDF tree.
    ///
    /// This path is not affected by the tree's entity path prefix.
    pub fn with_static_transform_entity(mut self, entity_path: impl Into<EntityPath>) -> Self {
        self.log_paths.transforms = entity_path.into();
        self
    }

    /// Name of the robot defined in the URDF.
    pub fn name(&self) -> &str {
        &self.name
    }

    /// The frame prefix, if set.
    pub fn frame_prefix(&self) -> Option<&str> {
        self.frame_prefix.as_deref()
    }

    /// Applies [`Self::frame_prefix`] to the given name, if set.
    pub fn apply_frame_prefix(&self, name: &str) -> String {
        match &self.frame_prefix {
            Some(prefix) => format!("{prefix}{name}"),
            None => name.to_owned(),
        }
    }

    /// Computes a [`Transform3D`] for a joint at the given value.
    ///
    /// If [`Self::frame_prefix`] is set, the frame IDs of the transform are prefixed with it.
    pub fn compute_joint_transform(
        &self,
        joint: &Joint,
        value: f64,
        clamp: bool,
    ) -> Result<Transform3D, joint_transform::Error> {
        let result = joint_transform::internal::compute_joint_transform(joint, value, clamp)?;

        if let Some(warning) = &result.warning {
            re_log::warn!("{warning}");
        }

        Ok(Transform3D::update_fields()
            .with_translation(result.translation.to_array())
            .with_quaternion(result.quaternion.to_array())
            .with_parent_frame(self.apply_frame_prefix(&result.parent_frame))
            .with_child_frame(self.apply_frame_prefix(&result.child_frame)))
    }

    /// Computes batches of 3D transform components from Arrow list arrays containing joint names and values.
    ///
    /// `names` must be a `ListArray` with `Utf8` values.
    /// `values` must be a `ListArray` with values castable to `Float64`.
    ///
    /// The output is a `ListArray` with `translation`, `quaternion`, `parent_frame`, and
    /// `child_frame` fields and the same outer row count as the inputs.
    ///
    /// Note: this is intended as a helper for lens pipelines, where you would usually pipe this output
    /// through an additional lens that scatters each batch into final [`Transform3D`] rows.
    pub fn compute_joint_transform_batches(
        &self,
        names: &ListArray,
        values: &ListArray,
        clamp: bool,
    ) -> anyhow::Result<ListArray> {
        if names.len() != values.len() {
            bail!(
                "joint name and value arrays must have the same number of rows, got {} and {}",
                names.len(),
                values.len()
            );
        }

        let joint_names = names
            .values()
            .try_downcast_array_ref::<StringArray>()
            .with_context(|| {
                format!(
                    "joint names must be list<utf8>, got {:?}",
                    names.data_type()
                )
            })?;

        let joint_values = cast(values.values().as_ref(), &DataType::Float64)
            .context("failed to cast joint values to float64")?
            .try_downcast_array::<Float64Array>()?;

        let mut offsets = Vec::<i32>::with_capacity(names.len() + 1);
        let mut translations = Vec::<Vec3D>::new();
        let mut quaternions = Vec::<Quaternion>::new();
        let mut parent_frames = Vec::<String>::new();
        let mut child_frames = Vec::<String>::new();

        offsets.push(0);

        for row in 0..names.len() {
            if names.is_null(row) || values.is_null(row) {
                bail!(
                    "joint name and value lists must not contain null rows, got null at row {row}"
                );
            }

            let names_start = names.value_offsets()[row] as usize;
            let names_end = names.value_offsets()[row + 1] as usize;
            let values_start = values.value_offsets()[row] as usize;
            let values_end = values.value_offsets()[row + 1] as usize;

            let num_names = names_end - names_start;
            let num_values = values_end - values_start;
            if num_names != num_values {
                bail!(
                    "joint name and value lists must have the same length at row {row}, got {num_names} and {num_values}"
                );
            }

            for (name_index, value_index) in (names_start..names_end).zip(values_start..values_end)
            {
                if joint_names.is_null(name_index) || joint_values.is_null(value_index) {
                    bail!(
                        "joint name and value lists must not contain null values, got null at row {row}"
                    );
                }

                let joint_name = joint_names.value(name_index);
                let joint = self
                    .get_joint_by_name(joint_name)
                    .with_context(|| format!("URDF does not contain joint {joint_name:?}"))?;

                let result = joint_transform::internal::compute_joint_transform(
                    joint,
                    joint_values.value(value_index),
                    clamp,
                )?;

                if let Some(warning) = &result.warning {
                    re_log::warn!("{warning}");
                }

                translations.push(Vec3D([
                    result.translation.x,
                    result.translation.y,
                    result.translation.z,
                ]));
                quaternions.push(Quaternion([
                    result.quaternion.x,
                    result.quaternion.y,
                    result.quaternion.z,
                    result.quaternion.w,
                ]));
                parent_frames.push(self.apply_frame_prefix(&result.parent_frame));
                child_frames.push(self.apply_frame_prefix(&result.child_frame));
            }

            offsets.push(
                i32::try_from(parent_frames.len())
                    .context("too many joint transforms for Arrow list offsets")?,
            );
        }

        let translation_array = Vec3D::to_arrow_opt(translations.iter().map(Some))?;
        let quaternion_array = Quaternion::to_arrow_opt(quaternions.iter().map(Some))?;

        let struct_fields = Fields::from(vec![
            Field::new("translation", translation_array.data_type().clone(), false),
            Field::new("quaternion", quaternion_array.data_type().clone(), false),
            Field::new("parent_frame", DataType::Utf8, false),
            Field::new("child_frame", DataType::Utf8, false),
        ]);

        let struct_array = StructArray::try_new(
            struct_fields.clone(),
            vec![
                translation_array,
                quaternion_array,
                Arc::new(StringArray::from(parent_frames)),
                Arc::new(StringArray::from(child_frames)),
            ],
            None,
        )?;

        Ok(ListArray::new(
            Arc::new(Field::new_list_field(
                DataType::Struct(struct_fields),
                false,
            )),
            OffsetBuffer::new(offsets.into()),
            Arc::new(struct_array),
            None,
        ))
    }

    /// The root [`Link`] in the URDF hierarchy.
    pub fn root(&self) -> &Link {
        &self.root
    }

    pub fn joints(&self) -> impl Iterator<Item = &Joint> {
        self.joints.iter()
    }

    pub fn get_joint_by_name(&self, joint_name: &str) -> Option<&Joint> {
        self.joints.iter().find(|j| j.name == joint_name)
    }

    /// Returns the [`Link`] with the given `name`, if it exists.
    pub fn get_link(&self, link_name: &str) -> Option<&Link> {
        self.links.get(link_name)
    }

    /// Returns the child [`Joint`]s of the link with the given `name`, if any.
    pub fn get_children(&self, link_name: &str) -> Option<&Vec<Joint>> {
        self.children.get(link_name)
    }

    /// Get the visual geometries of a link and their entity paths, if any.
    pub fn get_visual_geometries(
        &self,
        link: &Link,
    ) -> Option<Vec<(EntityPath, &urdf_rs::Visual)>> {
        let link = self.links.get(&link.name)?;
        if link.visual.is_empty() {
            return None;
        }

        // The base path for all visual geometries of this link.
        // We use flat paths under `visual_root` since link names have to be unique and to avoid deep nesting.
        let visual_base_path_for_link =
            self.log_paths.visual_root.clone() / EntityPathPart::new(&link.name);

        // Collect all the link's visual geometries and build their entity paths.
        link.visual
            .iter()
            .enumerate()
            .map(|(i, visual)| {
                let visual_name = visual.name.clone().unwrap_or_else(|| format!("visual_{i}"));
                (
                    visual_base_path_for_link.clone() / EntityPathPart::new(visual_name),
                    visual,
                )
            })
            .collect::<Vec<_>>()
            .into()
    }

    /// Get the collision geometries of a link and their entity paths, if any.
    ///
    /// Collision geometries are organized by geometry type under `collision_root`:
    /// `collision_geometries/{geometry_type}/{link_name}/{collision_name}`.
    /// This makes it easy to toggle visibility per geometry type (e.g. hide meshes but keep primitives).
    pub fn get_collision_geometries(
        &self,
        link: &Link,
    ) -> Option<Vec<(EntityPath, &urdf_rs::Collision)>> {
        let link = self.links.get(&link.name)?;
        if link.collision.is_empty() {
            return None;
        }

        // Collect all the link's collision geometries and build their entity paths,
        // organized by geometry type for easy per-type visibility toggling.
        link.collision
            .iter()
            .enumerate()
            .map(|(i, collision)| {
                let geometry_type = geometry_type_name(&collision.geometry);
                let collision_name = collision
                    .name
                    .clone()
                    .unwrap_or_else(|| format!("collision_{i}"));
                let path = self.log_paths.collision_root.clone()
                    / EntityPathPart::new(geometry_type)
                    / EntityPathPart::new(&link.name)
                    / EntityPathPart::new(collision_name);
                (path, collision)
            })
            .collect::<Vec<_>>()
            .into()
    }

    pub fn get_joint_child(&self, joint: &Joint) -> &Link {
        &self.links[&joint.child.link] // Safe because we checked that the joint's child link exists in `new()`
    }

    /// Get a material by name, if it exists.
    pub(crate) fn get_material(&self, name: &str) -> Option<&Material> {
        self.materials.get(name)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_minimal_link(name: &str) -> urdf_rs::Link {
        urdf_rs::Link {
            name: name.to_owned(),
            inertial: Default::default(),
            visual: vec![],
            collision: vec![],
        }
    }

    #[test]
    fn test_apply_frame_prefix_without_prefix() {
        let robot = urdf_rs::Robot {
            name: "test".to_owned(),
            links: vec![make_minimal_link("base")],
            joints: vec![],
            materials: vec![],
        };
        let tree = UrdfTree::new(robot, None, None).unwrap();
        assert_eq!(tree.apply_frame_prefix("base"), "base");
        assert!(tree.frame_prefix().is_none());
    }

    #[test]
    fn test_apply_frame_prefix_with_prefix() {
        let robot = urdf_rs::Robot {
            name: "test".to_owned(),
            links: vec![make_minimal_link("base")],
            joints: vec![],
            materials: vec![],
        };
        let tree = UrdfTree::new(robot, None, None)
            .unwrap()
            .with_frame_prefix("left_arm/");
        assert_eq!(tree.apply_frame_prefix("base"), "left_arm/base");
        assert_eq!(tree.frame_prefix(), Some("left_arm/"));
    }

    #[test]
    fn test_with_static_transform_entity_overrides_default_path() {
        let robot = urdf_rs::Robot {
            name: "test".to_owned(),
            links: vec![make_minimal_link("base")],
            joints: vec![],
            materials: vec![],
        };
        let tree = UrdfTree::new(robot, None, None)
            .unwrap()
            .with_static_transform_entity("robot/tf");

        assert_eq!(tree.log_paths.transforms.to_string(), "/robot/tf");
    }

    #[test]
    fn test_static_transforms_path_defaults_to_tf_static() {
        let paths = UrdfLogPaths::new("test", None);

        assert_eq!(paths.transforms.to_string(), "/tf_static");
    }

    #[test]
    fn test_static_transforms_path_is_unaffected_by_prefix() {
        let paths = UrdfLogPaths::new("test", Some(EntityPath::parse_forgiving("robots/left_arm")));

        assert_eq!(paths.transforms.to_string(), "/tf_static");
    }
}