# Simple compute shader example.
#
# An OpenGL 4.3 context, a texture for the compute shader to write
# and the fragment shader to read, and two program objects are created.
# One object is for the compute shader and the other is for rendering
# (vertex + fragment shaders).
# After that we go into a loop where we update a counter in the compute shader,
# fill in the texture (as image2D), and blit the texture onto the screen.
#
# Author: Ville Timonen (http://wili.cc/blog/opengl-cs.html) 2012
#
# Modified and extended for Tcl3D by Paul Obermeier 2018/07/30
# See www.tcl3d.org for the Tcl3D extension.

package require Tk
package require tcl3d

# Font to be used in the Tk listbox.
set g_Demo(listFont) {-family {Courier} -size 10}

# Window size.
set g_Demo(winWidth)  512
set g_Demo(winHeight) 512

set g_Demo(appName) "Tcl3D demo: Compute Shader Demo"

# Show errors occuring in the Togl callbacks.
proc bgerror { msg } {
    tk_messageBox -icon error -type ok -message "Error: $msg\n\n$::errorInfo"
    exit
}

# Print info message into widget at the bottom of the window.
proc PrintInfo { msg } {
    if { [winfo exists .fr.info] } {
        .fr.info configure -text $msg
    }
}

proc NextFrame {} {
    .fr.toglwin postredisplay
}

proc Animate {} {
    global g_Demo

    .fr.toglwin postredisplay
    set g_Demo(animateId) [tcl3dAfterIdle Animate]
}

proc StartAnimation {} {
    global g_Demo

    if { ! [info exists g_Demo(animateId)] } {
        Animate
    }
}

proc StopAnimation {} {
    global g_Demo

    if { [info exists g_Demo(animateId)] } {
        after cancel $g_Demo(animateId)
        unset g_Demo(animateId)
    }
}

proc CheckGLErrors { msg } {
    set errMsg [tcl3dOglGetError]
    if { $errMsg eq "" } {
        return
    }
    puts "$msg: $errMsg"
}

proc GenTexture {} {
    # We create a single float channel 512^2 texture
    set texHandle [tcl3dVector GLuint 1]
    glGenTextures 1 $texHandle

    glActiveTexture GL_TEXTURE0
    glBindTexture GL_TEXTURE_2D [$texHandle get 0]
    glTexParameteri GL_TEXTURE_2D GL_TEXTURE_MIN_FILTER $::GL_LINEAR
    glTexParameteri GL_TEXTURE_2D GL_TEXTURE_MAG_FILTER $::GL_LINEAR
    glTexImage2D GL_TEXTURE_2D 0 $::GL_R32F 512 512 0 GL_RED GL_FLOAT "NULL"

    # Because we're also using this tex as an image (in order to write to it),
    # we bind it to an image unit as well
    glBindImageTexture 0 [$texHandle get 0] 0 GL_FALSE 0 GL_WRITE_ONLY $::GL_R32F
    CheckGLErrors "Gen texture"
}

proc GenRenderProg {} {
    set progHandle [glCreateProgram]
    set vp [glCreateShader GL_VERTEX_SHADER]
    set fp [glCreateShader GL_FRAGMENT_SHADER]

    set vpSrc "
        #version 430
        in vec2 pos;
        out vec2 texCoord;
        void main() {
            texCoord = pos*0.5f + 0.5f;
            gl_Position = vec4(pos.x, pos.y, 0.0, 1.0);
        }
    "

    set fpSrc "
        #version 430
        uniform sampler2D srcTex;
        in vec2 texCoord;
        out vec4 color;
        void main() {
            float c = texture(srcTex, texCoord).x;
            color = vec4(c, 1.0, 1.0, 1.0);
        }
    "

    tcl3dOglShaderSource $vp $vpSrc
    tcl3dOglShaderSource $fp $fpSrc

    glCompileShader $vp
    set rvalue [tcl3dOglGetShaderState $vp GL_COMPILE_STATUS]
    if { ! $rvalue } {
        puts stderr "Error in compiling vp"
        exit 30
    }
    glAttachShader $progHandle $vp

    glCompileShader $fp
    set rvalue [tcl3dOglGetShaderState $fp GL_COMPILE_STATUS]
    if { ! $rvalue } {
        puts stderr "Error in compiling fp"
        exit 31
    }
    glAttachShader $progHandle $fp

    glBindFragDataLocation $progHandle 0 "color"
    glLinkProgram $progHandle

    set rvalue [tcl3dOglGetProgramState $progHandle GL_LINK_STATUS]
    if { ! $rvalue } {
        puts stderr "Error in linking sp"
        exit 32
    }
    
    glUseProgram $progHandle
    glUniform1i [glGetUniformLocation $progHandle "srcTex"] 0

    set vertArray [tcl3dVector GLuint 1]
    glGenVertexArrays 1 $vertArray
    glBindVertexArray [$vertArray get 0]

    set posBuf [tcl3dVector GLuint 1]
    glGenBuffers 1 $posBuf
    glBindBuffer GL_ARRAY_BUFFER [$posBuf get 0]

    set dataList [list \
        -1.0 -1.0 \
        -1.0  1.0 \
         1.0 -1.0 \
         1.0  1.0 \
    ]
    set dataVec [tcl3dVectorFromList GLfloat $dataList]
    set dataVecSize [expr [llength $dataList] * [$dataVec elemsize]]

    glBufferData GL_ARRAY_BUFFER $dataVecSize $dataVec GL_STREAM_DRAW

    set posPtr [glGetAttribLocation $progHandle "pos"]
    glVertexAttribPointer $posPtr 2 GL_FLOAT GL_FALSE 0 "NULL"
    glEnableVertexAttribArray $posPtr

    CheckGLErrors "Render shaders"
    return $progHandle
}

