shinespark/src/fighter/baked_animation/asset_baking.rs

use super::{
    BAKED_ANIMATION_FPS,
    runtime::{BakedAnimationClip, BakedFighterAnimations},
};
use crate::math::fixed_transform::FixedMat4;
use crate::fighter::{animation::AnimKind, manifest::FighterManifest};
use bevy::{
    asset::{
        AssetLoader, AssetPath, AsyncWriteExt, LoadContext,
        io::{Reader, Writer},
        processor::LoadTransformAndSave,
        saver::{AssetSaver, SavedAsset},
        transformer::{AssetTransformer, TransformedAsset},
    },
    prelude::*,
    reflect::TypePath,
};
use serde::Deserialize;
use std::{collections::HashMap, io, path::Path};

#[derive(Deserialize)]
struct FighterBakeDeclaration {
    manifest_path: String,
}
#[derive(Asset, TypePath)]
struct FighterBakeSource {
    model_bytes: Vec<u8>,
    buffers: Vec<Vec<u8>>,
    animations: HashMap<AnimKind, String>,
}
#[derive(Default, TypePath)]
struct FighterBakeSourceLoader;

impl AssetLoader for FighterBakeSourceLoader {
    type Asset = FighterBakeSource;
    type Settings = ();
    type Error = io::Error;
    async fn load(
        &self,
        reader: &mut dyn Reader,
        _: &Self::Settings,
        load_context: &mut LoadContext<'_>,
    ) -> Result<Self::Asset, Self::Error> {
        let mut declaration = Vec::new();
        reader.read_to_end(&mut declaration).await?;
        let declaration: FighterBakeDeclaration = ron::de::from_bytes(&declaration)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        let manifest_bytes = load_context
            .read_asset_bytes(&declaration.manifest_path)
            .await
            .map_err(|error| io::Error::other(error.to_string()))?;
        let manifest: FighterManifest = ron::de::from_bytes(&manifest_bytes)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        let model_bytes = load_context
            .read_asset_bytes(&manifest.model_path)
            .await
            .map_err(|error| io::Error::other(error.to_string()))?;
        let gltf = gltf::Gltf::from_slice(&model_bytes)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        let mut buffers = Vec::with_capacity(gltf.buffers().len());
        for buffer in gltf.buffers() {
            let bytes = match buffer.source() {
                gltf::buffer::Source::Uri(uri) => load_context
                    .read_asset_bytes(AssetPath::from_path_buf(
                        Path::new(&manifest.model_path)
                            .parent()
                            .unwrap_or_else(|| Path::new(""))
                            .join(uri),
                    ))
                    .await
                    .map_err(|error| io::Error::other(error.to_string()))?,
                gltf::buffer::Source::Bin => gltf.blob.clone().ok_or_else(|| {
                    io::Error::new(io::ErrorKind::InvalidData, "missing glTF binary blob")
                })?,
            };
            buffers.push(bytes);
        }
        Ok(FighterBakeSource {
            model_bytes,
            buffers,
            animations: manifest.animations,
        })
    }
    fn extensions(&self) -> &[&str] {
        &["fighterbake"]
    }
}

#[derive(Default, TypePath)]
struct BakedFighterAnimationsLoader;
impl AssetLoader for BakedFighterAnimationsLoader {
    type Asset = BakedFighterAnimations;
    type Settings = ();
    type Error = io::Error;
    async fn load(
        &self,
        reader: &mut dyn Reader,
        _: &Self::Settings,
        _: &mut LoadContext<'_>,
    ) -> Result<Self::Asset, Self::Error> {
        let mut bytes = Vec::new();
        reader.read_to_end(&mut bytes).await?;
        bincode::deserialize(&bytes)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))
    }
    fn extensions(&self) -> &[&str] {
        &["fighterbake"]
    }
}

#[derive(Default, TypePath)]
struct FighterBakeTransformer;
impl AssetTransformer for FighterBakeTransformer {
    type AssetInput = FighterBakeSource;
    type AssetOutput = BakedFighterAnimations;
    type Settings = ();
    type Error = io::Error;
    async fn transform<'a>(
        &'a self,
        source: TransformedAsset<Self::AssetInput>,
        _: &'a Self::Settings,
    ) -> Result<TransformedAsset<Self::AssetOutput>, Self::Error> {
        let baked = bake_gltf(&source.model_bytes, &source.buffers, &source.animations)?;
        Ok(source.replace_asset(baked))
    }
}

#[derive(Default, TypePath)]
struct BakedFighterAnimationsSaver;
impl AssetSaver for BakedFighterAnimationsSaver {
    type Asset = BakedFighterAnimations;
    type Settings = ();
    type OutputLoader = BakedFighterAnimationsLoader;
    type Error = io::Error;
    async fn save(
        &self,
        writer: &mut Writer,
        asset: SavedAsset<'_, '_, Self::Asset>,
        _: &Self::Settings,
        _: AssetPath<'_>,
    ) -> Result<(), Self::Error> {
        let bytes = bincode::serialize(&*asset)
            .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
        writer.write_all(&bytes).await
    }
}
type FighterBakeProcessor = LoadTransformAndSave<
    FighterBakeSourceLoader,
    FighterBakeTransformer,
    BakedFighterAnimationsSaver,
