/******************************************************************/
/* Program for a simple, interactive 2D Harr Wavelet Transform    */
/*     using an OpenGL shaderfor the transform and inverse        */
/*     transform.                                                 */
/*                                                                */
/* Key intermediate stages are shown and labeled to better help   */
/*     understand what is going on during this process.           */
/*                                                                */
/* Also demonstrates a simple use of an IGLUWidgetWindow to open  */
/*     a window containing FLTK widgets that control variables    */
/*     that modify rendering behavior.                            */
/*                                                                */
/* Chris Wyman (02/23/2012)                                       */
/******************************************************************/

#include <GL/glew.h>
#include <stdio.h>

// Include IGLU and set it up to avoid the iglu:: prefix on types
#include "iglu.h"
using namespace iglu;

// IGLU data structures this program uses
IGLUTexture::Ptr        inputImage  = 0;         // The input texture to play with
IGLUShaderProgram::Ptr  harr        = 0;         // The Harr transform shader
IGLUShaderProgram::Ptr  invharr     = 0;         // The inverse Harr transform
IGLUShaderProgram::Ptr  truncate    = 0;         // A shader to truncate wavelet coefs
IGLUVertexArray::Ptr    vertArray   = 0;         // Geometry for a quad
IGLUFramebuffer::Ptr    pingPong[2] = { 0, 0 };  // Buffers storing intermediate data

// IGLU variables that can be changed via a FLTK-window user-interface.
//    Param #1 = var value, Param #2 = var range, Param #3 = incr/decr delta, Param #4 = var name
//    NOTE:  IGLUBool don't take range or delta parameters (since range and delta is fixed).
IGLUInt                 numPasses  ( 5,    IGLURange<int>  (1, 8),   1,     "Pyramid levels:" );
IGLUFloat               threshold  ( 0.05, IGLURange<float>(0, 0.4), 0.001, "Truncation threshold:" );
IGLUBool                displayFPS ( true,                                  "Display framerate" );

// These are handles to the two windows, should you need to modify their
//   behavior after window creation.
IGLUWindow::Ptr         myWin = 0;
IGLUWidgetWindow::Ptr   uiWin = 0;



// We do lots of transform passes (4 of these per wavelet level), so we 
//    encapsulate it here to make the display routine more legible.
void TransformPass( IGLUFramebuffer::Ptr &fbo, IGLUShaderProgram::Ptr &shd, 
				    const IGLUTexture &inTex, const ivec2 &offset )
{
	fbo->Bind();                  // Bind the temp framebuffer we'll draw into this step
	shd["inputTex"]  = inTex;     // Setup the texture we'll be reading from (last temp FBO)
	shd["dirOffset"] = offset;    // The offset controls the direction -- a pass in X?  Or Y?
		vertArray->DrawArrays( GL_TRIANGLE_STRIP, 0, 4 );  // Draw a quad to invoke shader
	fbo->Unbind();                // Done, so unbind temp buffer (to enable use elsewhere)
}

// The display callback that is invoked by our IGLUWindow with the screen is redrawn
void display ( void )
{
	// Draw our viewport 1 (Original image) on screen
	glViewport( 0, 512, 512, 512 );
	IGLUDraw::Fullscreen( inputImage, IGLU_DRAW_FLIP_Y );

	// Compute the Harr transform with an appropriate number of passes
	harr->Enable();
	float passScale = 1.0f;
	for ( int i = 0; i < numPasses; i++, passScale*=0.5 )
	{
		harr["passScale"] = passScale;
		if (i==0) TransformPass( pingPong[0], harr, *inputImage, ivec2( 1, 0 ) );      // Pass in X
		else      TransformPass( pingPong[0], harr, pingPong[1][0], ivec2( 1, 0 ) );
		TransformPass( pingPong[1], harr, pingPong[0][0], ivec2( 0, 1 ) );             // Pass in Y
	}
	harr->Disable();

	// Draw our viewport 2 (Original image under Harr Wavelet Transform) on screen
	glViewport( 0, 0, 512, 512 );
	IGLUDraw::Fullscreen( pingPong[1][0], IGLU_DRAW_FLIP_Y );

	// Perform residual/coefficient thresholding
	pingPong[0]->Bind();
		truncate->Enable();
		truncate["inputTex"]   = pingPong[1][0];
		truncate["threshold"]  = threshold;
			vertArray->DrawArrays( GL_TRIANGLE_STRIP, 0, 4 );
		truncate->Disable();
	pingPong[0]->Unbind();

	// Draw our viewport 3 (Thresholded wavelet coefs) on screen
	glViewport( 512, 0, 512, 512 );
	IGLUDraw::Fullscreen( pingPong[0][0], IGLU_DRAW_FLIP_Y );

	// Invert the compressed/thresholded coeficients to get our
	//    "compressed" image (though we're not storing it in
	//    any sort of compressed format)
	invharr->Enable();
	passScale = 1.0f/pow(2.0,numPasses-1);
	for ( int i = 0; i < numPasses; i++, passScale*=2.0 )
	{
		invharr["passScale"] = passScale;
		TransformPass( pingPong[1], invharr, pingPong[0][0], ivec2( 0, 1 ) );  // Pass in Y
		TransformPass( pingPong[0], invharr, pingPong[1][0], ivec2( 1, 0 ) );  // Pass in X
	}
	invharr->Disable();

	// Draw our viewport 4 (Compressed image) on screen
	glViewport( 512, 512, 512, 512 );
	IGLUDraw::Fullscreen( pingPong[0][0], IGLU_DRAW_FLIP_Y );

	// Draw miscellaneous frame information.  This needs a full-screen viewport
	glViewport( 0, 0, 1024, 1024 );

	// Draw labels for images
	IGLUDraw::DrawText( IGLU_FONT_VARIABLE, 350, 993, "Original Image" );
	IGLUDraw::DrawText( IGLU_FONT_VARIABLE, 260, 481, "Harr Wavelet Transform" );
	IGLUDraw::DrawText( IGLU_FONT_VARIABLE, 755, 481, "Thesholded Wavelet Coefs" );
	IGLUDraw::DrawText( IGLU_FONT_VARIABLE, 830, 993, "Compressed Image" );
}

