shinespark/src/fighter/baked_animation/runtime.rs

use std::collections::HashMap;

use bevy::prelude::*;
use serde::{Deserialize, Serialize};

use super::BAKED_ANIMATION_FPS;

use crate::{
    fighter::{
        Fighter,
        animation::{AnimKind, FighterAnimationFrame},
        manifest::FighterManifest,
        visual::FighterAnimations,
    },
    math::fixed_transform::FixedMat4,
    math::int::FGi32,
};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BakedAnimationClip {
    pub frames: Vec<Vec<FixedMat4>>,
}

#[derive(Asset, bevy::reflect::TypePath, Clone, Debug, Serialize, Deserialize)]
#[type_path = "shinespark::fighter::baked_animation"]
pub struct BakedFighterAnimations {
    pub frame_rate: u32,
    pub bone_names: Vec<String>,
    pub parents: Vec<Option<usize>>,
    pub(super) clips: Vec<BakedAnimationClip>,
    pub(super) clip_indices: HashMap<AnimKind, usize>,
}

impl BakedFighterAnimations {
    pub fn contains_bone(&self, bone: &str) -> bool {
        self.bone_names.iter().any(|name| name == bone)
    }
    pub fn has_unique_bone_names(&self) -> bool {
        let mut names = std::collections::HashSet::with_capacity(self.bone_names.len());
        self.bone_names.iter().all(|name| names.insert(name))
    }
    pub fn sample_pose(&self, frame: &FighterAnimationFrame) -> Option<&[FixedMat4]> {
        let frames = &self.clip(frame.kind)?.frames;
        let frame_index = if frame.repeat {
            let loop_frame_count = frames.len().saturating_sub(1).max(1) as u32;
            frame.frame % loop_frame_count
        } else {
            frame.frame
        };
        frames
            .get(frame_index as usize)
            .or_else(|| frames.last())
            .map(Vec::as_slice)
    }
    pub fn frame_count(&self, kind: AnimKind) -> Option<u32> {
        self.clip(kind).map(|clip| clip.frames.len() as u32)
    }
    /// Samples one bone in model space without depending on the runtime's
    /// previous-frame buffers. Tooling uses this when the user scrubs between
    /// unrelated frames.
    pub fn sample_bone_model_transform(
        &self,
        frame: &FighterAnimationFrame,
        bone: &str,
        model_scale: FGi32,
    ) -> Option<FixedMat4> {
        let local = self.sample_pose(frame)?;
        let index = self.bone_names.iter().position(|name| name == bone)?;
        let mut model = vec![None; local.len()];
        resolve_model_transform(index, &self.parents, local, &mut model).ok()?;
        let scale = FixedMat4 {
            cols: [
                [model_scale, FGi32::ZERO, FGi32::ZERO, FGi32::ZERO],
                [FGi32::ZERO, model_scale, FGi32::ZERO, FGi32::ZERO],
                [FGi32::ZERO, FGi32::ZERO, model_scale, FGi32::ZERO],
                [FGi32::ZERO, FGi32::ZERO, FGi32::ZERO, FGi32::ONE],
            ],
        };
        Some(scale.mul(model[index]?))
    }
    pub(super) fn clip(&self, kind: AnimKind) -> Option<&BakedAnimationClip> {
        self.clip_indices
            .get(&kind)
            .and_then(|index| self.clips.get(*index))
    }
}

#[derive(Component, Clone, Debug, Default)]
pub struct FighterBoneMatrices {
    bone_names: Vec<String>,
    parents: Vec<Option<usize>>,
    matrices: Vec<FixedMat4>,
    previous_frame_matrices: Vec<FixedMat4>,
    local_matrices: Vec<FixedMat4>,
    previous_local_matrices: Vec<FixedMat4>,
    model_matrices: Vec<Option<FixedMat4>>,
    presentation_kind: Option<AnimKind>,
}

#[derive(Component, Clone, Debug, Default)]
pub struct FighterBoneIndexMap(pub HashMap<String, usize>);

impl FighterBoneMatrices {
    /// Samples a baked pose for a preview or other non-gameplay caller.
    pub fn sample_baked_pose(
        &mut self,
        baked: &BakedFighterAnimations,
        frame: &FighterAnimationFrame,
        model_scale: FGi32,
    ) -> bool {
        let Some(pose) = baked.sample_pose(frame) else {
            return false;
        };
        self.set_pose(
            frame.kind,
            &baked.bone_names,
            &baked.parents,
            pose,
            model_scale,
        );
        true
    }

