zendo-protocol 0.2.0

//! Decoded frame structs with one named field per joint or landmark.

use crate::constants::{
    BODY_ISB_ANGLE_COUNT, BODY_JOINT_COUNT, BODY_LANDMARK_COUNT, HAND_JOINT_COUNT,
    HAND_LANDMARK_COUNT,
};
use crate::types::{
    HandJoint, HandLandmarkName, IsbAngleName, Joint, Landmark, LandmarkName, Quaternion,
};

/// Generates a frame struct: one public named field per joint/landmark, plus
/// wire-order conversion, name lookup, and iteration.
macro_rules! named_frame {
    (
        $(#[$attr:meta])*
        $name:ident, elem = $elem:ty, key = $key:ty, count = $count:expr,
        { $( $field:ident => $variant:ident ),+ $(,)? }
    ) => {
        $(#[$attr])*
        #[derive(Clone, Copy, Debug, Default, PartialEq)]
        pub struct $name {
            $( pub $field: $elem ),+
        }

        impl $name {
            /// Builds a frame from its values in wire order.
            ///
            /// Useful for the Zendo server, which produces frames to encode
            /// rather than decoding them.
            pub const fn from_array(values: [$elem; $count]) -> Self {
                let [ $( $field ),+ ] = values;
                Self { $( $field ),+ }
            }

            /// Returns the frame's values in wire order.
            pub const fn to_array(&self) -> [$elem; $count] {
                [ $( self.$field ),+ ]
            }

            /// Looks up a single value by name.
            pub const fn get(&self, key: $key) -> $elem {
                match key {
                    $( <$key>::$variant => self.$field ),+
                }
            }

            /// Iterates over `(name, value)` pairs in wire order.
            pub fn iter(&self) -> impl Iterator<Item = ($key, $elem)> {
                [ $( (<$key>::$variant, self.$field) ),+ ].into_iter()
            }
        }
    };
}

named_frame! {
    /// Body-joint orientations for one frame (13 joints).
    BodyQuaternionFrame, elem = Quaternion, key = Joint, count = BODY_JOINT_COUNT,
    {
        hips => Hips,
        spine => Spine,
        neck => Neck,
        right_arm => RightArm,
        right_forearm => RightForearm,
        left_arm => LeftArm,
        left_forearm => LeftForearm,
        right_upleg => RightUpLeg,
        right_leg => RightLeg,
        right_foot => RightFoot,
        left_upleg => LeftUpLeg,
        left_leg => LeftLeg,
        left_foot => LeftFoot,
    }
}

named_frame! {
    /// Body-landmark positions for one frame (19 MAIA landmarks).
    BodyLandmarkFrame, elem = Landmark, key = LandmarkName, count = BODY_LANDMARK_COUNT,
    {
        sacroiliac_joint => SacroiliacJoint,
        suprasternal_notch => SuprasternalNotch,
        nose => Nose,
        left_ear => LeftEar,
        right_ear => RightEar,
        left_shoulder => LeftShoulder,
        right_shoulder => RightShoulder,
        left_elbow => LeftElbow,
        right_elbow => RightElbow,
        left_wrist => LeftWrist,
        right_wrist => RightWrist,
        left_hip => LeftHip,
        right_hip => RightHip,
        left_knee => LeftKnee,
        right_knee => RightKnee,
        left_ankle => LeftAnkle,
        right_ankle => RightAnkle,
        left_foot_index => LeftFootIndex,
        right_foot_index => RightFootIndex,
    }
}

named_frame! {
    /// Hand-joint orientations for one frame (16 joints).
    HandQuaternionFrame, elem = Quaternion, key = HandJoint, count = HAND_JOINT_COUNT,
    {
        wrist => Wrist,
        thumb_mcp => ThumbMcp,
        thumb_pip => ThumbPip,
        thumb_dip => ThumbDip,
        index_mcp => IndexMcp,
        index_pip => IndexPip,
        index_dip => IndexDip,
        middle_mcp => MiddleMcp,
        middle_pip => MiddlePip,
        middle_dip => MiddleDip,
        ring_mcp => RingMcp,
        ring_pip => RingPip,
        ring_dip => RingDip,
        pinky_mcp => PinkyMcp,
        pinky_pip => PinkyPip,
        pinky_dip => PinkyDip,
    }
}

named_frame! {
    /// Hand-landmark positions for one frame (21 BlazePose hand landmarks).
    HandLandmarkFrame, elem = Landmark, key = HandLandmarkName, count = HAND_LANDMARK_COUNT,
    {
        wrist => Wrist,
        thumb_cmc => ThumbCmc,
        thumb_mcp => ThumbMcp,
        thumb_ip => ThumbIp,
        thumb_tip => ThumbTip,
        index_mcp => IndexMcp,
        index_pip => IndexPip,
        index_dip => IndexDip,
        index_tip => IndexTip,
        middle_mcp => MiddleMcp,
        middle_pip => MiddlePip,
        middle_dip => MiddleDip,
        middle_tip => MiddleTip,
        ring_mcp => RingMcp,
        ring_pip => RingPip,
        ring_dip => RingDip,
        ring_tip => RingTip,
        pinky_mcp => PinkyMcp,
        pinky_pip => PinkyPip,
        pinky_dip => PinkyDip,
        pinky_tip => PinkyTip,
    }
}

named_frame! {
    /// Body ISB joint angles for one frame (21 scalars, radians).
    ///
    /// Shoulder `plane_of_elevation` fields are `NaN` when the shoulder is near
    /// neutral (singularity). All other values are always finite after a solve.
    BodyIsbAnglesFrame, elem = f64, key = IsbAngleName, count = BODY_ISB_ANGLE_COUNT,
    {
        thorax_lateral_bend => ThoraxLateralBend,
        thorax_axial_rotation => ThoraxAxialRotation,
        neck_flexion => NeckFlexion,
        neck_lateral_bend => NeckLateralBend,
        neck_axial_rotation => NeckAxialRotation,
        right_shoulder_plane_of_elevation => RightShoulderPlaneOfElevation,
        right_shoulder_elevation => RightShoulderElevation,
        left_shoulder_plane_of_elevation => LeftShoulderPlaneOfElevation,
        left_shoulder_elevation => LeftShoulderElevation,
        right_elbow_flexion => RightElbowFlexion,
        left_elbow_flexion => LeftElbowFlexion,
        right_hip_flexion => RightHipFlexion,
        right_hip_adduction => RightHipAdduction,
        right_hip_internal_rotation => RightHipInternalRotation,
        left_hip_flexion => LeftHipFlexion,
        left_hip_adduction => LeftHipAdduction,
        left_hip_internal_rotation => LeftHipInternalRotation,
        right_knee_flexion => RightKneeFlexion,
        left_knee_flexion => LeftKneeFlexion,
        right_ankle_dorsiflexion => RightAnkleDorsiflexion,
        left_ankle_dorsiflexion => LeftAnkleDorsiflexion,
    }
}

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

    #[test]
    fn from_array_then_to_array_round_trips() {
        // Arrange
        let mut values = [Quaternion::default(); BODY_JOINT_COUNT];
        for (i, q) in values.iter_mut().enumerate() {
            *q = Quaternion {
                w: i as f64,
                x: 0.0,
                y: 0.0,
                z: 0.0,
            };
        }

        // Act
        let frame = BodyQuaternionFrame::from_array(values);

        // Assert
        assert_eq!(frame.to_array(), values);
        assert_eq!(frame.hips.w, 0.0);
        assert_eq!(frame.left_foot.w, 12.0);
    }

    #[test]
    fn get_matches_field_access() {
        // Arrange
        let frame = HandLandmarkFrame {
            thumb_tip: Landmark {
                x: 1.0,
                y: 2.0,
                z: 3.0,
                confidence: 0.5,
            },
            ..HandLandmarkFrame::default()
        };

        // Act / Assert
        assert_eq!(frame.get(HandLandmarkName::ThumbTip), frame.thumb_tip);
    }

    #[test]
    fn iter_yields_wire_order() {
        // Arrange
        let frame = BodyLandmarkFrame::default();

        // Act
        let names: Vec<LandmarkName> = frame.iter().map(|(name, _)| name).collect();

        // Assert
        assert_eq!(names, LandmarkName::ALL.to_vec());
    }
}