>;

pub struct BakedAnimationPlugin;
impl Plugin for BakedAnimationPlugin {
    fn build(&self, app: &mut App) {
        app.init_asset::<BakedFighterAnimations>()
            .register_asset_loader(FighterBakeSourceLoader)
            .register_asset_loader(BakedFighterAnimationsLoader)
            .register_asset_processor(FighterBakeProcessor::new(
                FighterBakeTransformer,
                BakedFighterAnimationsSaver,
            ))
            .set_default_asset_processor::<FighterBakeProcessor>("fighterbake")
            .add_systems(
                bevy_ggrs::prelude::GgrsSchedule,
                super::runtime::update_fighter_bone_matrices
                    .in_set(crate::schedule::GameplaySet::Animation),
            )
            .add_systems(
                Update,
                super::visual_binding::apply_baked_fighter_visual_pose,
            )
            .add_observer(super::visual_binding::bind_fighter_visual_bones);
    }
}

fn bake_gltf(
    bytes: &[u8],
    buffer_bytes: &[Vec<u8>],
    animations: &HashMap<AnimKind, String>,
) -> io::Result<BakedFighterAnimations> {
    let gltf = gltf::Gltf::from_slice(bytes)
        .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error))?;
    let buffers = gltf
        .buffers()
        .zip(buffer_bytes)
        .map(|(buffer, bytes)| {
            if bytes.len() < buffer.length() {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!(
                        "glTF buffer {} is too short: expected {}, got {}",
                        buffer.index(),
                        buffer.length(),
                        bytes.len()
                    ),
                ));
            }
            Ok(gltf::buffer::Data(bytes.clone()))
        })
        .collect::<io::Result<Vec<_>>>()?;
    if buffers.len() != gltf.buffers().len() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "glTF buffer data is incomplete",
        ));
    }
    bake_gltf_document(gltf.document, buffers, animations)
}

pub(super) fn bake_gltf_document(
    document: gltf::Document,
    buffers: Vec<gltf::buffer::Data>,
    animations: &HashMap<AnimKind, String>,
) -> io::Result<BakedFighterAnimations> {
    let nodes: Vec<_> = document.nodes().collect();
    if nodes.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "fighter model has no nodes",
        ));
    }
    let mut parents = vec![None; nodes.len()];
    for node in &nodes {
        for child in node.children() {
            parents[child.index()] = Some(node.index());
        }
    }
    let names = nodes
        .iter()
        .map(|node| {
            node.name().map(str::to_owned).ok_or_else(|| {
                io::Error::new(
                    io::ErrorKind::InvalidData,
                    format!("node {} has no name", node.index()),
                )
            })
        })
        .collect::<io::Result<Vec<_>>>()?;
    let mut pending = HashMap::<&str, Vec<AnimKind>>::new();
    for (kind, name) in animations {
        pending.entry(name).or_default().push(*kind);
    }
    let mut clips = Vec::with_capacity(pending.len());
    let mut clip_indices = HashMap::with_capacity(animations.len());
    for animation in document.animations() {
        let name = animation.name().ok_or_else(|| {
            io::Error::new(
                io::ErrorKind::InvalidData,
                format!("animation {} has no name", animation.index()),
            )
        })?;
        let Some(kinds) = pending.remove(name) else {
            continue;
        };
        let channels = animation
            .channels()
            .map(|channel| read_channel(channel, &buffers))
            .collect::<io::Result<Vec<_>>>()?;
        let duration = channels
            .iter()
            .filter_map(Channel::duration)
            .fold(0.0_f32, f32::max);
        let frame_count = baked_frame_count(duration, BAKED_ANIMATION_FPS);
        let mut frames = Vec::with_capacity(frame_count);
        for frame in 0..frame_count {
            let time = (frame as f32 / BAKED_ANIMATION_FPS as f32).min(duration);
            frames.push(
                nodes
                    .iter()
                    .map(|node| local_transform(node.clone(), &channels, time))
                    .collect::<io::Result<Vec<_>>>()?,
            );
        }
        let index = clips.len();
        clips.push(BakedAnimationClip { frames });
        for kind in kinds {
            clip_indices.insert(kind, index);
        }
    }
    if clips.is_empty() {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "fighter model has no animations",
        ));
    }
    if !pending.is_empty() {
        let mut missing = pending.keys().copied().collect::<Vec<_>>();
        missing.sort_unstable();
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "fighter model has no animations named {}",
                missing.join(", ")
            ),
        ));
    }
    Ok(BakedFighterAnimations {
        frame_rate: BAKED_ANIMATION_FPS,
        bone_names: names,
        parents,
        clips,
        clip_indices,
    })
}

