# Simple compute shader example.
#
# An example of the OpenGL Red Book 9th Edition modified to work with Tcl3D.
#
# Modified and extended for Tcl3D by Paul Obermeier 2018/08/04
# 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: Red Book 9 example 12-simplecompute"

set g_Demo(localSizeX)  32
set g_Demo(localSizeY)  16
set g_Demo(globalSizeX)  8 
set g_Demo(globalSizeY) 16

set g_Demo(frameCount)  1

# 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 vglAttachShaderSource { prog type src } {
    set sh [glCreateShader $type]
    tcl3dOglShaderSource $sh $src
    glCompileShader $sh
    set infoLog [tcl3dOglGetProgramInfoLog $prog]
    if { $infoLog ne "" } {
        puts $infoLog
    }
    glAttachShader $prog $sh
    glDeleteShader $sh
}

proc Animate {} {
    global g_Demo

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

proc StartAnimation {} {
    global g_Demo

    if { [info commands .fr.toglwin] eq "" } {
        return
    }
    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 Init {} {
    global g_Demo

    # Initialize our compute program
    set g_Demo(compute_prog) [glCreateProgram]

    set compute_shader_source "
        #version 430 core
        layout (local_size_x = $g_Demo(localSizeX), local_size_y = $g_Demo(localSizeY)) in;
        layout (rgba32f) uniform image2D output_image;
        void main(void)
        {
            imageStore(output_image,
            ivec2(gl_GlobalInvocationID.xy),
            vec4(vec2(gl_LocalInvocationID.xy) / vec2(gl_WorkGroupSize.xy), 0.0, 0.0));
        }
    "

    vglAttachShaderSource $g_Demo(compute_prog) $::GL_COMPUTE_SHADER $compute_shader_source

    glLinkProgram $g_Demo(compute_prog)

    # This is the texture that the compute program will write into
    set g_Demo(output_image) [tcl3dVector GLuint 1]
    glGenTextures 1 $g_Demo(output_image)
    glBindTexture GL_TEXTURE_2D [$g_Demo(output_image) get 0]
    glTexStorage2D GL_TEXTURE_2D 8 GL_RGBA32F 256 256

    # Now create a simple program to visualize the result
    set g_Demo(render_prog) [glCreateProgram]

    set render_vs "
        #version 430 core
        in vec4 vert;
        void main(void)
        {
            gl_Position = vert;
        }
    "

    set render_fs "
        #version 430 core
        layout (location = 0) out vec4 color;
        uniform sampler2D output_image;
        void main(void)
        {
            color = texture(output_image, vec2(gl_FragCoord.xy) / vec2(textureSize(output_image, 0)));
        }
    "

    vglAttachShaderSource $g_Demo(render_prog) $::GL_VERTEX_SHADER   $render_vs
    vglAttachShaderSource $g_Demo(render_prog) $::GL_FRAGMENT_SHADER $render_fs

    glLinkProgram $g_Demo(render_prog)

    # This is the VAO containing the data to draw the quad (including its associated VBO)
    set render_vao [tcl3dVector GLuint 1]
    glGenVertexArrays 1 $render_vao
    glBindVertexArray [$render_vao get 0]

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

    set verts [list \
        -1.0 -1.0 0.5 1.0 \
         1.0 -1.0 0.5 1.0 \
         1.0  1.0 0.5 1.0 \
        -1.0  1.0 0.5 1.0 \
    ]
    set vertsVec [tcl3dVectorFromList GLfloat $verts]
    set vertsVecSize [expr [llength $verts] * [$vertsVec elemsize]]
    glBufferData GL_ARRAY_BUFFER $vertsVecSize $vertsVec GL_STATIC_DRAW
    glVertexAttribPointer 0 4 GL_FLOAT GL_FALSE 0 "NULL"
}

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

    # Activate the compute program and bind the output texture image
    glUseProgram $g_Demo(compute_prog)
    glBindImageTexture 0 [$g_Demo(output_image) get 0] 0 GL_FALSE 0 GL_WRITE_ONLY $::GL_RGBA32F
    glDispatchCompute $g_Demo(globalSizeX) $g_Demo(globalSizeY) 1

    # Now bind the texture for rendering _from_
    glBindTexture GL_TEXTURE_2D [$g_Demo(output_image) get 0]

    # Clear, select the rendering program and draw a full screen quad
    glClear [expr $::GL_COLOR_BUFFER_BIT | $::GL_DEPTH_BUFFER_BIT]
    glUseProgram $g_Demo(render_prog)
    glDrawArrays GL_TRIANGLE_FAN 0 4

    $toglwin swapbuffer
    incr g_Demo(frameCount)
}

proc Cleanup {} {
    global g_Demo

    glUseProgram 0
    if { [info exists g_Demo(compute_prog)] } {
        tcl3dOglDestroyProgram $g_Demo(compute_prog)
    }
    if { [info exists g_Demo(render_prog)] } {
        tcl3dOglDestroyProgram $g_Demo(render_prog)
    }
    glDeleteTextures 1 [$::g_Demo(output_image) get 0]
    $g_Demo(output_image) delete

    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

    # Create a Togl window using a core profile.
    # Reshape and Display callbacks are configured later after knowing if
    # the needed core profile is available.
    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 1
    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-Escape> "ExitProg"

    .fr.usage insert end "Key-Escape Exit"
}

CreateWindow

PrintInfo [tcl3dOglGetInfoString]

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