import pyglet

from pyglet.gl import (

GL_RGBA32F, GL_FLOAT, GL_TEXTURE_2D, GL_READ_WRITE, GL_PIXEL_PACK_BUFFER, GLuint, GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0,

glTexParameteri, GL_TEXTURE_MIN_FILTER, GL_TEXTURE_MAG_FILTER, GL_RGBA, glReadPixels, glBindFramebuffer, glGenFramebuffers,

glFramebufferTexture2D, glGenTextures, glTexImage2D, glBindTexture, GL_LINEAR, glUseProgram, glActiveTexture, GL_TEXTURE0, GL_RGBA32UI, GL_UNSIGNED_INT, GL_RED,

GL_RGBA32I, GL_INT, GL_R32UI, GL_FALSE, GL_READ_ONLY, glBindImageTexture

)

import numpy as np

import ctypes

from pyglet import gl

global window

window = pyglet.window.Window(700, 700)

# Compute shader source code

compute_src = """#version 430 core

layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in;

layout(r32ui, binding=0) uniform uimage2D img_input;

layout(r32ui, binding=1) uniform uimage2D img_output;

void main(void)

{

uvec2 texel_coord = uvec2(gl_GlobalInvocationID.xy);

uvec4 loadvalue = imageLoad(img_input, ivec2(texel_coord));

imageStore(img_output, ivec2(texel_coord), uvec4(loadvalue.r + 12, 0, 0, 0));

}

"""

# Create and initialize program

program = pyglet.graphics.shader.ComputeShaderProgram(compute_src)

width, height = 2, 2

# Initialize input data

dataq = np.zeros((height, width), dtype=np.uint32)

dataq = np.array([[0, 1],

[2, 3]])

dataq = dataq.astype(np.uint32)

# Print initialized input data

print("Input data:")

print(dataq)

# Create input texture

texture = GLuint(0)

glGenTextures(1, ctypes.byref(texture))

glBindTexture(GL_TEXTURE_2D, texture)

glTexImage2D(GL_TEXTURE_2D, 0, GL_R32UI, width, height, 0, GL_RED, GL_UNSIGNED_INT, dataq.ctypes.data)

glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)

glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)

texture.bind_image_texture(0)

#glBindTexture(GL_TEXTURE_2D, 0)

glBindImageTexture(0, texture, 0, GL_FALSE, 0, GL_READ_ONLY, GL_R32UI)

# Output texture

out_texture = pyglet.image.Texture.create(width, height, GL_TEXTURE_2D, internalformat=GL_R32UI)

out_texture.bind_image_texture(unit=1) # Bind output texture to unit 1 (image binding in the shader)

# Create framebuffer and attach output texture

framebuffer = GLuint(0)

glGenFramebuffers(1, ctypes.byref(framebuffer))

glBindFramebuffer(GL_FRAMEBUFFER, framebuffer)

glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, out_texture.id, 0)

# Bind the input texture to texture unit 0 (as expected by the shader)

glActiveTexture(GL_TEXTURE0)

glBindTexture(GL_TEXTURE_2D, texture)

# Dispatch compute shader with the correct number of work groups (based on texture size)

with program:

program.dispatch(width, height, 1)

pyglet.gl.glFinish()

data = (ctypes.cuint32 * (width * height))()

glReadPixels(0, 0, width, height, GL_RED, GL_UNSIGNED_INT, data)

np_data = np.array(data).reshape((height, width))

print("Output data:")

print(np_data)

glBindFramebuffer(GL_FRAMEBUFFER, 0)