1use crate::CartesianTreeError;
2use crate::Pose;
3use crate::orientation::IntoOrientation;
4use crate::tree::Walking;
5use crate::tree::{HasChildren, HasParent, NodeEquality};
6
7use nalgebra::UnitQuaternion;
8use nalgebra::{Isometry3, Translation3, Vector3};
9use std::cell::RefCell;
10use std::rc::{Rc, Weak};
11
12use serde::{Deserialize, Serialize};
13use serde_json;
14
15#[derive(Clone, Debug)]
22pub struct Frame {
23    pub(crate) data: Rc<RefCell<FrameData>>,
24}
25
26#[derive(Debug)]
27pub(crate) struct FrameData {
28    pub(crate) name: String,
30    parent: Option<Weak<RefCell<FrameData>>>,
32    transform_to_parent: Isometry3<f64>,
34    children: Vec<Frame>,
36}
37
38#[derive(Serialize, Deserialize, Debug, Clone)]
39struct SerialFrame {
40    name: String,
41    position: Vector3<f64>,
42    orientation: UnitQuaternion<f64>,
43    children: Vec<SerialFrame>,
44}
45
46impl Frame {
47    pub fn new_origin(name: impl Into<String>) -> Self {
60        Self {
61            data: Rc::new(RefCell::new(FrameData {
62                name: name.into(),
63                parent: None,
64                children: Vec::new(),
65                transform_to_parent: Isometry3::identity(),
66            })),
67        }
68    }
69
70    pub(crate) fn borrow(&self) -> std::cell::Ref<'_, FrameData> {
71        self.data.borrow()
72    }
73
74    fn borrow_mut(&self) -> std::cell::RefMut<'_, FrameData> {
75        self.data.borrow_mut()
76    }
77
78    pub(crate) fn downgrade(&self) -> Weak<RefCell<FrameData>> {
79        Rc::downgrade(&self.data)
80    }
81
82    pub(crate) fn walk_up_and_transform(
83        &self,
84        target: &Self,
85    ) -> Result<Isometry3<f64>, CartesianTreeError> {
86        let mut transform = Isometry3::identity();
87        let mut current = self.clone();
88
89        while !current.is_same(target) {
90            let transform_to_its_parent = {
91                let current_data = current.borrow();
93
94                if current_data.parent.is_none() {
96                    return Err(CartesianTreeError::IsNoAncestor(target.name(), self.name()));
97                }
98                current_data.transform_to_parent
99            };
100
101            transform = transform_to_its_parent * transform;
102
103            let parent_frame_opt = current.parent();
104            current = parent_frame_opt
105                .ok_or_else(|| CartesianTreeError::IsNoAncestor(target.name(), self.name()))?;
106        }
107
108        Ok(transform)
109    }
110
111    #[must_use]
113    pub fn name(&self) -> String {
114        self.borrow().name.clone()
115    }
116
117    pub fn transform_to_parent(&self) -> Result<Isometry3<f64>, CartesianTreeError> {
126        if self.parent().is_none() {
127            return Err(CartesianTreeError::RootHasNoParent(self.name()));
128        }
129        Ok(self.borrow().transform_to_parent)
130    }
131
132    pub fn update_transform<O>(
161        &self,
162        position: Vector3<f64>,
163        orientation: O,
164    ) -> Result<(), CartesianTreeError>
165    where
166        O: IntoOrientation,
167    {
168        if self.parent().is_none() {
169            return Err(CartesianTreeError::CannotUpdateRootTransform(self.name()));
170        }
171        self.borrow_mut().transform_to_parent =
172            Isometry3::from_parts(Translation3::from(position), orientation.into_orientation());
173        Ok(())
174    }
175
176    pub fn add_child<O>(
205        &self,
206        name: impl Into<String>,
207        position: Vector3<f64>,
208        orientation: O,
209    ) -> Result<Self, CartesianTreeError>
210    where
211        O: IntoOrientation,
212    {
213        let child_name = name.into();
214        {
215            let frame = self.borrow();
216            if frame
217                .children
218                .iter()
219                .any(|child| child.borrow().name == child_name)
220            {
221                return Err(CartesianTreeError::ChildNameConflict(
222                    child_name,
223                    self.name(),
224                ));
225            }
226        }
227        let quat = orientation.into_orientation();
228        let transform = Isometry3::from_parts(Translation3::from(position), quat);
229
230        let child = Self {
231            data: Rc::new(RefCell::new(FrameData {
232                name: child_name,
233                parent: Some(Rc::downgrade(&self.data)),
234                children: Vec::new(),
235                transform_to_parent: transform,
236            })),
237        };
238
239        self.borrow_mut().children.push(child.clone());
240        Ok(child)
241    }
242
243    pub fn calibrate_child<O>(
276        &self,
277        name: impl Into<String>,
278        desired_position: Vector3<f64>,
279        desired_orientation: O,
280        reference_pose: &Pose,
281    ) -> Result<Self, CartesianTreeError>
282    where
283        O: IntoOrientation,
284    {
285        let reference_frame = reference_pose.frame().ok_or_else(|| {
286            CartesianTreeError::FrameDropped("Reference pose frame has been dropped".to_string())
287        })?;
288
289        let ancestor = self.lca_with(&reference_frame).ok_or_else(|| {
290            CartesianTreeError::NoCommonAncestor(self.name(), reference_frame.name())
291        })?;
292
293        let t_reference_to_ancestor = reference_frame.walk_up_and_transform(&ancestor)?;
294        let t_pose_to_reference = reference_pose.transformation();
295        let t_pose_to_ancestor = t_reference_to_ancestor * t_pose_to_reference;
296
297        let t_parent_to_ancestor = self.walk_up_and_transform(&ancestor)?;
298        let t_ancestor_to_parent = t_parent_to_ancestor.inverse();
299
300        let desired_pose = Isometry3::from_parts(
301            Translation3::from(desired_position),
302            desired_orientation.into_orientation(),
303        );
304
305        let t_calibrated_to_parent =
306            t_pose_to_ancestor * desired_pose.inverse() * t_ancestor_to_parent;
307
308        self.add_child(
309            name,
310            t_calibrated_to_parent.translation.vector,
311            t_calibrated_to_parent.rotation,
312        )
313    }
314
315    pub fn add_pose<O>(&self, position: Vector3<f64>, orientation: O) -> Pose
333    where
334        O: IntoOrientation,
335    {
336        Pose::new(self.downgrade(), position, orientation)
337    }
338
339    pub fn to_json(&self) -> Result<String, CartesianTreeError> {
351        let serial = self.to_serial();
352        Ok(serde_json::to_string_pretty(&serial)?)
353    }
354
355    fn to_serial(&self) -> SerialFrame {
359        let (position, orientation) = if self.parent().is_some() {
360            let iso = self
361                .transform_to_parent()
362                .unwrap_or_else(|_| Isometry3::identity());
363            (iso.translation.vector, iso.rotation)
364        } else {
365            (Vector3::zeros(), UnitQuaternion::identity())
366        };
367
368        SerialFrame {
369            name: self.name(),
370            position,
371            orientation,
372            children: self.children().into_iter().map(|c| c.to_serial()).collect(),
373        }
374    }
375
376    pub fn apply_config(&self, json: &str) -> Result<(), CartesianTreeError> {
394        let serial: SerialFrame = serde_json::from_str(json)?;
395        self.apply_serial(&serial)
396    }
397
398    fn apply_serial(&self, serial: &SerialFrame) -> Result<(), CartesianTreeError> {
399        if self.name() != serial.name {
400            return Err(CartesianTreeError::Mismatch(format!(
401                "Frame names do not match: {} vs {}",
402                self.name(),
403                serial.name
404            )));
405        }
406
407        if self.parent().is_some() {
409            self.update_transform(serial.position, serial.orientation)?;
410        }
411
412        for potential_child in &serial.children {
413            if let Some(child) = self
414                .children()
415                .into_iter()
416                .find(|c| c.name() == potential_child.name)
417            {
418                child.apply_serial(potential_child)?;
419            }
420        }
421
422        Ok(())
423    }
424}
425
426impl HasParent for Frame {
427    type Node = Self;
428
429    fn parent(&self) -> Option<Self::Node> {
430        self.borrow()
431            .parent
432            .clone()
433            .and_then(|data_weak| data_weak.upgrade().map(|data_rc| Self { data: data_rc }))
434    }
435}
436
437impl NodeEquality for Frame {
438    fn is_same(&self, other: &Self) -> bool {
439        Rc::ptr_eq(&self.data, &other.data)
440    }
441}
442
443impl HasChildren for Frame {
444    type Node = Self;
445    fn children(&self) -> Vec<Self> {
446        self.borrow().children.clone()
447    }
448}
449
450#[cfg(test)]
451mod tests {
452    use super::*;
453    use nalgebra::{UnitQuaternion, Vector3};
454
455    #[test]
456    fn create_origin_frame() {
457        let root = Frame::new_origin("world");
458        let root_borrow = root.borrow();
459        assert_eq!(root_borrow.name, "world");
460        assert!(root_borrow.parent.is_none());
461        assert_eq!(root_borrow.children.len(), 0);
462    }
463
464    #[test]
465    fn add_child_frame_with_quaternion() {
466        let root = Frame::new_origin("world");
467        let child = root
468            .add_child(
469                "dummy",
470                Vector3::new(1.0, 0.0, 0.0),
471                UnitQuaternion::identity(),
472            )
473            .unwrap();
474
475        let root_borrow = root.borrow();
476        assert_eq!(root_borrow.children.len(), 1);
477
478        let child_borrow = child.borrow();
479        assert_eq!(child_borrow.name, "dummy");
480        assert!(child_borrow.parent.is_some());
481
482        let parent_name = child_borrow
483            .parent
484            .as_ref()
485            .unwrap()
486            .upgrade()
487            .unwrap()
488            .borrow()
489            .name
490            .clone();
491        assert_eq!(parent_name, "world");
492    }
493
494    #[test]
495    fn add_child_frame_with_rpy() {
496        let root = Frame::new_origin("world");
497        let child = root
498            .add_child(
499                "dummy",
500                Vector3::new(0.0, 1.0, 0.0),
501                (0.0, 0.0, std::f64::consts::FRAC_PI_2),
502            )
503            .unwrap();
504
505        let child_borrow = child.borrow();
506        assert_eq!(child_borrow.name, "dummy");
507
508        let rotation = child_borrow.transform_to_parent.rotation;
509        let expected = UnitQuaternion::from_euler_angles(0.0, 0.0, std::f64::consts::FRAC_PI_2);
510        assert!((rotation.angle() - expected.angle()).abs() < 1e-10);
511    }
512
513    #[test]
514    fn multiple_child_frames() {
515        let root = Frame::new_origin("world");
516
517        let a = root
518            .add_child("a", Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
519            .unwrap();
520        let b = root
521            .add_child("b", Vector3::new(0.0, 1.0, 0.0), UnitQuaternion::identity())
522            .unwrap();
523
524        let root_borrow = root.borrow();
525        assert_eq!(root_borrow.children.len(), 2);
526
527        let a_borrow = a.borrow();
528        let b_borrow = b.borrow();
529
530        assert_eq!(
531            a_borrow
532                .parent
533                .as_ref()
534                .unwrap()
535                .upgrade()
536                .unwrap()
537                .borrow()
538                .name,
539            "world"
540        );
541        assert_eq!(
542            b_borrow
543                .parent
544                .as_ref()
545                .unwrap()
546                .upgrade()
547                .unwrap()
548                .borrow()
549                .name,
550            "world"
551        );
552    }
553
554    #[test]
555    fn reject_duplicate_child_name() {
556        let root = Frame::new_origin("world");
557
558        let _ = root
559            .add_child(
560                "duplicate",
561                Vector3::new(1.0, 0.0, 0.0),
562                UnitQuaternion::identity(),
563            )
564            .unwrap();
565
566        let result = root.add_child(
567            "duplicate",
568            Vector3::new(2.0, 0.0, 0.0),
569            UnitQuaternion::identity(),
570        );
571        assert!(result.is_err());
572    }
573
574    #[test]
575    #[should_panic(expected = "already borrowed")]
576    fn test_borrow_conflict() {
577        let frame = Frame::new_origin("root");
578        let _borrow = frame.borrow(); frame.borrow_mut(); }
581
582    #[test]
583    fn test_add_pose_to_frame() {
584        let frame = Frame::new_origin("dummy");
585        let pose = frame.add_pose(Vector3::new(1.0, 2.0, 3.0), UnitQuaternion::identity());
586
587        assert_eq!(pose.frame().unwrap().name(), "dummy");
588    }
589
590    #[test]
591    fn test_update_transform() {
592        let root = Frame::new_origin("root");
593        let child = root
594            .add_child(
595                "dummy",
596                Vector3::new(0.0, 0.0, 1.0),
597                UnitQuaternion::identity(),
598            )
599            .unwrap();
600        child
601            .update_transform(Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
602            .unwrap();
603        assert_eq!(
604            child.transform_to_parent().unwrap().translation.vector,
605            Vector3::new(1.0, 0.0, 0.0)
606        );
607
608        assert!(
610            root.update_transform(Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
611                .is_err()
612        );
613    }
614
615    #[test]
616    fn test_pose_transformation_between_frames() {
617        let root = Frame::new_origin("root");
618
619        let f1 = root
620            .add_child(
621                "f1",
622                Vector3::new(1.0, 0.0, 0.0),
623                UnitQuaternion::identity(),
624            )
625            .unwrap();
626
627        let f2 = f1
628            .add_child(
629                "f2",
630                Vector3::new(0.0, 2.0, 0.0),
631                UnitQuaternion::identity(),
632            )
633            .unwrap();
634
635        let pose_in_f2 = f2.add_pose(Vector3::new(1.0, 1.0, 0.0), UnitQuaternion::identity());
636
637        let pose_in_root = pose_in_f2.in_frame(&root).unwrap();
638        let pos = pose_in_root.transformation().translation.vector;
639
640        assert!((pos - Vector3::new(2.0, 3.0, 0.0)).norm() < 1e-6);
642    }
643
644    #[test]
645    fn test_calibrate_child() {
646        let root = Frame::new_origin("root");
647
648        let reference_pose = root.add_pose(
649            Vector3::new(1.0, 2.0, 3.0),
650            UnitQuaternion::from_euler_angles(0.0, 0.0, std::f64::consts::FRAC_PI_2),
651        );
652
653        let calibrated_frame = root
655            .calibrate_child(
656                "calibrated",
657                Vector3::zeros(),
658                UnitQuaternion::identity(),
659                &reference_pose,
660            )
661            .unwrap();
662
663        let pose_in_calibrated = reference_pose.in_frame(&calibrated_frame).unwrap();
664        let transformation = pose_in_calibrated.transformation();
665
666        assert!((transformation.translation.vector - Vector3::zeros()).norm() < 1e-6);
667        assert!((transformation.rotation.angle() - 0.0).abs() < 1e-6);
668
669        let calibrated_transformation = calibrated_frame.transform_to_parent().unwrap();
671        assert!(
672            (calibrated_transformation.translation.vector - Vector3::new(1.0, 2.0, 3.0)).norm()
673                < 1e-6
674        );
675        assert!(
676            (calibrated_transformation.rotation.angle() - std::f64::consts::FRAC_PI_2).abs() < 1e-6
677        );
678    }
679
680    #[test]
681    fn test_to_json_and_apply_config() {
682        let root = Frame::new_origin("root");
683        let _ = root
684            .add_child(
685                "child",
686                Vector3::new(1.0, 2.0, 3.0),
687                UnitQuaternion::from_euler_angles(0.1, 0.2, 0.3),
688            )
689            .unwrap();
690
691        let json = root.to_json().unwrap();
692        assert!(json.contains(r#""name": "root""#));
694        assert!(json.contains(r#""name": "child""#));
695
696        let default_root = Frame::new_origin("root");
698        default_root
699            .add_child(
700                "child",
701                Vector3::new(0.0, 0.0, 0.0),
702                UnitQuaternion::identity(),
703            )
704            .unwrap();
705
706        default_root.apply_config(&json).unwrap();
708
709        let updated_child = default_root
711            .children()
712            .into_iter()
713            .find(|c| c.name() == "child")
714            .unwrap();
715        let iso = updated_child.transform_to_parent().unwrap();
716        assert_eq!(iso.translation.vector, Vector3::new(1.0, 2.0, 3.0));
717        let (r, p, y) = iso.rotation.euler_angles();
718        assert!((r - 0.1).abs() < 1e-6);
719        assert!((p - 0.2).abs() < 1e-6);
720        assert!((y - 0.3).abs() < 1e-6);
721
722        let partial_json = r#"
724        {
725            "name": "root",
726            "position": [0.0, 0.0, 0.0],
727            "orientation": [0.0, 0.0, 0.0, 1.0],
728            "children": [
729                {
730                    "name": "child",
731                    "position": [4.0, 5.0, 6.0],
732                    "orientation": [0.0, 0.0, 0.0, 1.0],
733                    "children": []
734                },
735                {
736                    "name": "extra",
737                    "position": [0.0, 0.0, 0.0],
738                    "orientation": [0.0, 0.0, 0.0, 1.0],
739                    "children": []
740                }
741            ]
742        }
743        "#;
744        default_root.apply_config(partial_json).unwrap();
745        let updated_child = default_root
746            .children()
747            .into_iter()
748            .find(|c| c.name() == "child")
749            .unwrap();
750        assert_eq!(
751            updated_child
752                .transform_to_parent()
753                .unwrap()
754                .translation
755                .vector,
756            Vector3::new(4.0, 5.0, 6.0)
757        );
758
759        let mismatch_json = r#"
761        {
762            "name": "wrong_root",
763            "position": [0.0, 0.0, 0.0],
764            "orientation": [0.0, 0.0, 0.0, 1.0],
765            "children": []
766        }
767        "#;
768        assert!(default_root.apply_config(mismatch_json).is_err());
769    }
770}