pub(super) fn baked_frame_count(duration: f32, frame_rate: u32) -> usize {
    let frame = duration as f64 * frame_rate as f64;
    let nearest_frame = frame.round();
    let last_frame = if (frame - nearest_frame).abs() < 0.0001 {
        nearest_frame
    } else {
        frame.ceil()
    };
    last_frame as usize + 1
}

#[derive(Clone)]
pub(super) struct Channel {
    pub(super) node: usize,
    pub(super) property: gltf::animation::Property,
    pub(super) interpolation: gltf::animation::Interpolation,
    pub(super) times: Vec<f32>,
    pub(super) values: ChannelValues,
}
impl Channel {
    fn duration(&self) -> Option<f32> {
        self.times.last().copied()
    }
}
#[derive(Clone)]
pub(super) enum ChannelValues {
    Vec3(Vec<[f32; 3]>),
    Rotation(Vec<[f32; 4]>),
}

fn read_channel(
    channel: gltf::animation::Channel<'_>,
    buffers: &[gltf::buffer::Data],
) -> io::Result<Channel> {
    let reader = channel.reader(|buffer| Some(&buffers[buffer.index()].0));
    let times = reader
        .read_inputs()
        .ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidData, "animation channel has no input")
        })?
        .collect();
    let values = match reader.read_outputs().ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            "animation channel has no output",
        )
    })? {
        gltf::animation::util::ReadOutputs::Translations(values)
        | gltf::animation::util::ReadOutputs::Scales(values) => {
            ChannelValues::Vec3(values.collect())
        }
        gltf::animation::util::ReadOutputs::Rotations(values) => {
            ChannelValues::Rotation(values.into_f32().collect())
        }
        gltf::animation::util::ReadOutputs::MorphTargetWeights(_) => {
            return Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "morph target animation is not supported for bone baking",
            ));
        }
    };
    if channel.sampler().interpolation() == gltf::animation::Interpolation::CubicSpline {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "cubic spline animation is not supported for bone baking",
        ));
    }
    Ok(Channel {
        node: channel.target().node().index(),
        property: channel.target().property(),
        interpolation: channel.sampler().interpolation(),
        times,
        values,
    })
}

fn local_transform(node: gltf::Node<'_>, channels: &[Channel], time: f32) -> io::Result<FixedMat4> {
    let (mut translation, mut rotation, mut scale) = match node.transform() {
        gltf::scene::Transform::Decomposed {
            translation,
            rotation,
            scale,
        } => (
            Vec3::from(translation),
            Quat::from_array(rotation),
            Vec3::from(scale),
        ),
        gltf::scene::Transform::Matrix { matrix } => {
            return FixedMat4::from_mat4(Mat4::from_cols_array_2d(&matrix));
        }
    };
    for channel in channels
        .iter()
        .filter(|channel| channel.node == node.index())
    {
        match (&channel.values, channel.property) {
            (ChannelValues::Vec3(values), gltf::animation::Property::Translation) => {
                translation = Vec3::from(sample_vec3(channel, values, time))
            }
            (ChannelValues::Vec3(values), gltf::animation::Property::Scale) => {
                scale = Vec3::from(sample_vec3(channel, values, time))
            }
            (ChannelValues::Rotation(values), gltf::animation::Property::Rotation) => {
                rotation = Quat::from_array(sample_rotation(channel, values, time))
            }
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "animation property does not match its values",
                ));
            }
        }
    }
    FixedMat4::from_mat4(Mat4::from_scale_rotation_translation(
        scale,
        rotation,
        translation,
    ))
}
fn sample_index(channel: &Channel, time: f32) -> (usize, usize, f32) {
    let next = channel.times.partition_point(|key| *key <= time);
    if next == 0 {
        return (0, 0, 0.0);
    }
    if next == channel.times.len() {
        let last = next - 1;
        return (last, last, 0.0);
    }
    let previous = next - 1;
    let factor = (time - channel.times[previous]) / (channel.times[next] - channel.times[previous]);
    (previous, next, factor)
}
pub(super) fn sample_vec3(channel: &Channel, values: &[[f32; 3]], time: f32) -> [f32; 3] {
    let (from, to, factor) = sample_index(channel, time);
    if channel.interpolation == gltf::animation::Interpolation::Step {
        return values[from];
    }
    Vec3::from(values[from])
        .lerp(Vec3::from(values[to]), factor)
        .to_array()
}
fn sample_rotation(channel: &Channel, values: &[[f32; 4]], time: f32) -> [f32; 4] {
    let (from, to, factor) = sample_index(channel, time);
    if channel.interpolation == gltf::animation::Interpolation::Step {
        return values[from];
    }
    Quat::from_array(values[from])
        .slerp(Quat::from_array(values[to]), factor)
        .to_array()
}