#version 430 core
layout(local_size_x = 16, local_size_y = 16) in;

// state_curr: read-only; state_next: write-only
layout(r8ui,  binding = 0) readonly  uniform uimage2D state_curr;
layout(r8ui,  binding = 1) writeonly uniform uimage2D state_next;

// claim buffer to resolve conflicts
layout(r32ui, binding = 2) coherent uniform uimage2D claim;

uniform ivec2 u_gridSize;
uniform uint  u_frame;

const uint AIR  = 0u;
const uint SAND = 1u;

// How far down we scan to approximate column thickness
const int MAX_SCAN = 6;

// How much “taller” this column must be vs neighbor before we slide
// Smaller -> shallower, more flowy piles
const int SLOPE_THRESHOLD = 2;

uint hash(uvec2 p, uint seed) {
    p = p * 1664525u + 1013904223u + seed;
    p.x ^= p.y >> 16;
    p.y ^= p.x >> 16;
    return p.x ^ p.y;
}

// Approximate thickness of the sand column under (x, y)
int columnThickness(int x, int y) {
    int t = 0;
    for (int dy = 0; dy < MAX_SCAN; ++dy) {
        int yy = y + dy;
        if (yy >= u_gridSize.y) break;
        uint v = imageLoad(state_curr, ivec2(x, yy)).r;
        if (v == SAND) {
            t++;
        } else {
            break;
        }
    }
    return t;
}

void main() {
    ivec2 pos = ivec2(gl_GlobalInvocationID.xy);
    if (pos.x >= u_gridSize.x || pos.y >= u_gridSize.y) return;

    uint self = imageLoad(state_curr, pos).r;
    if (self != SAND) return;

    // Lock bottom row in place
    if (pos.y == u_gridSize.y - 1) {
        imageStore(state_next, pos, uvec4(SAND, 0u, 0u, 0u));
        return;
    }

    ivec2 dest = pos;

    // Only relax grains that are resting on something
    ivec2 down = pos + ivec2(0, 1);
    uint below = imageLoad(state_curr, down).r;
    if (below == SAND) {
        int hSelf = columnThickness(pos.x, pos.y);

        bool canLeft  = false;
        bool canRight = false;

        ivec2 left  = pos + ivec2(-1, 0);
        ivec2 right = pos + ivec2( 1, 0);

        int hLeft  = 0;
        int hRight = 0;

        // Check left side
        if (left.x >= 0) {
            uint leftCell = imageLoad(state_curr, left).r;
            if (leftCell == AIR) {
                hLeft = columnThickness(left.x, pos.y);
                int slopeL = hSelf - hLeft;
                if (slopeL > SLOPE_THRESHOLD) {
                    canLeft = true;
                }
            }
        }

        // Check right side
        if (right.x < u_gridSize.x) {
            uint rightCell = imageLoad(state_curr, right).r;
            if (rightCell == AIR) {
                hRight = columnThickness(right.x, pos.y);
                int slopeR = hSelf - hRight;
                if (slopeR > SLOPE_THRESHOLD) {
                    canRight = true;
                }
            }
        }

        if (canLeft || canRight) {
            if (canLeft && canRight) {
                // Prefer the *lower* column; randomize if equal
                if (hLeft < hRight) {
                    dest = left;
                } else if (hRight < hLeft) {
                    dest = right;
                } else {
                    uint h = hash(uvec2(pos), u_frame);
                    bool chooseLeft = (h & 1u) != 0u;
                    dest = chooseLeft ? left : right;
                }
            } else if (canLeft) {
                dest = left;
            } else { // canRight
                dest = right;
            }
        }
    }

    bool staying = (dest.x == pos.x && dest.y == pos.y);

    if (staying) {
        imageStore(state_next, pos, uvec4(SAND, 0u, 0u, 0u));
    } else {
        // Moving to new cell; resolve conflicts via claim buffer
        uint old = imageAtomicCompSwap(claim, dest, 0u, 1u);
        if (old == 0u) {
            imageStore(state_next, dest, uvec4(SAND, 0u, 0u, 0u));
        } else {
            // lost: stay put
            imageStore(state_next, pos, uvec4(SAND, 0u, 0u, 0u));
        }
    }
}