    #[cfg(test)]
    pub(super) fn buffer_capacities(&self) -> [usize; 5] {
        [
            self.matrices.capacity(),
            self.local_matrices.capacity(),
            self.previous_frame_matrices.capacity(),
            self.previous_local_matrices.capacity(),
            self.model_matrices.capacity(),
        ]
    }

    pub fn get_by_idx(&self, index: usize) -> Option<(FixedMat4, FixedMat4)> {
        Some((
            *self.previous_frame_matrices.get(index)?,
            *self.matrices.get(index)?,
        ))
    }
    pub fn get(&self, bone: &str) -> Option<(FixedMat4, FixedMat4)> {
        let index = self.bone_names.iter().position(|name| name == bone)?;
        Some((
            *self.previous_frame_matrices.get(index)?,
            *self.matrices.get(index)?,
        ))
    }
    pub fn get_current(&self, bone: &str) -> Option<FixedMat4> {
        let index = self.bone_names.iter().position(|name| name == bone)?;
        self.matrices.get(index).copied()
    }
    pub fn get_index(&self, index: usize) -> Option<(FixedMat4, FixedMat4)> {
        Some((
            *self.previous_frame_matrices.get(index)?,
            *self.matrices.get(index)?,
        ))
    }
    pub fn get_current_index(&self, index: usize) -> Option<FixedMat4> {
        self.matrices.get(index).copied()
    }
    pub fn iter(&self) -> impl Iterator<Item = (&str, FixedMat4)> {
        self.bone_names
            .iter()
            .map(String::as_str)
            .zip(self.matrices.iter().copied())
    }
    pub fn local(&self, index: usize) -> Option<FixedMat4> {
        self.local_matrices.get(index).copied()
    }
    pub(super) fn presentation_local(
        &self,
        index: usize,
        overstep_fraction: f32,
    ) -> Option<Transform> {
        let current = self.local_matrices.get(index).copied()?;
        let previous = self
            .previous_local_matrices
            .get(index)
            .copied()
            .unwrap_or(current);
        Some(interpolate_transform(previous, current, overstep_fraction))
    }
    pub(super) fn set_pose(
        &mut self,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
    ) {
        self.sample_pose(kind, bone_names, parents, local, model_scale, true);
    }
    pub(super) fn resample_pose(
        &mut self,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
    ) {
        self.sample_pose(kind, bone_names, parents, local, model_scale, false);
    }
    fn sample_pose(
        &mut self,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
        advance_history: bool,
    ) {
        assert_eq!(
            bone_names.len(),
            local.len(),
            "baked animation pose does not match its bone names"
        );
        assert_eq!(
            bone_names.len(),
            parents.len(),
            "baked bone parents do not match names"
        );
        if self.bone_names.is_empty() {
            self.bone_names.extend_from_slice(bone_names);
            self.parents.extend_from_slice(parents);
        } else {
            assert_eq!(self.bone_names, bone_names, "baked bone names changed");
            assert_eq!(self.parents, parents, "baked bone parents changed");
        }
        self.matrices.reserve(local.len());
        self.previous_frame_matrices.reserve(local.len());
        self.local_matrices.reserve(local.len());
        self.previous_local_matrices.reserve(local.len());
        self.model_matrices.reserve(local.len());
        if advance_history {
            std::mem::swap(&mut self.matrices, &mut self.previous_frame_matrices);
            std::mem::swap(&mut self.local_matrices, &mut self.previous_local_matrices);
        }
        let animation_changed = self.presentation_kind != Some(kind);
        self.local_matrices.clear();
        self.local_matrices.extend_from_slice(local);
        self.model_matrices.resize(local.len(), None);
        self.model_matrices.fill(None);
        for index in 0..local.len() {
            resolve_model_transform(index, parents, local, &mut self.model_matrices)
                .expect("baked bone hierarchy should be valid");
        }
        let model_scale = FixedMat4 {
            cols: [
                [model_scale, FGi32::ZERO, FGi32::ZERO, FGi32::ZERO],
                [FGi32::ZERO, model_scale, FGi32::ZERO, FGi32::ZERO],
                [FGi32::ZERO, FGi32::ZERO, model_scale, FGi32::ZERO],
                [FGi32::ZERO, FGi32::ZERO, FGi32::ZERO, FGi32::ONE],
            ],
        };
        self.matrices.clear();
        self.matrices.extend(
            self.model_matrices.iter().map(|matrix| {
                model_scale.mul(matrix.expect("baked bone hierarchy should be valid"))
            }),
        );
        // A state transition is a discontinuity, not an animation frame to
        // interpolate from the previous clip.  Keep the presentation pose at
        // the new clip's first/current frame immediately.
        if animation_changed {
            self.previous_local_matrices
                .clone_from(&self.local_matrices);
        }
        self.presentation_kind = Some(kind);
    }
}