proc GenComputeProg {} {
    # Creating the compute shader, and the program object containing the shader
    set progHandle [glCreateProgram]
    set cs [glCreateShader GL_COMPUTE_SHADER]

    # In order to write to a texture, we have to introduce it as image2D.
    # local_size_x/y/z layout variables define the work group size.
    # gl_GlobalInvocationID is a uvec3 variable giving the global ID of the thread,
    # gl_LocalInvocationID is the local index within the work group, and
    # gl_WorkGroupID is the work group's index
    set csSrc "
        #version 430
        uniform float roll;
        layout (binding=0, r32f) uniform image2D destTex;
        layout (local_size_x = 16, local_size_y = 16) in;
        void main() {
            ivec2 storePos = ivec2(gl_GlobalInvocationID.xy);
            float localCoef = length(vec2(ivec2(gl_LocalInvocationID.xy)-8)/8.0);
            float globalCoef = sin(float(gl_WorkGroupID.x+gl_WorkGroupID.y)*0.1 + roll)*0.5;
            imageStore(destTex, storePos, vec4(1.0-globalCoef*localCoef, 0.0, 0.0, 0.0));
        }
    "

    tcl3dOglShaderSource $cs $csSrc
    glCompileShader $cs
    set rvalue [tcl3dOglGetShaderState $cs GL_COMPILE_STATUS]
    if { ! $rvalue } {
        puts stderr "Error in compiling cs"
        puts "Compiler log: [tcl3dOglGetShaderInfoLog $cs]"
        exit 40
    }
    glAttachShader $progHandle $cs

    glLinkProgram $progHandle
    set rvalue [tcl3dOglGetProgramState $progHandle GL_LINK_STATUS]
    if { ! $rvalue } {
        puts stderr "Error in linking cs"
        puts "Compiler log: [tcl3dOglGetShaderInfoLog $cs]"
        exit 42
    }

    glUseProgram $progHandle
    
    glUniform1i [glGetUniformLocation $progHandle "destTex"] 0

    CheckGLErrors "Compute shader"
    return $progHandle
}

proc Init {} {
    global g_Demo

    GenTexture
    set g_Demo(renderHandle)  [GenRenderProg]
    set g_Demo(computeHandle) [GenComputeProg]
    set g_Demo(curFrame) 0
}

proc CreateCallback { toglwin } {
    global g_Demo

}

proc ReshapeCallback { toglwin { w -1 } { h -1 } } {
    set w [$toglwin width]
    set h [$toglwin height]
    glViewport 0 0 $w $h
}

proc DisplayCallback { toglwin } {
    global g_Demo

    glClear [expr $::GL_COLOR_BUFFER_BIT | $::GL_DEPTH_BUFFER_BIT]

    # Update texture
    glUseProgram $g_Demo(computeHandle)
    glUniform1f [glGetUniformLocation $g_Demo(computeHandle) "roll"] [expr $g_Demo(curFrame)*0.01]
    # 512^2 threads in blocks of 16^2
    glDispatchCompute [expr 512/16] [expr 512/16] 1
    CheckGLErrors "Dispatch compute shader"

    # Draw screen
    glUseProgram $g_Demo(renderHandle)
    glDrawArrays GL_TRIANGLE_STRIP 0 4
    CheckGLErrors "Draw screen"

    $toglwin swapbuffer
    incr g_Demo(curFrame)
}

proc Cleanup {} {
    global g_Demo

    if { [info exists g_Program] } {
        tcl3dOglDestroyProgram $g_Program
    }

    foreach var [info globals g_*] {
        uplevel #0 unset $var
    }
}

# Put all exit related code here.
proc ExitProg {} {
    exit
}

# Create the OpenGL window and some Tk helper widgets.
proc CreateWindow {} {
    global g_Demo

    frame .fr
    pack .fr -expand 1 -fill both

    set retVal [catch { togl .fr.toglwin \
            -width $g_Demo(winWidth) -height $g_Demo(winHeight) \
            -double true -depth true \
            -createcommand  CreateCallback \
            -coreprofile true -major 4 -minor 3 } errMsg]

    if { $retVal != 0 } {
        tk_messageBox -icon error -type ok -title "Missing Togl feature" \
                      -message "Demo needs core profile: $errMsg"
        proc ::Cleanup {} {}
        exit 1
        return
    }

    Init

    .fr.toglwin configure \
        -reshapecommand ReshapeCallback \
        -displaycommand DisplayCallback

    listbox .fr.usage -font $g_Demo(listFont) -height 3
    label .fr.status
    label .fr.info
    grid .fr.toglwin -row 0 -column 0 -sticky news
    grid .fr.usage   -row 1 -column 0 -sticky news
    grid .fr.info    -row 2 -column 0 -sticky news
    grid rowconfigure .fr 0 -weight 1
    grid columnconfigure .fr 0 -weight 1
    wm title . $g_Demo(appName)

    # Watch for Esc key and Quit messages
    wm protocol . WM_DELETE_WINDOW "ExitProg"
    bind . <Key-n>      "NextFrame"
    bind . <Key-Escape> "ExitProg"

    bind .fr.toglwin <1> "StartAnimation"
    bind .fr.toglwin <2> "StopAnimation"
    bind .fr.toglwin <3> "StopAnimation"
    bind .fr.toglwin <Control-Button-1> "StopAnimation" 

    .fr.usage insert end "Key-Escape Exit"
    .fr.usage insert end "Key-n      Next frame"
    .fr.usage insert end "Mouse-L|MR Start|Stop animation"
}

CreateWindow
PrintInfo [tcl3dOglGetInfoString]

if { [file tail [info script]] eq [file tail $::argv0] } {
    # If started directly from tclsh or wish, then start animation.
    update
    StartAnimation
}

