#version 460 core

// Compute shader for classifying pixels as either:
// edge or internal

layout(local_size_x = 16, local_size_y = 16) in;

layout(rgba16f, binding = 0) uniform restrict readonly image2D g_normal;

layout(r8ui, binding = 1) uniform restrict writeonly uimage2D classification;

const uint INTERNAL = 0;
const uint EDGE = 1;

// approx cos(10°)
const float THRESH = 0.984;

void main() {
    ivec2 img_coord = ivec2(gl_GlobalInvocationID.xy);
    ivec2 img_size = imageSize(g_normal);

    vec3 normals[9] = vec3[9](
        vec3(0, 0, 0), vec3(0, 0, 0), vec3(0, 0, 0),
        vec3(0, 0, 0), vec3(0, 0, 0), vec3(0, 0, 0),
        vec3(0, 0, 0), vec3(0, 0, 0), vec3(0, 0, 0)
    );
    int i = 0;
    for (int dy = -1; dy <= 1; ++dy) {
        int y = img_coord.y + dy;
        if (y < 0 || y == img_size.y)
            continue;
        for (int dx = -1; dx <= 1; ++dx) {
            int x = img_coord.x + dx;
            if (x < 0 || x == img_size.x)
                continue;
            normals[i] = imageLoad(g_normal, ivec2(x, y)).xyz;
        }
    }

    uint classif = INTERNAL;
    for (int i = 0; i < 9; ++i) {
        float angle = abs(dot(normals[i], normals[4]));
        if (angle >= THRESH)
            classif = EDGE;
    }

    imageStore(classification, img_coord, uvec4(classif));
}