// Code that initializes our OpenGL state.  This is guaranteed (by IGLUWindow) to be 
//    called after the OpenGL context & all extensions have been initialized
void OpenGLInitialization( void )
{
	// Load the image to "compress" in our wavelet basis
	inputImage = new IGLUTexture2D( "../../CommonSampleFiles/images/lena512x512.ppm",  
		                            IGLU_MIN_NEAREST | IGLU_MAG_NEAREST | 
									IGLU_CLAMP_TO_EDGE_S | IGLU_CLAMP_TO_EDGE_T );	                        

	// Set up our full-screen quad geometry.
	vertArray = new IGLUVertexArray();
	float vertArrayData[] = { 0, 0, 1, 0, 0, 1, 1, 1 };
	vertArray->SetVertexArray( sizeof( vertArrayData ), vertArrayData );
	vertArray->EnableAttribute( IGLU_ATTRIB_VERTEX, 2, GL_FLOAT );
	
	// Create the framebuffers we'll ping-pong between to perform the wavelet transform
	pingPong[0] = IGLUFramebuffer::Create( GL_RGBA16F, 512, 512 ); 
	pingPong[1] = IGLUFramebuffer::Create( GL_RGBA16F, 512, 512 );
	
	// Load the shader that performs the wavelet transform
	harr         = new IGLUShaderProgram( "shaders/harrTransform.vert.glsl", 
		                                  "shaders/harrTransform.frag.glsl");
	harr["mvp"]  = IGLUMatrix4x4::Ortho( 0, 1, 0, 1 );

	// Load the shader that performs the wavelet transform
	invharr         = new IGLUShaderProgram( "shaders/harrTransformInverse.vert.glsl", 
		                                     "shaders/harrTransformInverse.frag.glsl");
	invharr["mvp"]  = IGLUMatrix4x4::Ortho( 0, 1, 0, 1 );

	// Crop out small residual coefficients
	truncate     = new IGLUShaderProgram( "shaders/cropCoefs.vert.glsl", 
		                                  "shaders/cropCoefs.frag.glsl");
	truncate["mvp"]  = IGLUMatrix4x4::Ortho( 0, 1, 0, 1 );
}

// A method that asks our main window to start/stop displaying the 
//     current framerate.  Is setup, below, as a callback so the 
//     appropriate variable automatically calls this when modified.
void ToggleFramerate( void *uiBool )
{
	myWin->DisplayFramerate( ((IGLUBool *)uiBool)->GetValue() );
}

// Our main entry point
int main(int argc, char** argv)
{
	// Put together our main window
	myWin = new IGLUWindow( 1024, 1024, "Simple IGLU Example:  GPU-Based Harr Wavelet Transform" );
	myWin->SetWindowProperties( IGLU_WINDOW_NO_RESIZE |	
								IGLU_WINDOW_DOUBLE |
								IGLU_WINDOW_REDRAW_ON_IDLE |
								IGLU_WINDOW_W_FRAMERATE ); 
	myWin->SetDisplayCallback( display );  // Setup re-display callback 

	// We don't actually do anything when idling, except redraw.  Using NullIdle()
	//    along with IGLU_WINDOW_REDRAW_ON_IDLE accomplishes this automatic redraw
	myWin->SetIdleCallback( IGLUWindow::NullIdle );

	// Give a callback to execute GL setup once a context has been created
	myWin->SetPreprocessOnGLInit( OpenGLInitialization );

	// Create / show the main window
	myWin->CreateWindow( argc, argv );     

	// Create a widget window to allow us to interact/change our IGLUVariables
	uiWin = new IGLUWidgetWindow( 300, 130, "UI Widget Window" );
	uiWin->AddWidget( &threshold );
	uiWin->AddWidget( &numPasses );
	uiWin->AddWidget( &displayFPS );
	//uiWin->CreateWindow();

	// Tell our main window about the widget window... 
	myWin->SetWidgetWindow( uiWin );

	// Since displayFPS changes the state of our window (not our display() function --
	//     which automatically gets called), we need to setup a callback to get notified
	//     when we need to toggle the window properties to start/stop displaying FPS.
	displayFPS.SetCallback( ToggleFramerate );

	// Start the FLTK interaction loop
	IGLUWindow::Run();
	return 0;
}
