// compute.glsl - 简单的向量加法
#version 460 core

layout (local_size_x = 256) in;  // 每个工作组 256 个线程

layout (std430, binding = 0) buffer InputA {
    float dataA[];
};

layout (std430, binding = 1) buffer InputB {
    float dataB[];
};

layout (std430, binding = 2) buffer Output {
    float dataOut[];
};

void main() {
    uint idx = gl_GlobalInvocationID.x;
    if (idx < dataA.length()) {
        dataOut[idx] = dataA[idx] + dataB[idx];
    }
}

// 加载计算着色器
GLuint computeShader = glCreateShader(GL_COMPUTE_SHADER);
std::string source = readFile("compute.glsl");
const char* src = source.c_str();
glShaderSource(computeShader, 1, &src, NULL);
glCompileShader(computeShader);

GLuint program = glCreateProgram();
glAttachShader(program, computeShader);
glLinkProgram(program);

// 创建 SSBO
const int N = 1024 * 1024;  // 100 万个元素
std::vector<float> a(N, 1.0f), b(N, 2.0f);

GLuint ssboA, ssboB, ssboOut;
glGenBuffers(1, &ssboA);
glGenBuffers(1, &ssboB);
glGenBuffers(1, &ssboOut);

glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssboA);
glBufferData(GL_SHADER_STORAGE_BUFFER, N * sizeof(float), a.data(), GL_STATIC_DRAW);

glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssboB);
glBufferData(GL_SHADER_STORAGE_BUFFER, N * sizeof(float), b.data(), GL_STATIC_DRAW);

glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssboOut);
glBufferData(GL_SHADER_STORAGE_BUFFER, N * sizeof(float), NULL, GL_STATIC_DRAW);

// 绑定到绑定点
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 0, ssboA);
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 1, ssboB);
glBindBufferBase(GL_SHADER_STORAGE_BUFFER, 2, ssboOut);

// 分派计算
glUseProgram(program);
glDispatchCompute(N / 256, 1, 1);  // N/256 个工作组

// 等待 GPU 完成
glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT);

// 读回结果
glBindBuffer(GL_SHADER_STORAGE_BUFFER, ssboOut);
float* result = (float*)glMapBuffer(GL_SHADER_STORAGE_BUFFER, GL_READ_ONLY);
// result[i] = 3.0f for all i
glUnmapBuffer(GL_SHADER_STORAGE_BUFFER);

// 一维
layout (local_size_x = 256) in;

// 二维
layout (local_size_x = 16, local_size_y = 16) in;

// 三维
layout (local_size_x = 4, local_size_y = 4, local_size_z = 4) in;

// 查询限制
int maxWorkGroupSize[3];
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 0, &maxWorkGroupSize[0]); // X
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 1, &maxWorkGroupSize[1]); // Y
glGetIntegeri_v(GL_MAX_COMPUTE_WORK_GROUP_SIZE, 2, &maxWorkGroupSize[2]); // Z

int maxTotalSize;
glGetIntegerv(GL_MAX_COMPUTE_WORK_GROUP_INVOCATIONS, &maxTotalSize); // 通常 1024

// parallel_reduce.glsl - 计算数组总和
#version 460 core

layout (local_size_x = 256) in;

layout (std430, binding = 0) buffer InputBuffer {
    float inputData[];
};

layout (std430, binding = 1) buffer OutputBuffer {
    float outputData[];  // 每个工作组输出一个和
};

shared float sharedData[256];  // 共享内存

void main() {
    uint tid = gl_LocalInvocationID.x;
    uint gid = gl_GlobalInvocationID.x;

    // 加载数据到共享内存
    sharedData[tid] = inputData[gid];
    barrier();  // 同步：确保所有线程加载完毕

    // 归约：树形求和
    for (uint stride = gl_WorkGroupSize.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sharedData[tid] += sharedData[tid + stride];
        }
        barrier();  // 每步同步
    }

    // 线程 0 写出结果
    if (tid == 0) {
        outputData[gl_WorkGroupID.x] = sharedData[0];
    }
}

barrier();  // 工作组内同步屏障

// gaussian_blur.glsl
#version 460 core

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

layout (rgba8, binding = 0) uniform image2D inputImage;
layout (rgba8, binding = 1) uniform image2D outputImage;

// 5×5 高斯核
const float kernel[25] = float[](
    1.0/256,  4.0/256,  6.0/256,  4.0/256, 1.0/256,
    4.0/256, 16.0/256, 24.0/256, 16.0/256, 4.0/256,
    6.0/256, 24.0/256, 36.0/256, 24.0/256, 6.0/256,
    4.0/256, 16.0/256, 24.0/256, 16.0/256, 4.0/256,
    1.0/256,  4.0/256,  6.0/256,  4.0/256, 1.0/256
);

void main() {
    ivec2 pixel = ivec2(gl_GlobalInvocationID.xy);
    ivec2 size = imageSize(inputImage);

    if (pixel.x >= size.x || pixel.y >= size.y) return;

    vec4 sum = vec4(0.0);
    for (int y = -2; y <= 2; y++) {
        for (int x = -2; x <= 2; x++) {
            ivec2 samplePos = clamp(pixel + ivec2(x, y), ivec2(0), size - 1);
            sum += imageLoad(inputImage, samplePos) * kernel[(y+2)*5 + (x+2)];
        }
    }

    imageStore(outputImage, pixel, sum);
}

// histogram.glsl
#version 460 core

layout (local_size_x = 256) in;

layout (rgba8, binding = 0) uniform image2D inputImage;

layout (std430, binding = 0) buffer HistogramBuffer {
    uint bins[256];
};

shared uint localBins[256];  // 工作组局部直方图

void main() {
    uint tid = gl_LocalInvocationID.x;

    // 初始化局部直方图
    localBins[tid] = 0;
    barrier();

    // 每个线程处理多个像素
    ivec2 size = imageSize(inputImage);
    uint totalPixels = uint(size.x * size.y);
    uint threadsPerGroup = gl_WorkGroupSize.x;

    for (uint i = tid; i < totalPixels; i += threadsPerGroup) {
        ivec2 pos = ivec2(i % uint(size.x), i / uint(size.x));
        vec4 color = imageLoad(inputImage, pos);
        uint luminance = uint(dot(color.rgb, vec3(0.299, 0.587, 0.114)) * 255.0);
        atomicAdd(localBins[luminance], 1u);
    }
    barrier();

    // 累加到全局直方图
    atomicAdd(bins[tid], localBins[tid]);
}

// 将 OpenGL 纹理绑定为计算着色器的图像
glBindImageTexture(0, inputTexture, 0, GL_FALSE, 0, GL_READ_ONLY, GL_RGBA8);
glBindImageTexture(1, outputTexture, 0, GL_FALSE, 0, GL_WRITE_ONLY, GL_RGBA8);

// 分派
glUseProgram(blurProgram);
glDispatchCompute(width / 16, height / 16, 1);
glMemoryBarrier(GL_SHADER_IMAGE_ACCESS_BARRIER_BIT);

// 之后可以直接用 outputTexture 进行渲染

// ❌ 差：不同线程走不同分支
if (gl_LocalInvocationID.x % 2 == 0) {
    // 偶数线程走这里
} else {
    // 奇数线程走这里 → GPU 需要两路都执行
}

// ✅ 好：数据预排序，所有线程走同一路径

// ❌ 差：步长访问
float val = data[gl_GlobalInvocationID.x * stride];

// ✅ 好：连续访问
float val = data[gl_GlobalInvocationID.x];