#version 430
#extension GL_ARB_gpu_shader_fp64 : enable
layout (local_size_x=50, local_size_y=1, local_size_z=1) in;

layout(std430, binding = 0) buffer PositionBuffer {
    dvec2 X[];
};

layout(std430, binding = 1) buffer VelocityBuffer {
    dvec2 V[];
};

layout(std430, binding = 2) buffer MassBuffer {
    double M[];
};

layout(std430, binding = 3) buffer LengthBuffer {
    double L[];
};

layout(std430, binding = 4) buffer IndexBuffer {
    uint indices[];
};

layout(std430, binding = 5) buffer SegmentCountBuffer {
    uint segmentCounts[];
};

uniform double h;
uniform double gravity;
uniform uint substeps;

shared dvec2 P[50];

void solveConstraint(uint segmentId, uint gIndex){
    double w1 = segmentId == 0 ? 0.0 : 1.0 / M[gIndex - 1];
    double w2 = 1.0 / M[gIndex];

    dvec2 s = (segmentId == 0 ? dvec2(0, 0) : P[segmentId - 1]) - P[segmentId];
    dvec2 n = normalize(s);
    double l = length(s);

    dvec2 deltaP1 = n * -w1 / (w1 + w2) * (l - L[gIndex]);
    dvec2 deltaP2 = n *  w2 / (w1 + w2) * (l - L[gIndex]);

    if (segmentId > 0)
        P[segmentId - 1] += deltaP1;
    P[segmentId] += deltaP2;
}

void main() {
    uint pendulumId = gl_WorkGroupID.x;
    uint i_off = indices[pendulumId];
    uint count = segmentCounts[pendulumId];
    uint localId = gl_LocalInvocationID.x;
    uint index = i_off + localId;

    if (localId >= count)
        return;

    for (int k = 0; k < substeps; k++){

        // explicit integration with gravity as external force
        V[index] = V[index] + dvec2(0, -gravity * h);
        P[localId] = X[index] + V[index] * h;

        memoryBarrierShared();

        if (localId % 2 == 0){
            solveConstraint(localId, i_off + localId);
            if (localId + 1 < count){
                memoryBarrierShared();
                solveConstraint(localId + 1, i_off + localId + 1);
            }
        }

        barrier();

        // update position and velocity
        V[index] = (P[localId] - X[index]) / h;
        X[index] = P[localId];

        memoryBarrierBuffer();
    }
}