TiedUp-/src/main/java/com/tiedup/remake/client/gltf/GltfSkinningEngine.java

package com.tiedup.remake.client.gltf;

import net.minecraftforge.api.distmarker.Dist;
import net.minecraftforge.api.distmarker.OnlyIn;
import org.joml.Matrix4f;
import org.joml.Quaternionf;
import org.joml.Vector3f;
import org.joml.Vector4f;

/**
 * CPU-based Linear Blend Skinning (LBS) engine.
 * Computes joint matrices purely from glTF data (rest translations + animation rotations).
 * All data is in MC-converted space (consistent with IBMs and vertex positions).
 */
@OnlyIn(Dist.CLIENT)
public final class GltfSkinningEngine {

    private GltfSkinningEngine() {}

    /**
     * Compute joint matrices from glTF animation/rest data (default animation).
     * Each joint matrix = worldTransform * inverseBindMatrix.
     * Uses MC-converted glTF data throughout for consistency.
     *
     * @param data parsed glTF data (MC-converted)
     * @return array of joint matrices ready for skinning
     */
    public static Matrix4f[] computeJointMatrices(GltfData data) {
        return computeJointMatricesFromClip(data, data.animation());
    }

    /**
     * Compute joint matrices with frame interpolation for animated entities.
     * Uses SLERP for rotations and LERP for translations between adjacent keyframes.
     *
     * <p>The {@code time} parameter is in frame-space: 0.0 corresponds to the first
     * keyframe and {@code frameCount - 1} to the last. Values between integer frames
     * are interpolated. Out-of-range values are clamped.</p>
     *
     * @param data the parsed glTF data (MC-converted)
     * @param clip the animation clip to sample (null = rest pose for all joints)
     * @param time time in frame-space (0.0 = first frame, N-1 = last frame)
     * @return interpolated joint matrices ready for skinning
     */
    public static Matrix4f[] computeJointMatricesAnimated(
        GltfData data, GltfData.AnimationClip clip, float time
    ) {
        int jointCount = data.jointCount();
        Matrix4f[] jointMatrices = new Matrix4f[jointCount];
        Matrix4f[] worldTransforms = new Matrix4f[jointCount];

        int[] parents = data.parentJointIndices();

        for (int j = 0; j < jointCount; j++) {
            // Build local transform: translate(interpT) * rotate(interpQ)
            Matrix4f local = new Matrix4f();
            local.translate(getInterpolatedTranslation(data, clip, j, time));
            local.rotate(getInterpolatedRotation(data, clip, j, time));

            // Compose with parent
            if (parents[j] >= 0 && worldTransforms[parents[j]] != null) {
                worldTransforms[j] = new Matrix4f(worldTransforms[parents[j]]).mul(local);
            } else {
                worldTransforms[j] = new Matrix4f(local);
            }

            // Final joint matrix = worldTransform * inverseBindMatrix
            jointMatrices[j] = new Matrix4f(worldTransforms[j])
                .mul(data.inverseBindMatrices()[j]);
        }

        return jointMatrices;
    }

    /**
     * Internal: compute joint matrices from a specific animation clip.
     */
    private static Matrix4f[] computeJointMatricesFromClip(GltfData data, GltfData.AnimationClip clip) {
        int jointCount = data.jointCount();
        Matrix4f[] jointMatrices = new Matrix4f[jointCount];
        Matrix4f[] worldTransforms = new Matrix4f[jointCount];

        int[] parents = data.parentJointIndices();

        for (int j = 0; j < jointCount; j++) {
            // Build local transform: translate(animT or restT) * rotate(animQ or restQ)
            Matrix4f local = new Matrix4f();
            local.translate(getAnimTranslation(data, clip, j));
            local.rotate(getAnimRotation(data, clip, j));

            // Compose with parent
            if (parents[j] >= 0 && worldTransforms[parents[j]] != null) {
                worldTransforms[j] = new Matrix4f(worldTransforms[parents[j]]).mul(local);
            } else {
                worldTransforms[j] = new Matrix4f(local);
            }

            // Final joint matrix = worldTransform * inverseBindMatrix
            jointMatrices[j] = new Matrix4f(worldTransforms[j])
                .mul(data.inverseBindMatrices()[j]);
        }

        return jointMatrices;
    }

    /**
     * Get the animation rotation for a joint (MC-converted).
     * Falls back to rest rotation if no animation.
     */
    private static Quaternionf getAnimRotation(GltfData data, GltfData.AnimationClip clip, int jointIndex) {
        if (clip != null && jointIndex < clip.rotations().length
                && clip.rotations()[jointIndex] != null) {
            return clip.rotations()[jointIndex][0]; // first frame
        }
        return data.restRotations()[jointIndex];
    }

    /**
     * Get the animation translation for a joint (MC-converted).
     * Falls back to rest translation if no animation translation exists.
     */
    private static Vector3f getAnimTranslation(GltfData data, GltfData.AnimationClip clip, int jointIndex) {
        if (clip != null && clip.translations() != null
                && jointIndex < clip.translations().length
                && clip.translations()[jointIndex] != null) {
            return clip.translations()[jointIndex][0]; // first frame
        }
        return data.restTranslations()[jointIndex];
    }

    // ---- Interpolated accessors (for computeJointMatricesAnimated) ----

