metaverse_avatar 0.3.3

use std::io::Error;

use crate::avatar::Avatar;
use benthic_default_asset_converter::generated::DEFAULT_SKELETON;
use benthic_protocol::skeleton::{Joint, JointName, Skeleton, Transform};
use glam::{Mat4, Vec4};
use indexmap::IndexMap;
use metaverse_messages::http::mesh::Skin;
use uuid::Uuid;

/// This function takes an object's skeleton, and applies it to the agent's combined skeleton. The
/// combined skeleton contains all of the transforms for all of the joints, ranked by how many
/// elements of the outfit contain the same transform.
/// Highest ranking joint values will always be last.
pub fn update_global_avatar_skeleton(avatar: &mut Avatar, skeleton: &Skeleton) {
    for joint in skeleton.joints.values() {
        // add the mentioned joints to the used joints set.
        avatar.used_joints.insert(joint.name);
        if let Some(global_joint) = avatar.skeleton.joints.get_mut(&joint.name) {
            for transform in &joint.global_transforms {
                set_rank(global_joint, transform, |j| &mut j.global_transforms);
            }
            for local_transform in &joint.local_transforms {
                set_rank(global_joint, local_transform, |j| &mut j.local_transforms);
            }
        }
    }
}

/// Determine the local and global joint transforms for a skinned SceneObject.
pub fn create_skeleton(object_name: String, id: Uuid, skin: &Skin) -> Result<Skeleton, Error> {
    // if the object has a mesh, handle the skeleton
    // fetch the default skeleton which is generated at compile time
    // this will be used for calculating transforms, and getting joint local transforms for the
    // root object.
    let default_skeleton = DEFAULT_SKELETON.clone();

    let mut joints = IndexMap::new();
    // for every joint in the skin
    for (i, name) in skin.joint_names.iter().enumerate() {
        // apply the rotations from the default skeleton to the object
        // the default skeleton's transforms are stored in [0]
        // these rotations need to be applied, because the IBMs from the server are mostly
        // the identity matrix, and only contain translation information. The default
        // skeleton contains the rotations.
        let default_joints = default_skeleton.joints.get(name).unwrap().clone();
        let mut default_transform = default_joints.global_transforms[0].transform;

        default_transform.w_axis = Vec4::new(0.0, 0.0, 0.0, 1.0);

        // this one generates correct global coordinates. Currently working out the bugs with this.
        let transform_matrix = Mat4::from_rotation_translation(
            default_transform.to_scale_rotation_translation().1,
            skin.inverse_bind_matrices[i].inverse().w_axis.truncate(),
        );
        let transform = Transform {
            name: object_name.clone(),
            id,
            transform: transform_matrix,
            // Default transforms are rank 0, and will stay at position 0 in the array.
            // Custom transforms start at rank 1.
            rank: 1,
        };

        // create the joint object that contains the calculted transforms
        let joint = Joint {
            name: *name,
            parent: default_joints.parent,
            children: default_joints.children.into_iter().collect(),
            global_transforms: vec![transform.clone()],
            // leave the local tranforms empty for now. They will be added after the loop.
            local_transforms: vec![],
        };
        joints.insert(*name, joint.clone());
    }

    // determine the root joints of the skeleton
    let root_joints: Vec<JointName> = joints
        .values()
        .filter(|joint| joint.parent.is_none())
        .map(|joint| joint.name)
        .collect();

    let last_transforms: IndexMap<JointName, Mat4> = joints
        .iter()
        .map(|(name, joint)| (*name, joint.global_transforms.last().unwrap().transform))
        .collect();

    for (_, joint) in joints.iter_mut() {
        if let Some(parent_name) = &joint.parent {
            // if last_transforms doesn't contain the parent, it means that is the root, and
            // the parent should be calculated via the default skeleton.
            let parent = last_transforms
                .get(parent_name)
                .or_else(|| {
                    Some(
                        &default_skeleton
                            .joints
                            .get(parent_name)
                            .unwrap()
                            .global_transforms[0]
                            .transform,
                    )
                })
                .unwrap();

            let child = last_transforms.get(&joint.name).unwrap();
            // IBM-based local matrix
            let local_matrix = parent.inverse() * child;
            let local_transform = Transform {
                name: object_name.clone(),
                id,
                transform: local_matrix,
                rank: 1,
            };
            joint.local_transforms.push(local_transform);
        } else {
            let child = last_transforms.get(&joint.name).unwrap();

            let local_transform = Transform {
                name: object_name.clone(),
                id,
                transform: *child,
                rank: 1,
            };

            joint.local_transforms.push(local_transform);
        }
    }

    Ok(Skeleton {
        root: root_joints,
        joints,
    })
}

// Modify the rank for local and global transforms.
fn set_rank<F>(joint: &mut Joint, transform: &Transform, mut get_vec: F)
where
    F: FnMut(&mut Joint) -> &mut Vec<Transform>,
{
    let transforms = get_vec(joint);

    // Track whether we matched only rank 0
    let mut matched_rank0 = false;

    for t in transforms.iter_mut().skip(1) {
        if t.transform.abs_diff_eq(transform.transform, 1e-4) {
            t.rank += 1;
            transforms.sort_by_key(|a| a.rank);
            return;
        }
    }

    // If we didn’t find a non-zero match, check rank 0.
    if let Some(first) = transforms.first()
        && first.transform.abs_diff_eq(transform.transform, 1e-4)
    {
        matched_rank0 = true;
    }

    // If it matched only rank 0 or none at all, create a new transform
    if matched_rank0
        || !transforms
            .iter()
            .any(|t| t.transform.abs_diff_eq(transform.transform, 1e-4))
    {
        let mut new_t = transform.clone();
        // If it matched only rank 0, start at rank 1
        if matched_rank0 {
            new_t.rank = 1;
        }
        transforms.push(new_t);
        transforms.sort_by_key(|a| a.rank);
    }
}