mirage-engine 0.1.1

use bytemuck::{Pod, Zeroable};

use crate::math::{Mat4, Quat, Vec3};
use crate::mesh::Mixed;

/// The joints a model is posed by, and what each of its vertices takes of
/// them.
///
/// A mesh that is not loaded from a model has neither.
#[derive(Clone, Debug, Default, PartialEq)]
pub(crate) struct Rig {
    joints: Vec<Joint>,
    /// One per vertex of the mesh, in the order the vertex buffer takes them.
    weights: Vec<Weighted>,
}

impl Rig {
    /// The rig of `joints`, which `weights` takes one entry per vertex of.
    pub(crate) fn new(joints: Vec<Joint>, weights: Vec<Weighted>) -> Self {
        Self { joints, weights }
    }

    /// The joints, a joint's parent always before it.
    pub(crate) fn joints(&self) -> &[Joint] {
        &self.joints
    }

    /// What each vertex takes of the joints.
    pub(crate) fn weights(&self) -> &[Weighted] {
        &self.weights
    }

    /// Whether the mesh is skinned: it has joints, and one entry per vertex
    /// of what that vertex takes of them.
    ///
    /// A mesh that is skinned draws through the skinned stages and holds the
    /// second vertex stream they read; one that is not holds neither.
    pub(crate) fn skins(&self) -> bool {
        !self.joints.is_empty() && !self.weights.is_empty()
    }

    /// Memory the rig holds: its joints and its weights.
    pub(crate) fn bytes(&self) -> usize {
        size_of_val(self.joints()) + size_of_val(self.weights())
    }
}

/// One joint of a model: where it is placed, its own rest transform, and the
/// transform into its own space.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct Joint {
    pub(crate) placed: Placed,
    pub(crate) rest: Local,
    /// What the source states for a joint its skin names, and the inverse of
    /// the joint's own rest transform for one it does not.
    pub(crate) bind: Mat4,
}

/// Where a joint is placed: under another joint of the model, or, where no
/// joint is above it, within what the nodes above it multiply to.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) enum Placed {
    Under(u32),
    Within(Mat4),
}

/// One joint's own transform within what it is placed in: where it is, how
/// it is turned, and its size.
///
/// Kept as the three parts a clip moves one at a time, never as one matrix:
/// a clip that moves one of them leaves the other two where they were.
#[derive(Clone, Copy, Debug, PartialEq)]
pub(crate) struct Local {
    pub(crate) position: Vec3,
    pub(crate) turn: Quat,
    pub(crate) scale: Vec3,
}

impl Local {
    /// The transform at `position`, turned by `turn`, at `scale`.
    pub(crate) fn new(position: Vec3, turn: Quat, scale: Vec3) -> Self {
        Self {
            position,
            turn,
            scale,
        }
    }

    /// The three parts as one matrix: scaled, then turned, then moved.
    pub(crate) fn matrix(self) -> Mat4 {
        Mat4::from_scale_rotation_translation(self.scale, self.turn, self.position)
    }

    /// This transform `amount` of the way to `other`, each of the three
    /// parts mixed as a key of its own is.
    pub(crate) fn mixed(self, other: Self, amount: f32) -> Self {
        Self {
            position: self.position.mixed(other.position, amount),
            turn: self.turn.mixed(other.turn, amount),
            scale: self.scale.mixed(other.scale, amount),
        }
    }
}

/// The joints one vertex takes and how much of each: the model's own joint
/// numbers, and weights that sum to one.
///
/// One stream of these is drawn beside the vertices of a skinned mesh,
/// which is what `forward.wgsl` reads at its last two lanes.
#[repr(C)]
#[derive(Clone, Copy, Debug, Default, PartialEq, Pod, Zeroable)]
pub(crate) struct Weighted {
    pub(crate) joints: [u32; 4],
    pub(crate) weights: [f32; 4],
}

impl Weighted {
    /// `position` moved by the joints this vertex takes, each at the weight
    /// it takes of them, in the matrices `palette` holds for its model.
    ///
    /// The same blend the vertex stage takes, which is what the sphere over
    /// every pose a clip reaches is measured with. Every joint a vertex
    /// takes is a joint of its own model — a source stating one past the rig
    /// is a decode error — so `palette` holds a matrix for each of them.
    pub(crate) fn skinned(self, position: Vec3, palette: &[Mat4]) -> Vec3 {
        let blend = self
            .joints
            .into_iter()
            .zip(self.weights)
            .fold(Mat4::ZERO, |blend, (joint, weight)| {
                blend + palette[joint as usize] * weight
            });

        blend.transform_point3(position)
    }

    /// A vertex that takes `joint` alone, which is what a part of a model no
    /// skin covers takes.
    pub(crate) fn whole(joint: u32) -> Self {
        Self {
            joints: [joint, 0, 0, 0],
            weights: [1.0, 0.0, 0.0, 0.0],
        }
    }

    /// A vertex that takes `joints` by `weights`, scaled to sum to one; a
    /// vertex no joint takes at all takes its first joint whole.
    pub(crate) fn scaled(joints: [u32; 4], weights: [f32; 4]) -> Self {
        let mut taken = Self { joints, weights };
        let sum: f32 = taken.weights.iter().sum();
        taken.weights = match sum > 0.0 {
            true => taken.weights.map(|weight| weight / sum),
            false => [1.0, 0.0, 0.0, 0.0],
        };
        taken
    }
}