    /**
     * Get an interpolated rotation for a joint at a fractional frame time.
     * Uses SLERP between the two bounding keyframes.
     *
     * <p>Falls back to rest rotation when the clip is null or has no rotation
     * data for the given joint. A single-frame channel returns that frame directly.</p>
     *
     * @param data       parsed glTF data
     * @param clip       animation clip (may be null)
     * @param jointIndex joint to query
     * @param time       frame-space time (clamped internally)
     * @return new Quaternionf with the interpolated rotation (never mutates source data)
     */
    private static Quaternionf getInterpolatedRotation(
        GltfData data, GltfData.AnimationClip clip, int jointIndex, float time
    ) {
        if (clip == null || jointIndex >= clip.rotations().length
                || clip.rotations()[jointIndex] == null) {
            // No animation data for this joint -- use rest pose (copy to avoid mutation)
            Quaternionf rest = data.restRotations()[jointIndex];
            return new Quaternionf(rest);
        }

        Quaternionf[] frames = clip.rotations()[jointIndex];
        if (frames.length == 1) {
            return new Quaternionf(frames[0]);
        }

        // Clamp time to valid range [0, frameCount-1]
        float clamped = Math.max(0.0f, Math.min(time, frames.length - 1));
        int f0 = (int) Math.floor(clamped);
        int f1 = Math.min(f0 + 1, frames.length - 1);
        float alpha = clamped - f0;

        if (alpha < 1e-6f || f0 == f1) {
            return new Quaternionf(frames[f0]);
        }

        // SLERP: create a copy of frame0 and slerp toward frame1
        return new Quaternionf(frames[f0]).slerp(frames[f1], alpha);
    }

    /**
     * Get an interpolated translation for a joint at a fractional frame time.
     * Uses LERP between the two bounding keyframes.
     *
     * <p>Falls back to rest translation when the clip is null, the clip has no
     * translation data at all, or has no translation data for the given joint.
     * A single-frame channel returns that frame directly.</p>
     *
     * @param data       parsed glTF data
     * @param clip       animation clip (may be null)
     * @param jointIndex joint to query
     * @param time       frame-space time (clamped internally)
     * @return new Vector3f with the interpolated translation (never mutates source data)
     */
    private static Vector3f getInterpolatedTranslation(
        GltfData data, GltfData.AnimationClip clip, int jointIndex, float time
    ) {
        if (clip == null || clip.translations() == null
                || jointIndex >= clip.translations().length
                || clip.translations()[jointIndex] == null) {
            // No animation data for this joint -- use rest pose (copy to avoid mutation)
            Vector3f rest = data.restTranslations()[jointIndex];
            return new Vector3f(rest);
        }

        Vector3f[] frames = clip.translations()[jointIndex];
        if (frames.length == 1) {
            return new Vector3f(frames[0]);
        }

        // Clamp time to valid range [0, frameCount-1]
        float clamped = Math.max(0.0f, Math.min(time, frames.length - 1));
        int f0 = (int) Math.floor(clamped);
        int f1 = Math.min(f0 + 1, frames.length - 1);
        float alpha = clamped - f0;

        if (alpha < 1e-6f || f0 == f1) {
            return new Vector3f(frames[f0]);
        }

        // LERP: create a copy of frame0 and lerp toward frame1
        return new Vector3f(frames[f0]).lerp(frames[f1], alpha);
    }

    /**
     * Skin a single vertex using Linear Blend Skinning.
     *
     * <p>Callers should pre-allocate {@code tmpPos} and {@code tmpNorm} and reuse
     * them across all vertices in a mesh to avoid per-vertex allocations (12k+
     * allocations per frame for a typical mesh).</p>
     *
     * @param data         parsed glTF data
     * @param vertexIdx    index into the vertex arrays
     * @param jointMatrices joint matrices from computeJointMatrices
     * @param outPos       output skinned position (3 floats)
     * @param outNormal    output skinned normal (3 floats)
     * @param tmpPos       pre-allocated scratch Vector4f for position transforms
     * @param tmpNorm      pre-allocated scratch Vector4f for normal transforms
     */
    public static void skinVertex(
        GltfData data, int vertexIdx, Matrix4f[] jointMatrices,
        float[] outPos, float[] outNormal,
        Vector4f tmpPos, Vector4f tmpNorm
    ) {
        float[] positions = data.positions();
        float[] normals = data.normals();
        int[] joints = data.joints();
        float[] weights = data.weights();

        // Rest position
        float vx = positions[vertexIdx * 3];
        float vy = positions[vertexIdx * 3 + 1];
        float vz = positions[vertexIdx * 3 + 2];

        // Rest normal
        float nx = normals[vertexIdx * 3];
        float ny = normals[vertexIdx * 3 + 1];
        float nz = normals[vertexIdx * 3 + 2];

        // LBS: v_skinned = Σ(w[i] * jointMatrix[j[i]] * v_rest)
        float sx = 0, sy = 0, sz = 0;
        float snx = 0, sny = 0, snz = 0;

        for (int i = 0; i < 4; i++) {
            int ji = joints[vertexIdx * 4 + i];
            float w = weights[vertexIdx * 4 + i];
            if (w <= 0.0f || ji >= jointMatrices.length) continue;

            Matrix4f jm = jointMatrices[ji];

            // Transform position
            tmpPos.set(vx, vy, vz, 1.0f);
            jm.transform(tmpPos);
            sx += w * tmpPos.x;
            sy += w * tmpPos.y;
            sz += w * tmpPos.z;

            // Transform normal (ignore translation)
            tmpNorm.set(nx, ny, nz, 0.0f);
            jm.transform(tmpNorm);
            snx += w * tmpNorm.x;
            sny += w * tmpNorm.y;
            snz += w * tmpNorm.z;
        }

        outPos[0] = sx;
        outPos[1] = sy;
        outPos[2] = sz;

        // Normalize the normal
        float len = (float) Math.sqrt(snx * snx + sny * sny + snz * snz);
        if (len > 0.0001f) {
            outNormal[0] = snx / len;
            outNormal[1] = sny / len;
            outNormal[2] = snz / len;
        } else {
            outNormal[0] = 0;
            outNormal[1] = 1;
            outNormal[2] = 0;
        }
    }
}