pub fn update_fighter_bone_matrices(
    mut fighters: Query<(
        &FighterAnimationFrame,
        &FighterAnimations,
        &Fighter,
        &mut FighterBoneMatrices,
        &mut FighterBoneIndexMap,
    )>,
    baked_assets: Res<Assets<BakedFighterAnimations>>,
    manifests: Res<Assets<FighterManifest>>,
) {
    for (frame, animations, fighter, mut matrices, mut bone_indices) in &mut fighters {
        sample_fighter_bone_matrices(
            frame,
            animations,
            fighter,
            &mut matrices,
            &mut bone_indices,
            &baked_assets,
            &manifests,
            true,
        );
    }
}

/// Resamples the current pose after collision-driven state transitions without
/// advancing the previous-frame pose used for swept gameplay volumes.
pub fn resample_fighter_bone_matrices(
    mut fighters: Query<(
        &FighterAnimationFrame,
        &FighterAnimations,
        &Fighter,
        &mut FighterBoneMatrices,
        &mut FighterBoneIndexMap,
    )>,
    baked_assets: Res<Assets<BakedFighterAnimations>>,
    manifests: Res<Assets<FighterManifest>>,
) {
    for (frame, animations, fighter, mut matrices, mut bone_indices) in &mut fighters {
        sample_fighter_bone_matrices(
            frame,
            animations,
            fighter,
            &mut matrices,
            &mut bone_indices,
            &baked_assets,
            &manifests,
            false,
        );
    }
}

fn sample_fighter_bone_matrices(
    frame: &FighterAnimationFrame,
    animations: &FighterAnimations,
    fighter: &Fighter,
    matrices: &mut FighterBoneMatrices,
    bone_indices: &mut FighterBoneIndexMap,
    baked_assets: &Assets<BakedFighterAnimations>,
    manifests: &Assets<FighterManifest>,
    advance_history: bool,
) {
    let baked = baked_assets
        .get(&animations.baked)
        .expect("fighter baked animation asset should be loaded before the match starts");
    assert_eq!(
        baked.frame_rate, BAKED_ANIMATION_FPS,
        "fighter baked animation rate does not match the simulation rate"
    );
    let pose = baked
        .sample_pose(frame)
        .expect("fighter animation frame should select a baked pose");
    let manifest = manifests
        .get(&fighter.manifest)
        .expect("fighter manifest should be loaded before the match starts");
    if advance_history {
        matrices.set_pose(
            frame.kind,
            &baked.bone_names,
            &baked.parents,
            pose,
            manifest.model_scale,
        );
    } else {
        matrices.resample_pose(
            frame.kind,
            &baked.bone_names,
            &baked.parents,
            pose,
            manifest.model_scale,
        );
    }
    if bone_indices.0.is_empty() {
        bone_indices.0.extend(
            baked
                .bone_names
                .iter()
                .enumerate()
                .map(|(index, name)| (name.clone(), index)),
        );
    }
}

fn resolve_model_transform(
    index: usize,
    parents: &[Option<usize>],
    local: &[FixedMat4],
    model: &mut [Option<FixedMat4>],
) -> std::io::Result<FixedMat4> {
    if let Some(value) = model[index] {
        return Ok(value);
    }
    let value = match parents[index] {
        Some(parent) => resolve_model_transform(parent, parents, local, model)?.mul(local[index]),
        None => local[index],
    };
    model[index] = Some(value);
    Ok(value)
}

pub(super) fn interpolate_transform(
    previous: FixedMat4,
    current: FixedMat4,
    factor: f32,
) -> Transform {
    let (previous_scale, previous_rotation, previous_translation) =
        previous.to_mat4().to_scale_rotation_translation();
    let (current_scale, current_rotation, current_translation) =
        current.to_mat4().to_scale_rotation_translation();
    let factor = factor.clamp(0.0, 1.0);
    Transform {
        translation: previous_translation.lerp(current_translation, factor),
        rotation: previous_rotation.slerp(current_rotation, factor),
        scale: previous_scale.lerp(current_scale, factor),
    }
}