From 22e0f65f591fdcdc89441181e8f2c37cd722c258 Mon Sep 17 00:00:00 2001 From: Danny McGee Date: Tue, 6 Feb 2024 19:06:41 -0500 Subject: [PATCH] =?UTF-8?q?Technically=20working,=20=F0=9F=92=A9?= =?UTF-8?q?=F0=9F=97=91=EF=B8=8F=F0=9F=94=A5=20code=20quality?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- studio/src/anim.rs | 39 +++- studio/src/debug.rs | 28 ++- studio/src/mixamo/mod.rs | 367 ++++++++++++++++++++++++++++---------- studio/src/mixamo/tree.rs | 177 ++++++++++++++++++ 4 files changed, 509 insertions(+), 102 deletions(-) create mode 100644 studio/src/mixamo/tree.rs diff --git a/studio/src/anim.rs b/studio/src/anim.rs index 64ab974..ae8af7d 100644 --- a/studio/src/anim.rs +++ b/studio/src/anim.rs @@ -6,11 +6,44 @@ pub struct SkeletalJoint(pub Color); #[derive(Component)] pub struct Skeleton; -#[derive(Component, Clone, Copy, Debug, Deref, DerefMut)] +#[derive(Component, Clone, Copy, Debug, Default, Deref, DerefMut)] pub struct BindPoseLocal(pub Transform); -#[derive(Component, Clone, Copy, Debug, Deref, DerefMut)] +#[derive(Component, Clone, Copy, Debug, Default, Deref, DerefMut)] pub struct BindPoseGlobal(pub Affine3A); -#[derive(Component, Clone, Copy, Debug, Deref, DerefMut)] +#[derive(Component, Clone, Copy, Debug, Default, Deref, DerefMut)] pub struct InvBindPose(pub Affine3A); + +pub trait EntityPathExtras { + fn parse(path: &str) -> Self; + fn new(name: Name) -> Self; + fn append(&self, name: Name) -> Self; + fn leaf(&self) -> &Name; +} + +impl EntityPathExtras for EntityPath { + fn parse(path: &str) -> Self { + Self { + parts: path + .split(" > ") + .map(|part| Name::new(part.to_string())) + .collect(), + } + } + + fn new(name: Name) -> Self { + Self { parts: vec![name] } + } + + fn append(&self, name: Name) -> Self { + let mut parts = self.parts.clone(); + parts.push(name); + + Self { parts } + } + + fn leaf(&self) -> &Name { + self.parts.last().unwrap() + } +} diff --git a/studio/src/debug.rs b/studio/src/debug.rs index f5616c6..a36050a 100644 --- a/studio/src/debug.rs +++ b/studio/src/debug.rs @@ -1,6 +1,6 @@ use std::fmt; -use bevy::{prelude::*, window::PrimaryWindow}; +use bevy::{math::Affine3A, prelude::*, window::PrimaryWindow}; use bevy_inspector_egui::quick::WorldInspectorPlugin; use bevy_mod_picking::debug::DebugPickingMode; @@ -139,19 +139,37 @@ fn draw_skeletal_joints( } } +#[rustfmt::skip] +#[allow(dead_code)] +pub fn write_vec3(w: &mut impl fmt::Write, v: &Vec3) -> fmt::Result { + writeln!(w, " {:+.3} {:+.3} {:+.3}", v.x, v.y, v.z) +} + #[rustfmt::skip] #[allow(dead_code)] pub fn write_mat3(w: &mut impl fmt::Write, m: &Mat3) -> fmt::Result { - writeln!(w, " {:+.3} {:+.3} {:+.3}", m.x_axis.x, m.x_axis.y, m.x_axis.z)?; - writeln!(w, " {:+.3} {:+.3} {:+.3}", m.y_axis.x, m.y_axis.y, m.y_axis.z)?; - writeln!(w, " {:+.3} {:+.3} {:+.3}", m.z_axis.x, m.z_axis.y, m.z_axis.z)?; - Ok(()) + write_vec3(w, &m.x_axis)?; + write_vec3(w, &m.y_axis)?; + write_vec3(w, &m.z_axis)?; + Ok(()) } #[allow(dead_code)] pub fn print_mat3(m: &Mat3, label: impl fmt::Display) -> fmt::Result { let mut s = String::new(); write_mat3(&mut s, m)?; + + info!("\n{label}\n{}", s.replace('+', " ")); + + Ok(()) +} + +#[allow(dead_code)] +pub fn print_affine3a(value: &Affine3A, label: impl fmt::Display) -> fmt::Result { + let mut s = String::new(); + write_mat3(&mut s, &Mat3::from(value.matrix3))?; + write_vec3(&mut s, &Vec3::from(value.translation))?; + info!("\n{label}\n{}", s.replace('+', " ")); Ok(()) diff --git a/studio/src/mixamo/mod.rs b/studio/src/mixamo/mod.rs index 4ebf2ba..58d7326 100644 --- a/studio/src/mixamo/mod.rs +++ b/studio/src/mixamo/mod.rs @@ -1,17 +1,23 @@ use bevy::{ gltf::Gltf, - math::{Affine3A, Mat3A}, + math::{Affine3A, Mat3A, Vec3A}, prelude::*, render::mesh::skinning::SkinnedMesh, scene::InstanceId, - utils::{HashMap, HashSet}, + transform::TransformSystem, + utils::{EntityHashMap, HashMap, HashSet}, }; use crate::{ - anim::{BindPoseGlobal, BindPoseLocal, InvBindPose, SkeletalJoint, Skeleton}, + anim::{BindPoseGlobal, BindPoseLocal, EntityPathExtras, InvBindPose, SkeletalJoint, Skeleton}, + math::NearlyEq, scene::SceneState, }; +use self::tree::{RetargetNode, RetargetTree}; + +mod tree; + pub struct MixamoPlugin; impl Plugin for MixamoPlugin { @@ -25,11 +31,25 @@ impl Plugin for MixamoPlugin { import_mixamo_scene, tick_mixamo_import.run_if(resource_exists::()), tick_scene_spawning.run_if(in_state(MixamoState::Spawning)), - retarget_mixamo_anim.run_if( - in_state(SceneState::Active).and_then(in_state(MixamoState::Processing)), + build_retarget_trees.run_if( + in_state(SceneState::Active) + .and_then(in_state(MixamoState::Processing).and_then(run_once())), ), + // build_retarget_trees.pipe(retarget).run_if( + // in_state(SceneState::Active) + // .and_then(in_state(MixamoState::Processing).and_then(run_once())), + // ), + // retarget_mixamo_anim.run_if( + // in_state(SceneState::Active).and_then(in_state(MixamoState::Processing)), + // ), ), ) + .add_systems( + PostUpdate, + copy_poses + .run_if(resource_exists::()) + .after(TransformSystem::TransformPropagate), + ) .add_systems(OnEnter(MixamoState::Spawning), spawn_mixamo_scene) .add_systems(OnEnter(MixamoState::Processing), tag_mixamo_elements); } @@ -55,9 +75,6 @@ pub struct MixamoScene(pub InstanceId); #[derive(Component)] pub struct MixamoRoot; -#[derive(Resource, Deref)] -pub struct MixamoRetargetMap(pub HashMap); - fn import_mixamo_scene( mut cmd: Commands, r_assets: Res, @@ -130,6 +147,8 @@ fn tag_mixamo_elements( q_skinned_meshes: Query<&SkinnedMesh>, mut q_xforms: ParamSet<(TransformHelper, Query<&Transform>)>, ) { + info!("tag_mixamo_elements"); + for skinned_mesh in r_scene_spawner .iter_instance_entities(**r_mixamo_scene) .filter_map(|ent| q_skinned_meshes.get(ent).ok()) @@ -226,6 +245,152 @@ fn retarget_mixamo_anim( st_next.set(MixamoState::None); } +fn build_retarget_trees( + mut cmd: Commands, + ra_gltfs: Res>, + r_mixamo_import: Res, + r_mixamo_scene: Res, + ra_animations: Res>, + r_scene_spawner: Res, + q_target: Query>, + q_joints: Query<&SkeletalJoint>, + q_names: Query<&Name>, + q_children: Query<&Children>, + q_bind_poses: Query<&BindPoseLocal>, +) { + let gltf = ra_gltfs.get(&**r_mixamo_import).unwrap(); + let src_anim_handle = gltf.animations[0].clone(); + let src_anim = ra_animations.get(src_anim_handle.clone()).unwrap(); + + let Some(src_root) = r_scene_spawner + .iter_instance_entities(**r_mixamo_scene) + .find(|ent| match q_names.get(*ent) { + Ok(name) => src_anim.compatible_with(name), + Err(_) => false, + }) + else { + error!("Failed to find compatible entity for animation"); + return; + }; + + let mut src_player = AnimationPlayer::default(); + src_player.play(src_anim_handle.clone()).repeat(); + cmd.entity(src_root).insert(src_player); + + let Ok(tgt_root) = q_target.get_single() else { + error!("Failed to find target root"); + return; + }; + + let tgt_tree = RetargetTree::build(tgt_root, &q_names, &q_children, &q_bind_poses, &q_joints); + + cmd.insert_resource(tgt_tree); +} + +fn copy_poses( + r_retarget_tree: Res, + r_retarget_map: Res, + q_ent_names: Query<(Entity, &Name), With>, + q_inv_bind_poses: Query<&InvBindPose>, + q_parents: Query<&Parent>, + q_world_bind_poses: Query<&BindPoseGlobal>, + q_local_bind_poses: Query<&BindPoseLocal>, + mut q_xforms: Query<&mut Transform>, + q_world_xforms: Query<&GlobalTransform>, +) { + r_retarget_tree.for_each(|node| { + let Some(src_ent) = r_retarget_map + .tgt_to_src + .get(node.path.leaf().as_str()) + .and_then(|src_name| { + q_ent_names.iter().find_map(|(ent, name)| { + if name.as_str() == src_name { + Some(ent) + } else { + None + } + }) + }) + else { + return; + }; + let Ok(src_inv_bind_pose) = q_inv_bind_poses.get(src_ent) else { + return; + }; + let Ok(src_world_xform) = q_world_xforms.get(src_ent) else { + return; + }; + let Ok(tgt_bind_pose) = q_world_bind_poses.get(node.id) else { + return; + }; + + let mut local_mat = src_world_xform.affine() * (**src_inv_bind_pose); + + let parent_xform = node + .parent + .and_then(|parent| q_xforms.get(parent).ok()) + .copied() + .unwrap_or_default(); + + let mut parent_world_xform = GlobalTransform::from(parent_xform); + if let Some(parent) = node.parent { + for ancestor in q_parents.iter_ancestors(parent) { + let xform = q_xforms.get(ancestor).copied().unwrap_or_default(); + parent_world_xform = xform * parent_world_xform; + } + } + + let tgt_inv_parent_mat = parent_world_xform.affine().inverse(); + local_mat = tgt_inv_parent_mat * local_mat * (**tgt_bind_pose); + + let src_local_xlation = q_xforms.get(src_ent).map(|xform| xform.translation); + + let Ok(mut tgt_xform) = q_xforms.get_mut(node.id) else { + return; + }; + tgt_xform.rotation = Quat::from_affine3(&local_mat); + + let Ok(src_local_bind_pose) = q_local_bind_poses.get(src_ent) else { + return; + }; + let Ok(src_world_bind_pose) = q_world_bind_poses.get(src_ent) else { + return; + }; + + if let Ok(src_local_xlation) = src_local_xlation { + if !src_local_xlation.nearly_eq_within(src_local_bind_pose.translation, 1e-3) { + let src_world_xlation = src_world_bind_pose.transform_vector3(src_local_xlation); + let tgt_local_xlation = parent_world_xform + .affine() + .inverse() + .transform_vector3(src_world_xlation); + + tgt_xform.translation = tgt_local_xlation; + } + } + }); +} + +fn retarget( + In(trees): In, RetargetTree, RetargetTree)>>, + mut cmd: Commands, + ra_animations: Res>, +) { + let Some((src_anim, src_tree, tgt_tree)) = trees else { + return; + }; + + let mut src_player = AnimationPlayer::default(); + src_player.play(src_anim.clone()).pause(); + src_player.seek_to(0.); + cmd.entity(src_tree.root).insert(src_player); + + tgt_tree.for_each(|node| { + info!("{}", node.path.leaf()); + // crate::debug::print_affine3a(&node.transform, node.path.leaf()).unwrap(); + }); +} + fn retarget_node( // Source src: Entity, @@ -251,7 +416,7 @@ fn retarget_node( break 'retarget; } - if let Some(tn) = retarget_map.get(src_name.as_str()) { + if let Some(tn) = retarget_map.src_to_tgt.get(src_name.as_str()) { if tn != tgt_name.as_str() { let Some(tgt_children) = tgt_children else { break 'retarget; @@ -346,7 +511,7 @@ fn retarget_node( let Ok((src_name, _)) = q_nodes.get(src_next) else { continue; }; - let Some(tn) = retarget_map.get(src_name.as_str()) else { + let Some(tn) = retarget_map.src_to_tgt.get(src_name.as_str()) else { continue; }; @@ -421,91 +586,105 @@ fn retarget_node( } } +#[derive(Resource)] +pub struct MixamoRetargetMap { + src_to_tgt: HashMap, + tgt_to_src: HashMap, +} + +impl MixamoRetargetMap { + const SRC_TO_TGT: [(&str, &str); 60] = [ + // Spine + ("mixamorig:Hips", "DEF-spine"), + ("mixamorig:Spine", "DEF-spine.001"), + ("mixamorig:Spine1", "DEF-spine.002"), + ("mixamorig:Spine2", "DEF-spine.003"), + ("mixamorig:Neck", "DEF-spine.004"), + ("mixamorig:Head", "DEF-spine.006"), + // Right arm + ("mixamorig:RightShoulder", "DEF-shoulder.R"), + ("mixamorig:RightArm", "DEF-upper_arm.R"), + ("mixamorig:RightForeArm", "DEF-forearm.R"), + ("mixamorig:RightHand", "DEF-hand.R"), + // Right fingers + // -- index + ("mixamorig:RightHandIndex1", "DEF-palm.01.R"), + ("mixamorig:RightHandIndex2", "DEF-f_index.01.R"), + ("mixamorig:RightHandIndex3", "DEF-f_index.02.R"), + ("mixamorig:RightHandIndex4", "DEF-f_index.03.R"), + // -- middle + ("mixamorig:RightHandMiddle1", "DEF-palm.02.R"), + ("mixamorig:RightHandMiddle2", "DEF-f_middle.01.R"), + ("mixamorig:RightHandMiddle3", "DEF-f_middle.02.R"), + ("mixamorig:RightHandMiddle4", "DEF-f_middle.03.R"), + // -- ring + ("mixamorig:RightHandRing1", "DEF-palm.03.R"), + ("mixamorig:RightHandRing2", "DEF-f_ring.01.R"), + ("mixamorig:RightHandRing3", "DEF-f_ring.02.R"), + ("mixamorig:RightHandRing4", "DEF-f_ring.03.R"), + // -- pinky + ("mixamorig:RightHandPinky1", "DEF-palm.04.R"), + ("mixamorig:RightHandPinky2", "DEF-f_pinky.01.R"), + ("mixamorig:RightHandPinky3", "DEF-f_pinky.02.R"), + ("mixamorig:RightHandPinky4", "DEF-f_pinky.03.R"), + // -- thumb + ("mixamorig:RightHandThumb1", "DEF-thumb.01.R"), + ("mixamorig:RightHandThumb2", "DEF-thumb.02.R"), + ("mixamorig:RightHandThumb3", "DEF-thumb.03.R"), + // Right Leg + ("mixamorig:RightUpLeg", "DEF-thigh.R"), + ("mixamorig:RightLeg", "DEF-shin.R"), + ("mixamorig:RightFoot", "DEF-foot.R"), + ("mixamorig:RightToeBase", "DEF-toe.R"), + // Left arm + ("mixamorig:LeftShoulder", "DEF-shoulder.L"), + ("mixamorig:LeftArm", "DEF-upper_arm.L"), + ("mixamorig:LeftForeArm", "DEF-forearm.L"), + ("mixamorig:LeftHand", "DEF-hand.L"), + // Left fingers + // -- index + ("mixamorig:LeftHandIndex1", "DEF-palm.01.L"), + ("mixamorig:LeftHandIndex2", "DEF-f_index.01.L"), + ("mixamorig:LeftHandIndex3", "DEF-f_index.02.L"), + ("mixamorig:LeftHandIndex4", "DEF-f_index.03.L"), + // -- middle + ("mixamorig:LeftHandMiddle1", "DEF-palm.02.L"), + ("mixamorig:LeftHandMiddle2", "DEF-f_middle.01.L"), + ("mixamorig:LeftHandMiddle3", "DEF-f_middle.02.L"), + ("mixamorig:LeftHandMiddle4", "DEF-f_middle.03.L"), + // -- ring + ("mixamorig:LeftHandRing1", "DEF-palm.03.L"), + ("mixamorig:LeftHandRing2", "DEF-f_ring.01.L"), + ("mixamorig:LeftHandRing3", "DEF-f_ring.02.L"), + ("mixamorig:LeftHandRing4", "DEF-f_ring.03.L"), + // -- pinky + ("mixamorig:LeftHandPinky1", "DEF-palm.04.L"), + ("mixamorig:LeftHandPinky2", "DEF-f_pinky.01.L"), + ("mixamorig:LeftHandPinky3", "DEF-f_pinky.02.L"), + ("mixamorig:LeftHandPinky4", "DEF-f_pinky.03.L"), + // -- thumb + ("mixamorig:LeftHandThumb1", "DEF-thumb.01.L"), + ("mixamorig:LeftHandThumb2", "DEF-thumb.02.L"), + ("mixamorig:LeftHandThumb3", "DEF-thumb.03.L"), + // Left Leg + ("mixamorig:LeftUpLeg", "DEF-thigh.L"), + ("mixamorig:LeftLeg", "DEF-shin.L"), + ("mixamorig:LeftFoot", "DEF-foot.L"), + ("mixamorig:LeftToeBase", "DEF-toe.L"), + ]; +} + impl Default for MixamoRetargetMap { fn default() -> Self { - Self( - [ - // Spine - ("mixamorig:Hips", "DEF-spine"), - ("mixamorig:Spine", "DEF-spine.001"), - ("mixamorig:Spine1", "DEF-spine.002"), - ("mixamorig:Spine2", "DEF-spine.003"), - ("mixamorig:Neck", "DEF-spine.004"), - ("mixamorig:Head", "DEF-spine.006"), - // Right arm - ("mixamorig:RightShoulder", "DEF-shoulder.R"), - ("mixamorig:RightArm", "DEF-upper_arm.R"), - ("mixamorig:RightForeArm", "DEF-forearm.R"), - ("mixamorig:RightHand", "DEF-hand.R"), - // Right fingers - // -- index - ("mixamorig:RightHandIndex1", "DEF-palm.01.R"), - ("mixamorig:RightHandIndex2", "DEF-f_index.01.R"), - ("mixamorig:RightHandIndex3", "DEF-f_index.02.R"), - ("mixamorig:RightHandIndex4", "DEF-f_index.03.R"), - // -- middle - ("mixamorig:RightHandMiddle1", "DEF-palm.02.R"), - ("mixamorig:RightHandMiddle2", "DEF-f_middle.01.R"), - ("mixamorig:RightHandMiddle3", "DEF-f_middle.02.R"), - ("mixamorig:RightHandMiddle4", "DEF-f_middle.03.R"), - // -- ring - ("mixamorig:RightHandRing1", "DEF-palm.03.R"), - ("mixamorig:RightHandRing2", "DEF-f_ring.01.R"), - ("mixamorig:RightHandRing3", "DEF-f_ring.02.R"), - ("mixamorig:RightHandRing4", "DEF-f_ring.03.R"), - // -- pinky - ("mixamorig:RightHandPinky1", "DEF-palm.04.R"), - ("mixamorig:RightHandPinky2", "DEF-f_pinky.01.R"), - ("mixamorig:RightHandPinky3", "DEF-f_pinky.02.R"), - ("mixamorig:RightHandPinky4", "DEF-f_pinky.03.R"), - // -- thumb - ("mixamorig:RightHandThumb1", "DEF-thumb.01.R"), - ("mixamorig:RightHandThumb2", "DEF-thumb.02.R"), - ("mixamorig:RightHandThumb3", "DEF-thumb.03.R"), - // Right Leg - ("mixamorig:RightUpLeg", "DEF-thigh.R"), - ("mixamorig:RightLeg", "DEF-shin.R"), - ("mixamorig:RightFoot", "DEF-foot.R"), - ("mixamorig:RightToeBase", "DEF-toe.R"), - // Left arm - ("mixamorig:LeftShoulder", "DEF-shoulder.L"), - ("mixamorig:LeftArm", "DEF-upper_arm.L"), - ("mixamorig:LeftForeArm", "DEF-forearm.L"), - ("mixamorig:LeftHand", "DEF-hand.L"), - // Left fingers - // -- index - ("mixamorig:LeftHandIndex1", "DEF-palm.01.L"), - ("mixamorig:LeftHandIndex2", "DEF-f_index.01.L"), - ("mixamorig:LeftHandIndex3", "DEF-f_index.02.L"), - ("mixamorig:LeftHandIndex4", "DEF-f_index.03.L"), - // -- middle - ("mixamorig:LeftHandMiddle1", "DEF-palm.02.L"), - ("mixamorig:LeftHandMiddle2", "DEF-f_middle.01.L"), - ("mixamorig:LeftHandMiddle3", "DEF-f_middle.02.L"), - ("mixamorig:LeftHandMiddle4", "DEF-f_middle.03.L"), - // -- ring - ("mixamorig:LeftHandRing1", "DEF-palm.03.L"), - ("mixamorig:LeftHandRing2", "DEF-f_ring.01.L"), - ("mixamorig:LeftHandRing3", "DEF-f_ring.02.L"), - ("mixamorig:LeftHandRing4", "DEF-f_ring.03.L"), - // -- pinky - ("mixamorig:LeftHandPinky1", "DEF-palm.04.L"), - ("mixamorig:LeftHandPinky2", "DEF-f_pinky.01.L"), - ("mixamorig:LeftHandPinky3", "DEF-f_pinky.02.L"), - ("mixamorig:LeftHandPinky4", "DEF-f_pinky.03.L"), - // -- thumb - ("mixamorig:LeftHandThumb1", "DEF-thumb.01.L"), - ("mixamorig:LeftHandThumb2", "DEF-thumb.02.L"), - ("mixamorig:LeftHandThumb3", "DEF-thumb.03.L"), - // Left Leg - ("mixamorig:LeftUpLeg", "DEF-thigh.L"), - ("mixamorig:LeftLeg", "DEF-shin.L"), - ("mixamorig:LeftFoot", "DEF-foot.L"), - ("mixamorig:LeftToeBase", "DEF-toe.L"), - ] - .iter() - .map(|(k, v)| (k.to_string(), v.to_string())) - .collect(), - ) + Self { + src_to_tgt: Self::SRC_TO_TGT + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect(), + tgt_to_src: Self::SRC_TO_TGT + .iter() + .map(|(v, k)| (k.to_string(), v.to_string())) + .collect(), + } } } diff --git a/studio/src/mixamo/tree.rs b/studio/src/mixamo/tree.rs new file mode 100644 index 0000000..f21054b --- /dev/null +++ b/studio/src/mixamo/tree.rs @@ -0,0 +1,177 @@ +use bevy::{ + ecs::query::QueryEntityError, + math::{Affine3A, Mat3A}, + prelude::*, + utils::{EntityHashMap, HashMap}, +}; + +use crate::anim::{BindPoseGlobal, BindPoseLocal, EntityPathExtras, SkeletalJoint}; + +#[derive(Resource)] +pub(super) struct RetargetTree { + // pub output: AnimationClip, + pub root: Entity, + pub nodes: EntityHashMap, + pub path_map: HashMap, + // pub timestamps: Vec, +} + +#[derive(Clone)] +pub(super) struct RetargetNode { + pub id: Entity, + pub path: EntityPath, + // pub curves: Vec, + pub children: Vec, + pub parent: Option, + // pub transform: Affine3A, +} + +impl RetargetTree { + pub fn build( + root: Entity, + q_names: &Query<&Name>, + q_children: &Query<&Children>, + q_bind_poses: &Query<&BindPoseLocal>, + q_joints: &Query<&SkeletalJoint>, + ) -> Self { + let mut result = Self { + root, + nodes: EntityHashMap::default(), + path_map: HashMap::default(), + }; + + let root_path = EntityPath::new(q_names.get(root).unwrap().clone()); + let root_children = q_children + .get(root) + .unwrap() + .iter() + .copied() + .filter(|ent| q_joints.contains(*ent)) + .collect(); + + let root_node = RetargetNode { + id: root, + path: root_path.clone(), + children: root_children, + parent: None, + // transform: Affine3A::IDENTITY, + }; + + result.nodes.insert(root, root_node); + result.path_map.insert(root_path, root); + + result.populate( + root, q_names, q_children, // q_bind_poses, + q_joints, + ); + + result + } + + // TODO: Use an actual iterator + // FIXME: Going to need mutable access to nodes, probably through interior mutability + #[inline] + pub fn for_each(&self, mut f: F) + where + F: FnMut(&RetargetNode), + { + self.walk(&self.nodes[&self.root], &mut f); + } + + fn walk(&self, node: &RetargetNode, f: &mut F) + where + F: FnMut(&RetargetNode), + { + f(node); + + for id in node.children.iter() { + self.walk(&self.nodes[id], f); + } + } + + fn populate( + &mut self, + node_id: Entity, + q_names: &Query<&Name>, + q_children: &Query<&Children>, + // q_bind_poses: &Query<&BindPoseLocal>, + q_joints: &Query<&SkeletalJoint>, + ) { + let node = self.nodes[&node_id].clone(); + for &child_id in node.children.iter() { + let name = q_names.get(child_id).unwrap(); + let path = node.path.append(name.clone()); + + // let transform = match q_bind_poses.get(child_id) { + // Ok(bind_pose) => **bind_pose, + // Err(_) => panic!("BindPoseLocal not found for '{name}'"), + // }; + + // let transform = Affine3A { + // matrix3: Mat3A::from_quat(transform.rotation), + // translation: transform.translation.into(), + // }; + + let children = q_children + .get(child_id) + .map(|children| &**children) + .unwrap_or_default() + .iter() + .copied() + .filter(|ent| q_joints.contains(*ent)) + .collect(); + + let child = RetargetNode { + id: child_id, + path: path.clone(), + children, + parent: Some(node_id), + // transform, + }; + + self.nodes.insert(child_id, child); + self.path_map.insert(path, child_id); + + self.populate( + child_id, q_names, q_children, // q_bind_poses, + q_joints, + ); + } + } +} + +impl RetargetNode { + pub fn compute_global_transform( + &self, + q_parents: &Query<&Parent>, + q_xforms: &Q, + ) -> GlobalTransform { + let xform = q_xforms.get_by_value(self.id).unwrap_or_default(); + let mut result = GlobalTransform::from(xform); + + if let Some(parent) = self.parent { + for ancestor in q_parents.iter_ancestors(parent) { + let xform = q_xforms.get_by_value(ancestor).unwrap_or_default(); + result = xform * result; + } + } + + result + } +} + +pub trait TransformQuery { + fn get_by_value(&self, entity: Entity) -> Result; +} + +impl TransformQuery for Query<'_, '_, &Transform> { + fn get_by_value(&self, entity: Entity) -> Result { + self.get(entity).copied() + } +} + +impl TransformQuery for Query<'_, '_, &mut Transform> { + fn get_by_value(&self, entity: Entity) -> Result { + self.get(entity).copied() + } +}