#version 430
#extension GL_ARB_gpu_shader_fp64 : enable

#define MAX_SEGMENTS 32

layout (local_size_x=MAX_SEGMENTS, 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 InverseMassBuffer {
    double W[];
};

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[MAX_SEGMENTS];

// distance constraint
void solveConstraint(uint segmentId, uint gIndex){
    double w1 = segmentId == 0 ? 0.0 : W[gIndex - 1];
    double w2 = W[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;
}

// No need for synchronization, because each pendulum gets its own warp
void main() {
    uint pendulumId = gl_WorkGroupID.x;
    uint i_off = indices[pendulumId];
    uint segmentCount = segmentCounts[pendulumId];
    uint localId = gl_LocalInvocationID.x;
    uint gIndex = i_off + localId;

    // discard unneeded local threads for segments
    if (localId >= segmentCount)
        return;

    // substepping is not parallelizable
    for (int k = 0; k < substeps; k++){

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

        // every segment depends on its neighbour, so we can not parallelize all constraints
        // only parallelize those, who are independent of each other, skipping every second segment
        if (localId % 2 == 0){
            solveConstraint(localId, i_off + localId);
            if (localId + 1 < segmentCount)
                solveConstraint(localId + 1, gIndex + 1);
        }

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