#include "./../../../GPUInstancer/Resources/Compute/Include/PlatformDefines.compute"

#pragma kernel CSAnimateBones
#pragma kernel CSAnimateBonesLerped
#pragma kernel CSBonesFixWeights

uniform RWStructuredBuffer<float4> gpuiAnimationData; // index: 0 x -> frameNo1, y -> frameNo2, z -> frameNo3, w -> frameNo4
                                                      // index: 1 x -> weight1, y -> weight2, z -> weight3, w -> weight4
uniform RWStructuredBuffer<float4x4> gpuiAnimationBuffer;
uniform Texture2D<float4> gpuiAnimationTexture;
uniform uint animationTextureSizeX;
uniform uint totalNumberOfFrames;
uniform uint totalNumberOfBones;
uniform uint instanceCount;
uniform float deltaTime;
uniform uint frameRate;

inline float4x4 getBoneMatrixFromTexture(uint frameIndex, uint boneIndex)
{
    // ordered by frames
    uint textureIndex = frameIndex * totalNumberOfBones + boneIndex;

    uint2 texIndex;
    texIndex.x = (textureIndex % animationTextureSizeX);
    texIndex.y = (floor(textureIndex / (float) animationTextureSizeX) * 4);
    
    float4x4 boneMatrix;
    boneMatrix._11_21_31_41 = gpuiAnimationTexture[texIndex];
    texIndex.y += 1;
    boneMatrix._12_22_32_42 = gpuiAnimationTexture[texIndex];
    texIndex.y += 1;
    boneMatrix._13_23_33_43 = gpuiAnimationTexture[texIndex];
    texIndex.y += 1;
    boneMatrix._14_24_34_44 = gpuiAnimationTexture[texIndex];
    return boneMatrix;
}

inline float4x4 CalculateBoneMatrix(uint instanceIndex, uint boneIndex)
{
    float4 animationData = gpuiAnimationData[instanceIndex];
    float4 weightData = gpuiAnimationData[instanceIndex + 1];
    
    float frameNo = animationData[0];
    float weight = weightData[0];
    // required for some shader compilers like PSSL
    float4x4 weightMatrix = float4x4(
    weight, 0, 0, 0,
    0, weight, 0, 0,
    0, 0, weight, 0,
    0, 0, 0, weight
    );
    uint frameStart = floor(frameNo);
    uint frameEnd = ceil(frameNo);
    float progress = frac(frameNo);

    float4x4 bone = mul(lerp(getBoneMatrixFromTexture(frameStart, boneIndex), getBoneMatrixFromTexture(frameEnd, boneIndex), progress), weightMatrix);

    // blending
    if (weight < 1)
    {
        for (uint i = 1; i < 4; i++)
        {
            weight = weightData[i];
            if (weight > 0)
            {
                weightMatrix._11 = weight;
                weightMatrix._22 = weight;
                weightMatrix._33 = weight;
                weightMatrix._44 = weight;

                frameNo = animationData[i];
                frameStart = floor(frameNo);
                frameEnd = ceil(frameNo);
                progress = frac(frameNo);
                
                bone += mul(lerp(getBoneMatrixFromTexture(frameStart, boneIndex), getBoneMatrixFromTexture(frameEnd, boneIndex), progress), weightMatrix);
            }
        }
    }

    return bone;
}

[numthreads(GPUI_THREADS_2D, GPUI_THREADS_2D, 1)]
void CSAnimateBones(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= instanceCount || id.y >= totalNumberOfBones)
        return;
    
    // lerp from previous state disabled because it is causing problem on initialization
    gpuiAnimationBuffer[id.y + totalNumberOfBones * id.x] = CalculateBoneMatrix(id.x * 2, id.y);
}

[numthreads(GPUI_THREADS_2D, GPUI_THREADS_2D, 1)]
void CSAnimateBonesLerped(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= instanceCount || id.y >= totalNumberOfBones)
        return;
    
    // smooth pass from previous bone data
    gpuiAnimationBuffer[id.y + totalNumberOfBones * id.x] = lerp(gpuiAnimationBuffer[id.y + totalNumberOfBones * id.x], CalculateBoneMatrix(id.x * 2, id.y), clamp(deltaTime * frameRate, 0.0001, 1));
}


[numthreads(GPUI_THREADS, 1, 1)]
void CSBonesFixWeights(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= instanceCount)
        return;
    
    uint weightIndex = id.x * 2 + 1;
    
    float4 weightData = gpuiAnimationData[weightIndex];
    float totalWeight = weightData.x;
    totalWeight += weightData.y;
    totalWeight += weightData.z;
    totalWeight += weightData.w;
    
    if (totalWeight <= 0)
    {
        weightData[0] = 1;
    }
    else
    {
        totalWeight = 1 / totalWeight;

        weightData.x *= totalWeight;
        weightData.y *= totalWeight;
        weightData.z *= totalWeight;
        weightData.w *= totalWeight;
    }

    gpuiAnimationData[weightIndex] = weightData;
}