roast2d_internal 0.4.0

//! Skeletal animation system for 3D meshes.
//!
//! This module provides types for skeletal animation including:
//! - `Skeleton` - Bone hierarchy
//! - `Animation` - Animation clips with keyframes
//! - `AnimationPlayer` - Playback controller

use glam::{Mat4, Quat, Vec3};

/// Maximum bones per skeleton (GPU shader array limit)
pub const MAX_BONES: usize = 128;

/// A single bone in the skeleton hierarchy
#[derive(Debug, Clone)]
pub struct Bone {
    /// Bone name (from glTF)
    pub name: String,
    /// Index of parent bone, None if root
    pub parent: Option<usize>,
    /// Inverse bind matrix - transforms from mesh space to bone local space
    pub inverse_bind_matrix: Mat4,
    /// Local bind pose (rest position relative to parent)
    pub local_bind_pose: BoneTransform,
}

/// Transform data for a bone (decomposed for interpolation)
#[derive(Debug, Clone, Copy)]
pub struct BoneTransform {
    pub translation: Vec3,
    pub rotation: Quat,
    pub scale: Vec3,
}

impl Default for BoneTransform {
    fn default() -> Self {
        Self {
            translation: Vec3::ZERO,
            rotation: Quat::IDENTITY,
            scale: Vec3::ONE,
        }
    }
}

impl BoneTransform {
    /// Create a new bone transform
    pub fn new(translation: Vec3, rotation: Quat, scale: Vec3) -> Self {
        Self {
            translation,
            rotation,
            scale,
        }
    }

    /// Convert to a 4x4 transformation matrix
    pub fn to_mat4(&self) -> Mat4 {
        Mat4::from_scale_rotation_translation(self.scale, self.rotation, self.translation)
    }

    /// Linearly interpolate between two transforms
    pub fn lerp(&self, other: &Self, t: f32) -> Self {
        Self {
            translation: self.translation.lerp(other.translation, t),
            rotation: self.rotation.slerp(other.rotation, t),
            scale: self.scale.lerp(other.scale, t),
        }
    }
}

/// A skeleton defines the bone hierarchy
#[derive(Debug, Clone, Default)]
pub struct Skeleton {
    pub bones: Vec<Bone>,
    /// Root bone indices (bones with no parent)
    pub roots: Vec<usize>,
}

impl Skeleton {
    /// Create a new empty skeleton
    pub fn new() -> Self {
        Self::default()
    }

    /// Compute the world-space transforms for all bones given local transforms
    pub fn compute_world_transforms(&self, local_transforms: &[BoneTransform]) -> Vec<Mat4> {
        let mut world_transforms = vec![Mat4::IDENTITY; self.bones.len()];

        // Process bones in hierarchy order (parents before children)
        for (i, bone) in self.bones.iter().enumerate() {
            let local_mat = local_transforms
                .get(i)
                .map(|t| t.to_mat4())
                .unwrap_or_else(|| bone.local_bind_pose.to_mat4());

            world_transforms[i] = match bone.parent {
                Some(parent_idx) => world_transforms[parent_idx] * local_mat,
                None => local_mat,
            };
        }

        world_transforms
    }

    /// Compute final skinning matrices (for shader upload)
    pub fn compute_skinning_matrices(&self, world_transforms: &[Mat4]) -> Vec<Mat4> {
        self.bones
            .iter()
            .enumerate()
            .map(|(i, bone)| world_transforms[i] * bone.inverse_bind_matrix)
            .collect()
    }
}

/// A keyframe in an animation channel
#[derive(Debug, Clone)]
pub struct Keyframe<T: Clone> {
    pub time: f32,
    pub value: T,
}

impl<T: Clone> Keyframe<T> {
    pub fn new(time: f32, value: T) -> Self {
        Self { time, value }
    }
}

/// Part of a bone transform that can be animated
#[derive(Debug, Clone, Copy)]
pub enum BoneTransformPart {
    Translation(Vec3),
    Rotation(Quat),
    Scale(Vec3),
}

/// Animation channel targeting a specific bone property
#[derive(Debug, Clone)]
pub enum AnimationChannel {
    Translation {
        bone_index: usize,
        keyframes: Vec<Keyframe<Vec3>>,
    },
    Rotation {
        bone_index: usize,
        keyframes: Vec<Keyframe<Quat>>,
    },
    Scale {
        bone_index: usize,
        keyframes: Vec<Keyframe<Vec3>>,
    },
}

