import opengl, pixie, shady, shady/compute

initOffscreenWindow()

# Setup the uniforms.
var
  inputCommandBuffer*: Uniform[SamplerBuffer]
  patternBuffer*: Uniform[Sampler2D]
  screenBuffer*: Uniform[Sampler2D]
  outputImageBuffer*: UniformWriteOnly[UImageBuffer]
  dimensions*: Uniform[IVec4] # ivec4(width, height, 0, 0)
  patternDimensions*: Uniform[IVec4] # ivec4(width, height, 0, 0)
  screenDimensions*: Uniform[IVec4] # ivec4(width, height, 0, 0)

# The shader itself.
proc commandsToImage() =
  var pos = gl_GlobalInvocationID
  for x in 0 ..< dimensions.x:
    pos.x = x.uint32
    let value = uint32(texelFetch(inputCommandBuffer, int32(pos.x mod 256)).x)
    #echo pos.x, " ", value
    let colorValue = uvec4(
      128,
      0,
      value,
      255
    )
    imageStore(outputImageBuffer, int32(pos.y * uint32(dimensions.x) + pos.x), colorValue)

proc patternDither() =
  var
    pos = gl_GlobalInvocationId
    maskPos = ivec2(pos.x.int32 mod patternDimensions.x, pos.y.int32 mod patternDimensions.y)

  #echo screenDimensions.x

  let
    mask = patternBuffer.texelFetch(maskPos, 0)
    input = screenBuffer.texelFetch(ivec2(pos.x.int32, pos.y.int32), 0)
    colorValue =
      if input.r + input.g + input.b > mask.r + mask.g + mask.b: uvec4(255, 255, 255, 255)
      else: uvec4(0, 0, 0, 255)
  imageStore(outputImageBuffer, int32(pos.y * uint32(screenDimensions.x) + pos.x), colorValue)

# Setup the input and output data.
patternBuffer.image = readImage("pattern.png")
screenBuffer.image = readImage("screen.png")
outputImageBuffer.image = newImage(screenBuffer.image.width, screenBuffer.image.height)
patternDimensions = ivec4(
  patternBuffer.image.width.int32,
  patternBuffer.image.height.int32,
  0,
  0
)
screenDimensions = ivec4(
  screenBuffer.image.width.int32,
  screenBuffer.image.height.int32,
  0,
  0
)
#for i in 0 ..< 256:
#  inputCommandBuffer.data.add(i.float32)
#outputImageBuffer.image = newImage(256*2, 256*2)
#dimensions = ivec4(
#  outputImageBuffer.image.width.int32,
#  outputImageBuffer.image.height.int32,
#  0,
#  0
#)

