shinespark/src/fighter/baked_animation/runtime.rs

use std::collections::HashMap;

use bevy::prelude::*;
use bevy_ggrs::Rollback;
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,
    math::vec3::FGVec3,
};

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct BakedAnimationClip {
    pub frames: Vec<Vec<FixedMat4>>,
    /// Unscaled model-space position of the configured root-motion bone at
    /// each baked frame.
    pub root_positions: Vec<FGVec3>,
}

#[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) root_bone_index: usize,
    pub(super) clips: Vec<BakedAnimationClip>,
    pub(super) clip_indices: HashMap<AnimKind, usize>,
}

impl BakedFighterAnimations {
    fn sample_index(&self, frame: &FighterAnimationFrame) -> Option<usize> {
        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
        };
        Some((frame_index as usize).min(frames.len().saturating_sub(1)))
    }
    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 clip = self.clip(frame.kind)?;
        clip.frames.get(self.sample_index(frame)?).map(Vec::as_slice)
    }
    pub fn sample_root_position(
        &self,
        frame: &FighterAnimationFrame,
        model_scale: FGi32,
    ) -> Option<FGVec3> {
        let clip = self.clip(frame.kind)?;
        let position = *clip.root_positions.get(self.sample_index(frame)?)?;
        Some(position * model_scale)
    }
    pub fn initial_root_position(&self, kind: AnimKind, model_scale: FGi32) -> Option<FGVec3> {
        self.sample_root_position(&FighterAnimationFrame::new(kind, false), model_scale)
    }
    pub fn frame_count(&self, kind: AnimKind) -> Option<u32> {
        self.clip(kind).map(|clip| clip.frames.len() as u32)
    }
    pub fn root_motion_cycle_delta(
        &self,
        kind: AnimKind,
        model_scale: FGi32,
    ) -> Option<FGVec3> {
        let clip = self.clip(kind)?;
        Some((*clip.root_positions.last()? - *clip.root_positions.first()?) * model_scale)
    }
    /// 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))
    }
    pub fn get_animation_duration(&self, kind: AnimKind) -> Option<usize> {
        self.clip(kind).map(|clip| clip.frames.len())
    }
}

#[derive(Component, Clone, Debug, Default, Hash)]
#[require(Rollback)]
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>,
    previous_animation_frame: Option<FighterAnimationFrame>,
    root_motion_position: Option<FGVec3>,
    /// The raw, scaled model-space root translation used as the baseline for
    /// root-motion extraction.  This is kept separately because the published
    /// pose has the root translation removed.
    previous_root_motion_translation: Option<FGVec3>,
    root_motion_delta: crate::math::vec3::FGVec3,
}