impl AnimationChannel {
    /// Get the bone index this channel targets
    pub fn bone_index(&self) -> usize {
        match self {
            AnimationChannel::Translation { bone_index, .. } => *bone_index,
            AnimationChannel::Rotation { bone_index, .. } => *bone_index,
            AnimationChannel::Scale { bone_index, .. } => *bone_index,
        }
    }

    /// Sample the channel at a given time
    pub fn sample(&self, time: f32) -> (usize, BoneTransformPart) {
        match self {
            AnimationChannel::Translation {
                bone_index,
                keyframes,
            } => {
                let value = Self::interpolate_vec3(keyframes, time);
                (*bone_index, BoneTransformPart::Translation(value))
            }
            AnimationChannel::Rotation {
                bone_index,
                keyframes,
            } => {
                let value = Self::interpolate_quat(keyframes, time);
                (*bone_index, BoneTransformPart::Rotation(value))
            }
            AnimationChannel::Scale {
                bone_index,
                keyframes,
            } => {
                let value = Self::interpolate_vec3(keyframes, time);
                (*bone_index, BoneTransformPart::Scale(value))
            }
        }
    }

    fn interpolate_vec3(keyframes: &[Keyframe<Vec3>], time: f32) -> Vec3 {
        if keyframes.is_empty() {
            return Vec3::ZERO;
        }
        if time <= keyframes[0].time {
            return keyframes[0].value;
        }
        if time >= keyframes.last().unwrap().time {
            return keyframes.last().unwrap().value;
        }

        for i in 0..keyframes.len() - 1 {
            if time >= keyframes[i].time && time < keyframes[i + 1].time {
                let t = (time - keyframes[i].time) / (keyframes[i + 1].time - keyframes[i].time);
                return keyframes[i].value.lerp(keyframes[i + 1].value, t);
            }
        }
        keyframes.last().unwrap().value
    }

    fn interpolate_quat(keyframes: &[Keyframe<Quat>], time: f32) -> Quat {
        if keyframes.is_empty() {
            return Quat::IDENTITY;
        }
        if time <= keyframes[0].time {
            return keyframes[0].value;
        }
        if time >= keyframes.last().unwrap().time {
            return keyframes.last().unwrap().value;
        }

        for i in 0..keyframes.len() - 1 {
            if time >= keyframes[i].time && time < keyframes[i + 1].time {
                let t = (time - keyframes[i].time) / (keyframes[i + 1].time - keyframes[i].time);
                return keyframes[i].value.slerp(keyframes[i + 1].value, t);
            }
        }
        keyframes.last().unwrap().value
    }
}

/// A complete animation clip
#[derive(Debug, Clone, Default)]
pub struct Animation {
    pub name: String,
    pub duration: f32,
    pub channels: Vec<AnimationChannel>,
}

impl Animation {
    /// Create a new animation
    pub fn new(name: impl Into<String>, duration: f32) -> Self {
        Self {
            name: name.into(),
            duration,
            channels: Vec::new(),
        }
    }

    /// Sample all channels at the given time, returning per-bone transforms
    pub fn sample(&self, time: f32, skeleton: &Skeleton) -> Vec<BoneTransform> {
        // Start with bind pose
        let mut transforms: Vec<BoneTransform> =
            skeleton.bones.iter().map(|b| b.local_bind_pose).collect();

        // Apply animated channels
        for channel in &self.channels {
            let (bone_idx, part) = channel.sample(time);
            if bone_idx < transforms.len() {
                match part {
                    BoneTransformPart::Translation(v) => transforms[bone_idx].translation = v,
                    BoneTransformPart::Rotation(q) => transforms[bone_idx].rotation = q,
                    BoneTransformPart::Scale(v) => transforms[bone_idx].scale = v,
                }
            }
        }

        transforms
    }
}

/// Plays animations on a skeleton
#[derive(Debug, Clone)]
pub struct AnimationPlayer {
    /// Index of current animation
    current_animation: Option<usize>,
    /// Current playback time
    current_time: f32,
    /// Playback speed multiplier
    pub speed: f32,
    /// Whether currently playing
    pub playing: bool,
    /// Whether to loop
    pub looping: bool,
    /// Cached bone matrices for GPU upload
    bone_matrices: Vec<Mat4>,
}

