engine/gltf-loader/src/animation.rs

use super::{
    channel::{Channel, InterpolatedTransforms},
    node::Node,
    skin::Skin,
    transforms::Transform,
};

use anyhow::Result;

use gltf;
use gltf::animation::util::ReadOutputs;

use utilities::prelude::cgmath;
use utilities::prelude::cgmath::{Matrix4, Quaternion, Vector3};

use std::collections::HashMap;
use std::ops::IndexMut;
use std::sync::Arc;
use std::time::Duration;

#[derive(Debug)]
pub struct GltfAnimation {
    name: String,
    channels: HashMap<usize, Channel>,

    duration: Duration,

    skins: Arc<Vec<Skin>>,
    nodes: Arc<HashMap<usize, Node>>,
}

impl GltfAnimation {
    pub fn new(
        gltf_animation: &gltf::Animation<'_>,
        buffers: &Vec<gltf::buffer::Data>,
        skins: Arc<Vec<Skin>>,
        nodes: Arc<HashMap<usize, Node>>,
    ) -> Result<Self> {
        let mut channels: HashMap<usize, Channel> = HashMap::new();
        let mut duration = Duration::from_secs(0);

        // collects all channels
        // all node ids (target nodes) are unique
        // that means all transformations at a given time are collected into one transform object
        for gltf_channel in gltf_animation.channels() {
            let reader = gltf_channel.reader(|buffer| Some(&buffers[buffer.index()]));

            let read_outputs = match reader.read_outputs() {
                Some(outputs) => outputs,
                None => return Err(anyhow::Error::msg("Gltf: Missing Outputs from Channel")),
            };

            let outputs: Vec<Transform> = match read_outputs {
                ReadOutputs::Translations(trans) => trans
                    .map(|t| Transform::Translation(Vector3::from(t)))
                    .collect(),
                ReadOutputs::Rotations(rots) => rots
                    .into_f32()
                    .map(|r| {
                        let rot = Quaternion::new(r[3], r[0], r[1], r[2]);
                        Transform::Rotation(rot)
                    })
                    .collect(),
                ReadOutputs::Scales(scales) => {
                    scales.map(|s| Transform::Scale(Vector3::from(s))).collect()
                }
                _ => todo!("MorphTarget not supported"),
            };

            match channels.get_mut(&gltf_channel.target().node().index()) {
                Some(channel) => {
                    channel.update(&gltf_channel, buffers, outputs)?;
                    channel.sort_by_time();
                }
                None => {
                    // if channel with target node does not exist, create a new one
                    let mut channel = Channel::new(&gltf_channel, buffers, outputs)?;

                    channel.sort_by_time();

                    let channel_duration = channel.duration();

                    if channel_duration > duration {
                        duration = channel_duration;
                    }

                    channels.insert(gltf_channel.target().node().index(), channel);
                }
            }
        }

        for (_, channel) in channels.iter_mut() {
            channel.compress();
        }

        let name = match gltf_animation.name() {
            Some(name) => name.to_string(),
            None => format!("Animation{}", gltf_animation.index()),
        };

        Ok(Self {
            name,
            channels,

            duration,

            skins,
            nodes,
        })
    }

    pub fn name(&self) -> &String {
        &self.name
    }

    pub fn duration(&self) -> Duration {
        self.duration
    }

    // shared reference to all animations of an asset
    pub fn skins(&self) -> &Arc<Vec<Skin>> {
        &self.skins
    }

    fn transform(&self, bone_id: usize, time: Duration) -> Option<InterpolatedTransforms> {
        match self.channels.get(&bone_id) {
            Some(channel) => Some(channel.transform(time)),
            None => None,
        }
    }

    #[inline]
    fn get_skin(&self, node_index: usize) -> Option<&Skin> {
        if let Some(node) = self.nodes.get(&node_index) {
            if let Some(skin_index) = node.skin_index {
                return Some(&self.skins[skin_index]);
            }
        }

        None
    }

    pub fn animate(
        &self,
        time: Duration,
        node_index: usize,
        buffer: &mut impl IndexMut<usize, Output = Matrix4<f32>>,
    ) -> Result<()> {
        debug_assert!(time <= self.duration);

        if let Some(skin) = self.get_skin(node_index) {
            self.skeletal_animation(time, skin.root_joint_id, None, skin, buffer)?;
        }

        Ok(())
    }

    pub fn reset(
        &self,
        node_index: usize,
        buffer: &mut impl IndexMut<usize, Output = Matrix4<f32>>,
    ) -> Result<()> {
        if let Some(skin) = self.get_skin(node_index) {
            self.skeletal_reset(skin.root_joint_id, &skin, buffer)?;
        }

        Ok(())
    }

    fn skeletal_reset(
        &self,
        bone_id: usize,
        skin: &Skin,
        buffer: &mut impl IndexMut<usize, Output = Matrix4<f32>>,
    ) -> Result<()> {
        if let Some(bone) = self.nodes.get(&bone_id) {
            if let Some(ibm) = skin.ibm(bone_id) {
                buffer[bone.index] =
                    skin.inverse_global_transform()? * bone.final_transform() * ibm;
            }

            for child_id in bone.children() {
                self.skeletal_reset(*child_id, skin, buffer)?;
            }
        }

        Ok(())
    }

    fn skeletal_animation(
        &self,
        time: Duration,
        bone_id: usize,
        parent_info: Option<cgmath::Matrix4<f32>>,
        skin: &Skin,
        buffer: &mut impl IndexMut<usize, Output = Matrix4<f32>>,
    ) -> Result<()> {
        if let Some(bone) = self.nodes.get(&bone_id) {
            let mut transform_matrix = match self.transform(bone.index, time) {
                Some(transform) => transform.collapse(),
                None => bone.matrix(),
            };

            if let Some(parent_transform) = parent_info {
                transform_matrix = parent_transform * transform_matrix;
            }

            if let Some(ibm) = skin.ibm(bone_id) {
                let joint_matrix = skin.inverse_global_transform()? * transform_matrix * ibm;
                buffer[bone.index] = joint_matrix;
            }

            for child_id in bone.children() {
                self.skeletal_animation(time, *child_id, Some(transform_matrix), skin, buffer)?;
            }
        }

        Ok(())
    }
}