#[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_with_frame(
            *frame,
            Some(baked.root_bone_index),
            baked.sample_root_position(frame, model_scale),
            FGVec3::ZERO,
            None,
            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()
    }
    /// Translation of the configured root-motion bone since the previous
    /// simulation animation frame, in scaled model-space units.
    pub fn root_motion_delta(&self) -> FGVec3 {
        self.root_motion_delta
    }
    /// Current raw, scaled model-space position of the configured root-motion
    /// bone. The published skeleton has this translation removed.
    pub fn root_motion_position(&self) -> Option<FGVec3> {
        self.root_motion_position
    }
    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))
    }
    #[cfg(test)]
    pub(super) fn set_pose(
        &mut self,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
    ) {
        self.set_pose_with_frame(
            FighterAnimationFrame::new(kind, false),
            None,
            None,
            FGVec3::ZERO,
            None,
            kind,
            bone_names,
            parents,
            local,
            model_scale,
        );
    }
    pub(super) fn set_pose_with_frame(
        &mut self,
        animation_frame: FighterAnimationFrame,
        root_bone_index: Option<usize>,
        root_motion_position: Option<FGVec3>,
        root_motion_cycle_delta: FGVec3,
        root_motion_loop_frame_count: Option<u32>,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
    ) {
        self.sample_pose(
            animation_frame,
            root_bone_index,
            root_motion_position,
            root_motion_cycle_delta,
            root_motion_loop_frame_count,
            kind,
            bone_names,
            parents,
            local,
            model_scale,
            true,
        );
    }
    #[cfg(test)]
    pub(super) fn resample_pose(
        &mut self,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
    ) {
        self.resample_pose_with_frame(
            FighterAnimationFrame::new(kind, false),
            None,
            None,
            FGVec3::ZERO,
            None,
            kind,
            bone_names,
            parents,
            local,
            model_scale,
        );
    }
    pub(super) fn resample_pose_with_frame(
        &mut self,
        animation_frame: FighterAnimationFrame,
        root_bone_index: Option<usize>,
        root_motion_position: Option<FGVec3>,
        _root_motion_cycle_delta: FGVec3,
        _root_motion_loop_frame_count: Option<u32>,
        kind: AnimKind,
        bone_names: &[String],
        parents: &[Option<usize>],
        local: &[FixedMat4],
        model_scale: FGi32,
    ) {
        self.sample_pose(
            animation_frame,
            root_bone_index,
            root_motion_position,
            FGVec3::ZERO,
            None,
            kind,
            bone_names,
            parents,
            local,
            model_scale,
            false,
        );
    }
    fn sample_pose(
        &mut self,
        animation_frame: FighterAnimationFrame,
        root_index: Option<usize>,
        raw_root_motion_translation: Option<FGVec3>,
        root_motion_cycle_delta: FGVec3,
        root_motion_loop_frame_count: Option<u32>,
        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"
        );
        let animation_changed = self.presentation_kind != Some(kind);
        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);
        }
        self.local_matrices.clear();
        self.local_matrices.extend_from_slice(local);

        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],
            ],
        };
        // Gameplay always owns root displacement. Keep rotation and scale in
        // the pose while removing the configured root bone's translation.
        if let Some(root_index) = root_index {
            self.local_matrices[root_index].translate(FGVec3::ZERO);
        }
        self.model_matrices.resize(local.len(), None);
        self.model_matrices.fill(None);
        for index in 0..local.len() {
            resolve_model_transform(index, parents, &self.local_matrices, &mut self.model_matrices)
                .expect("baked bone hierarchy should be valid");
        }
        self.matrices.clear();
        self.matrices.extend(
            self.model_matrices.iter().map(|matrix| {
                model_scale.mul(matrix.expect("baked bone hierarchy should be valid"))
            }),
        );
        self.root_motion_delta = if advance_history
            && !animation_changed
            && self.previous_animation_frame.is_some()
            && let (Some(current), Some(previous)) = (
                raw_root_motion_translation,
                self.previous_root_motion_translation,
            )
        {
            let mut delta = current - previous;
            if animation_frame.repeat
                && self.previous_animation_frame.is_some_and(|previous| {
                    previous.kind == animation_frame.kind
                        && previous.frame.saturating_add(1) == animation_frame.frame
                })
                && root_motion_loop_frame_count.is_some_and(|loop_frame_count| {
                    loop_frame_count > 0 && animation_frame.frame % loop_frame_count == 0
                })
            {
                delta = delta + root_motion_cycle_delta;
            }
            delta
        } else {
            FGVec3::ZERO
        };
        self.root_motion_position = raw_root_motion_translation;
        self.previous_root_motion_translation = raw_root_motion_translation;
        if advance_history {
            self.previous_animation_frame = Some(animation_frame);
        }
        // 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");
    let root_motion_position = baked.sample_root_position(frame, manifest.model_scale);
    let root_motion_cycle_delta = baked
        .root_motion_cycle_delta(frame.kind, manifest.model_scale)
        .unwrap_or(FGVec3::ZERO);
    let root_motion_loop_frame_count = baked
        .frame_count(frame.kind)
        .map(|count| count.saturating_sub(1).max(1));
    if advance_history {
        matrices.set_pose_with_frame(
            *frame,
            Some(baked.root_bone_index),
            root_motion_position,
            root_motion_cycle_delta,
            root_motion_loop_frame_count,
            frame.kind,
            &baked.bone_names,
            &baked.parents,
            pose,
            manifest.model_scale,
        );
    } else {
        matrices.resample_pose_with_frame(
            *frame,
            Some(baked.root_bone_index),
            root_motion_position,
            FGVec3::ZERO,
            None,
            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)),
        );
    }
}

pub(super) 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),
    }
}