impl Default for AnimationPlayer {
    fn default() -> Self {
        Self {
            current_animation: None,
            current_time: 0.0,
            speed: 1.0,
            playing: false,
            looping: true,
            bone_matrices: Vec::new(),
        }
    }
}

impl AnimationPlayer {
    /// Create a new animation player
    pub fn new() -> Self {
        Self::default()
    }

    /// Start playing an animation by index
    pub fn play(&mut self, animation_index: usize) {
        self.current_animation = Some(animation_index);
        self.current_time = 0.0;
        self.playing = true;
    }

    /// Stop playback and reset time
    pub fn stop(&mut self) {
        self.playing = false;
        self.current_time = 0.0;
    }

    /// Pause playback (keeps current time)
    pub fn pause(&mut self) {
        self.playing = false;
    }

    /// Resume playback
    pub fn resume(&mut self) {
        self.playing = true;
    }

    /// Update animation, returns true if bone matrices changed
    pub fn update(&mut self, delta: f32, skeleton: &Skeleton, animations: &[Animation]) -> bool {
        if !self.playing {
            return false;
        }

        let Some(anim_idx) = self.current_animation else {
            return false;
        };
        let Some(animation) = animations.get(anim_idx) else {
            return false;
        };

        // Advance time
        self.current_time += delta * self.speed;

        // Handle looping/completion
        if self.current_time >= animation.duration {
            if self.looping {
                self.current_time %= animation.duration;
            } else {
                self.current_time = animation.duration;
                self.playing = false;
            }
        }

        // Handle negative time (reverse playback)
        if self.current_time < 0.0 {
            if self.looping {
                self.current_time = animation.duration + (self.current_time % animation.duration);
            } else {
                self.current_time = 0.0;
                self.playing = false;
            }
        }

        // Sample animation
        let local_transforms = animation.sample(self.current_time, skeleton);
        let world_transforms = skeleton.compute_world_transforms(&local_transforms);
        self.bone_matrices = skeleton.compute_skinning_matrices(&world_transforms);

        true
    }

    /// Force update bone matrices without advancing time (useful for initial pose)
    pub fn update_pose(&mut self, skeleton: &Skeleton, animations: &[Animation]) {
        let Some(anim_idx) = self.current_animation else {
            // Use bind pose if no animation
            self.bone_matrices = skeleton
                .bones
                .iter()
                .map(|b| b.inverse_bind_matrix)
                .collect();
            return;
        };
        let Some(animation) = animations.get(anim_idx) else {
            return;
        };

        let local_transforms = animation.sample(self.current_time, skeleton);
        let world_transforms = skeleton.compute_world_transforms(&local_transforms);
        self.bone_matrices = skeleton.compute_skinning_matrices(&world_transforms);
    }

    /// Get the computed bone matrices for shader upload
    pub fn bone_matrices(&self) -> &[Mat4] {
        &self.bone_matrices
    }

    /// Get current playback time
    pub fn time(&self) -> f32 {
        self.current_time
    }

    /// Set playback time directly
    pub fn set_time(&mut self, time: f32) {
        self.current_time = time;
    }

    /// Get the current animation index
    pub fn current_animation(&self) -> Option<usize> {
        self.current_animation
    }

    /// Check if animation finished (only relevant when not looping)
    pub fn is_finished(&self) -> bool {
        !self.playing && !self.looping && self.current_animation.is_some()
    }
}

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

    #[test]
    fn test_bone_transform_interpolation() {
        let t1 = BoneTransform::new(Vec3::ZERO, Quat::IDENTITY, Vec3::ONE);
        let t2 = BoneTransform::new(Vec3::new(10.0, 0.0, 0.0), Quat::IDENTITY, Vec3::ONE);

        let lerped = t1.lerp(&t2, 0.5);
        assert!((lerped.translation.x - 5.0).abs() < 0.001);
    }

    #[test]
    fn test_animation_channel_interpolation() {
        let keyframes = vec![
            Keyframe::new(0.0, Vec3::ZERO),
            Keyframe::new(1.0, Vec3::new(10.0, 0.0, 0.0)),
        ];

        let channel = AnimationChannel::Translation {
            bone_index: 0,
            keyframes,
        };

        let (idx, part) = channel.sample(0.5);
        assert_eq!(idx, 0);
        if let BoneTransformPart::Translation(v) = part {
            assert!((v.x - 5.0).abs() < 0.001);
        } else {
            panic!("Expected Translation");
        }
    }
}