template runComputeOnGpu(computeShader: proc(), invocationSize: UVec3) =

  let computeShaderSrc = toGLSL(
    computeShader,
    "430",
    extra = "layout (local_size_x = 1, local_size_y = 1, local_size_z = 1) in;\n"
  )

  writeFile("examples/compute1_sh.comp", computeShaderSrc)

  var shaderId = compileComputeShader((
    "examples/compute1_sh.comp",
    computeShaderSrc
  ))
  glUseProgram(shaderId)

  # Setup outputImageBuffer.
  var
    outputBufferId: GLuint
    outputTextureId: GLuint
  glGenBuffers(1, outputBufferId.addr)
  glBindBuffer(GL_TEXTURE_BUFFER, outputBufferId)
  glBufferData(GL_TEXTURE_BUFFER, outputImageBuffer.image.data.len * 4, nil, GL_STATIC_DRAW)
  let outputImageBufferLoc = glGetUniformLocation(shaderId, "outputImageBuffer")
  glUniform1i(outputImageBufferLoc, 0)
  glActiveTexture(GL_TEXTURE0)
  glGenTextures(1, outputTextureId.addr)
  glBindTexture(GL_TEXTURE_BUFFER, outputTextureId)
  glTexBuffer(GL_TEXTURE_BUFFER, GL_RGBA8UI, outputBufferId)
  glBindImageTexture(
    0,
    outputTextureId,
    0,
    GL_FALSE,
    0,
    GL_WRITE_ONLY,
    GL_RGBA8UI
  )
  glBindTexture(GL_TEXTURE_BUFFER, outputTextureId)

  # Setup patternBuffer.
  var
    patternBufferId: GLuint
    patternTextureId: GLuint
  glGenTextures(1, patternTextureId.addr)
  glActiveTexture(GL_TEXTURE1)
  glGenBuffers(1, patternBufferId.addr)
  #let byteLength = patternBuffer.image.width * patternBuffer.image.height * 4
  glBindBuffer(GL_TEXTURE_BUFFER, patternBufferId)
  glBufferData(
    GL_TEXTURE_BUFFER,
    patternBuffer.image.data.len * 4,
    patternBuffer.image.data[0].addr,
    GL_STATIC_DRAW
  )
  glBindTexture(GL_TEXTURE_BUFFER, patternTextureId)
  glTexBuffer(
    GL_TEXTURE_BUFFER,
    #GL_R32F,
    GL_RGBA8UI,
    patternBufferId
  )
  glBindTexture(GL_TEXTURE_BUFFER, patternTextureId)
  let patternBufferLoc = glGetUniformLocation(shaderId, "patternBuffer")
  glUniform1i(patternBufferLoc, 1)

  # Setup screenBuffer.
  var
    screenBufferId: GLuint
    screenTextureId: GLuint
  glGenTextures(1, screenTextureId.addr)
  glActiveTexture(GL_TEXTURE2)
  glGenBuffers(1, screenBufferId.addr)
  glBindBuffer(GL_TEXTURE_BUFFER, screenBufferId)
  glBufferData(
    GL_TEXTURE_BUFFER,
    screenBuffer.image.data.len * 4,
    screenBuffer.image.data[0].addr,
    GL_STATIC_DRAW
  )
  glBindTexture(GL_TEXTURE_BUFFER, screenTextureId)
  glTexBuffer(
    GL_TEXTURE_BUFFER,
    #GL_R32F,
    GL_RGBA8UI,
    screenBufferId
  )
  glBindTexture(GL_TEXTURE_BUFFER, screenTextureId)
  let screenBufferLoc = glGetUniformLocation(shaderId, "screenBuffer")
  glUniform1i(screenBufferLoc, 2)

  # Setup patternDimensions uniform.
  let patternDimensionsLoc = glGetUniformLocation(shaderId, "patternDimensions")
  glUniform4i(patternDimensionsLoc, patternDimensions.x, patternDimensions.y, patternDimensions.z, patternDimensions.w)
  let screenDimensionsLoc = glGetUniformLocation(shaderId, "screenDimensions")
  glUniform4i(screenDimensionsLoc, screenDimensions.x, screenDimensions.y, screenDimensions.z, screenDimensions.w)

  # Run the shader.
  glDispatchCompute(
    invocationSize.x.GLuint,
    invocationSize.y.GLuint,
    invocationSize.z.GLuint
  )
  glMemoryBarrier(GL_SHADER_STORAGE_BARRIER_BIT)

  # Read back the outputImageBuffer.
  let p = cast[ptr UncheckedArray[uint8]](
    glMapNamedBuffer(outputBufferId, GL_READ_ONLY)
  )
  copyMem(
    outputImageBuffer.image.data[0].addr,
    p,
    outputImageBuffer.image.data.len * 4
  )
  discard glUnmapNamedBuffer(outputBufferId)

# Run it on CPU.
runComputeOnCpu(patternDither, uvec3(outputImageBuffer.image.width.uint32, outputImageBuffer.image.height.uint32, 1))
outputImageBuffer.image.writeFile("examples/compute1_output_cpu.png")

# Just in case clear the image before running GPU.
outputImageBuffer.image.fill(rgbx(0, 0, 0, 0))

# Run it on the GPU.
runComputeOnGpu(patternDither, uvec3(outputImageBuffer.image.width.uint32, outputImageBuffer.image.height.uint32, 1))
outputImageBuffer.image.writeFile("examples/compute1_output_gpu.png")