/**
 * Augmented Reality
 *
 * Assignment 2 - Camera calibration through a chessboard pattern
 *
 * @author   Cesar Palomo - cpalomo @ inf . puc-rio . br
 * @homepage http://www.inf.puc-rio.br/~cpalomo
 * @date     08/29/07
 *
 * @description Using the OpenCV library, a video of a chessboard is used for camera
 * calibration.
 * In a first step, the chessboard corners are found and identified in an output video.
 * In a further step, some camera parameters (position and focal distance) are determined. 
 * Then, using an OpenGL canvas, the chessboard is drawn along with three lines representing
 *  the 3 axis: X in red, Y in green and Z in blue.
 * Besides, some virtual objects are placed on top of the chessboard.
 *
 */


#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <conio.h>
#include <time.h>
#include "cvcam.h"
#include "cv.h"
#include "highgui.h"
#include "tsai.c"

#include "GL/gl.h"
#include "GL/glu.h"
#include "GL/glut.h"

#define DEBUG			 0
#define RET_OK			 0
#define RET_ERROR		-1

/************************************************************************
 * 
 ************************************************************************/

typedef struct {
	// OpenCV corners representation
	CvPoint2D32f* corners_openCV;
	// Corners in world coordinates
	float *Xw, *Yw;
	// Corners in camera coordinates
	float *Xf, *Yf;

} Corners;

// globals
struct glb {
	// OpenGL's use globals
	float mtsaiopengl_projection[16], mtsaiopengl_modelview[16];
	float NEARPLANE, FARPLANE;
	// OpenCV's use globals
	CvCapture* capture;
	IplImage *current_frame;
	Corners corners;
	CvSize pattern_size;
	int num_max_corners;
	int FRAME_W, FRAME_H, FPS, FRAMES;
	// focal point, translation and rotation matrices
	float f, T[3], r[9];
	// Indicates whether the camera's intrinsic parameters have been found
	bool foundInstrinsicParams;
	// Controls video's FPS
	clock_t start, end;

} _glb;


/************************************************************************
 * 
 ************************************************************************/

/**
 * Inits OpenCV's use globals.
 *
 * @param m chessboard is m X n
 * @param n chessboard is m X n
 * @param inputFile input AVI filename
 * @return RET_OK when initialization ok, RET_ERROR otherwise
 */
int InitOpenCV(int m, int n, char *inputFile) {
	// inits some globals
	_glb.pattern_size    = cvSize(m - 1, n - 1);
	_glb.num_max_corners = (m - 1) * (n - 1);
	_glb.NEARPLANE = 10.0;
	_glb.FARPLANE = 1000.0;
	_glb.foundInstrinsicParams = FALSE;

	// captures input file
	_glb.capture = cvCaptureFromAVI(inputFile); 
	
	if( !_glb.capture ) {
		fprintf(stderr,"Could not initialize capturing...\n");
		return RET_ERROR;
	}

	// reads capture's fps and number of frames properties to control video rate
	_glb.FPS		= (int) cvGetCaptureProperty(_glb.capture, CV_CAP_PROP_FPS);
	_glb.FRAMES  = (int) cvGetCaptureProperty(_glb.capture, CV_CAP_PROP_FRAME_COUNT);
	// reads capture's width and height for proper display
	_glb.FRAME_W = (int) cvGetCaptureProperty(_glb.capture, CV_CAP_PROP_FRAME_WIDTH);
	_glb.FRAME_H = (int) cvGetCaptureProperty(_glb.capture, CV_CAP_PROP_FRAME_HEIGHT);

	// inits corners
	_glb.corners.corners_openCV = (CvPoint2D32f*) calloc(_glb.num_max_corners, sizeof(CvPoint2D32f));
	_glb.corners.Xf = (float*) calloc(_glb.num_max_corners, sizeof(float));
	_glb.corners.Yf = (float*) calloc(_glb.num_max_corners, sizeof(float));
	_glb.corners.Xw = (float*) calloc(_glb.num_max_corners, sizeof(float));
	_glb.corners.Yw = (float*) calloc(_glb.num_max_corners, sizeof(float));
	int ind = 0;
	for (int h = 0; h < n - 1; h++) {
		for (int w = 0; w < m - 1; w++) {
			_glb.corners.Xw[ind] = w+1;		_glb.corners.Yw[ind] = h+1;
			ind++;
		}
	}

	return RET_OK;
}


/**
 * Releases memory from some OpenCV's use globals.
 *
 * @return RET_OK when initialization ok, RET_ERROR otherwise
 */
void EndOpenCV() {
	// releases corners struct
	if (_glb.corners.corners_openCV) {
		free(_glb.corners.corners_openCV);
		free(_glb.corners.Xf);
		free(_glb.corners.Yf);
		free(_glb.corners.Xw);
		free(_glb.corners.Yw);
	}
	// releases video capture
	if (_glb.capture)
		cvReleaseCapture( &(_glb.capture) );
}

/**
 * Sets corner points in camera coordinates
 *
 */
void SetCornersInCameraCoord() {

	int half_width  = _glb.FRAME_W * 0.5;
	int half_height = _glb.FRAME_H * 0.5;
	for(int c = 0; c < _glb.num_max_corners; c++) {
		_glb.corners.Xf[c] = _glb.corners.corners_openCV[c].x - half_width;
		_glb.corners.Yf[c] = _glb.corners.corners_openCV[c].y - half_height;		
	}
}

/**
 * OpenGL GLUT's Idle Function
 *
 * Retrieves video's current frame and identify chessboard's inner corners.
 * Uses TSAI for finding relation between expected and observed corners.
 * Also updates some global variables for further correct display in OpenGL.
 *
 */
void IdleFunction(void) {

	// retrieves next frame
	if(!cvGrabFrame(_glb.capture)) {
		EndOpenCV();
		exit( 0 );
	}
	IplImage *frm_original = cvRetrieveFrame(_glb.capture);
	_glb.current_frame = cvCloneImage(frm_original);

	// finds corners
	memset(_glb.corners.corners_openCV, 0, _glb.num_max_corners * sizeof(CvPoint2D32f)); 
	int cornerCount = 0;
	int result = cvFindChessboardCorners(frm_original, _glb.pattern_size, _glb.corners.corners_openCV, &cornerCount, 
										CV_CALIB_CB_ADAPTIVE_THRESH);
							//			CV_CALIB_CB_NORMALIZE_IMAGE);
							//			CV_CALIB_CB_FILTER_QUADS);
	if (result) {
		int time_init = clock();

		// sets corners in camera coordinates
		SetCornersInCameraCoord();

		// finds the camera's intrinsic parameters only once
		if (!_glb.foundInstrinsicParams) {
			
			if (tsaiCalibration2D(_glb.num_max_corners, 
								  _glb.corners.Xw, _glb.corners.Yw, 
								  _glb.corners.Xf, _glb.corners.Yf, 
								  &(_glb.f), _glb.T, _glb.r)) {
				_glb.foundInstrinsicParams = TRUE;

				// creates OpenGL's Projection and ModelView matrices
				tsaiOGLProjectionMatrix(_glb.f, _glb.NEARPLANE, _glb.FARPLANE, _glb.FRAME_W, _glb.FRAME_H, _glb.mtsaiopengl_projection);
				tsaiOGLModelViewMatrix(_glb.T, _glb.r, _glb.mtsaiopengl_modelview);

				if (DEBUG) {
					printf("\n Parametros intrinsecos e extrinsecos capturados em %d milisegundos\n", clock() - time_init);
				}
			}
		}
		else {
			// finds camera's extrinsic parameters
			if (tsaiCalibration2D_Tr(_glb.num_max_corners, 
								  _glb.corners.Xw, _glb.corners.Yw, 
								  _glb.corners.Xf, _glb.corners.Yf, 
								  _glb.f, _glb.T, _glb.r)) {
				
				// updates OpenGL's ModelView matrix
				tsaiOGLModelViewMatrix(_glb.T, _glb.r, _glb.mtsaiopengl_modelview);

				if (DEBUG) {
					printf("\n Parametros extrinsecos capturados em %d milisegundos\n", clock() - time_init);
				}
			}
		}
	}

	// releases original captured frame's memory
	// TODO cvReleaseImage(&frm_original);

	// triggers OpenGL GLUT's display callback
	glutPostRedisplay();
}

/**
 * OpenGL GLUT's Reshape Function
 *
 * Ignores changes to the window size. If the user tries
 * to change it, this function forces the resize to the
 * original size: the video's dimensions.
 */
void ReshapeFunction(int w, int h) {
	if (w != _glb.FRAME_W || h != _glb.FRAME_H)
		glutReshapeWindow(_glb.FRAME_W, _glb.FRAME_H);
	else {
		glViewport(0, 0, (GLsizei)w, (GLsizei)h);
		glMatrixMode(GL_PROJECTION);
		glLoadIdentity();
		gluOrtho2D (0.f, _glb.FRAME_W, 0.f, _glb.FRAME_H);
	}
}


/**
 * OpenGL GLUT's Display Function
 *
 * Displays video's current frame, the identified chessboard corners
 * and three lines representing the three axis: X in red, Y in green and Z in blue.
 */
void DisplayFunction() {
	
	if (!_glb.current_frame)
		return;

	// displays background image with the chessboard
	glMatrixMode(GL_MODELVIEW);
	glLoadIdentity();

	glRasterPos2f(0.f, 0.f);
	glClear( GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT );  // clears the back buffer
	glDrawPixels(_glb.FRAME_W, _glb.FRAME_H, GL_BGR_EXT, GL_UNSIGNED_BYTE, _glb.current_frame->imageData);

	// displays axis and corners if camera's params have already been found
	if (_glb.foundInstrinsicParams) {
		// applies projection and transform according to camera's parameters
		glMatrixMode(GL_PROJECTION);
		glPushMatrix();								// saves state
		glLoadIdentity();
		glLoadMatrixf(_glb.mtsaiopengl_projection);

		glMatrixMode(GL_MODELVIEW);
		glPushMatrix();								// saves state
		glLoadIdentity();
		glLoadMatrixf(_glb.mtsaiopengl_modelview);

		// displays corners as points
		glPointSize(10.0);
		glColor3f(1.0, 1.0, 0.0);
		glBegin(GL_POINTS);
		for(int c = 0; c < _glb.num_max_corners; c++) {
			glVertex3f(_glb.corners.Xw[c], _glb.corners.Yw[c], 0.0);
		}
		glEnd();

		// draws axis as lines
		glLineWidth(3);
		glBegin(GL_LINES);
			// X axis
			glColor3f(1,0,0);
			glVertex3f( 0.0, 0.0, 0.0);			glVertex3f( 3.0, 0.0, 0.0);	// main line
			glVertex3f( 3.0, 0.0, 0.0);			glVertex3f( 2.8,-0.2, 0.0); // arrow's segments
			glVertex3f( 3.0, 0.0, 0.0);			glVertex3f( 2.8, 0.2, 0.0);

			// Y axis
			glColor3f(0,1,0);
			glVertex3f( 0.0, 0.0, 0.0);			glVertex3f( 0.0, 3.0, 0.0);	// main line
			glVertex3f( 0.0, 3.0, 0.0);			glVertex3f(-0.2, 2.8, 0.0); // arrow's segments
			glVertex3f( 0.0, 3.0, 0.0);			glVertex3f( 0.2, 2.8, 0.0);

			// X axis
			glColor3f(0,0,1);
			glVertex3f( 0.0, 0.0, 0.0);			glVertex3f( 0.0, 0.0, 3.0);	// main line
			glVertex3f( 0.0, 0.0, 3.0);			glVertex3f(-0.2,+0.2, 2.8); // arrow's segments
			glVertex3f( 0.0, 0.0, 3.0);			glVertex3f(+0.2,-0.2, 2.8);
		glEnd();

		// draws some virtual objects on top of the chessboard
		// donut
		glColor3f(0,0,1);
		glPushMatrix();
		glTranslatef(2.5,2.0,0);
		glutSolidTorus(0.25,0.5,20,30);
		glPopMatrix();

		// sphere
		glColor3f(1,0,0);
		glPushMatrix();
		glTranslatef(3.5,5.5,0);
		glutSolidSphere(0.5,20,30);
		glPopMatrix();

		// cube
		glColor3f(0,1,0);
		glPushMatrix();
		glTranslatef(5.5,2.5,0);
		glutSolidCone(0.5,3,20,30);
		glPopMatrix();

		// restores previous state
		glMatrixMode(GL_PROJECTION);
		glPopMatrix();
		glMatrixMode(GL_MODELVIEW);
		glPopMatrix();
	}


	// finally swaps buffers
	glFlush();
	glutSwapBuffers();

	// releases current frame's memory
	// TODO cvReleaseImage(&(_glb.current_frame));
}

/**
 * Initializes OpenGL's states and controls.
 *
 */
void InitOpenGL() {
	glEnable (GL_DEPTH_TEST);
	glDepthFunc (GL_LESS);
	glLightModeli (GL_LIGHT_MODEL_LOCAL_VIEWER, GL_TRUE);
	glLightModeli(GL_LIGHT_MODEL_TWO_SIDE, GL_TRUE);
	glEnable (GL_LIGHTING);
	glDisable (GL_NORMALIZE);
	glShadeModel(GL_FLAT);  
	glClearColor(0.0, 0.0, 0.0, 0.0);
	glPixelStorei(GL_UNPACK_ALIGNMENT, 1);           // data alignment is 1  
}

/**
 * Main function.
 *
 * Initializes OpenGL, GLUT and video capture.
 * Loops until user closes the window or the video ends.
 *
 */
int main(int argc, char *argv[]) {
	
	// chessboard and input file definition 
	if ( InitOpenCV(9, 7, "c:/users/gattass/publichtml/ra/calibra/videos/videosTeoricamenteMelhores1.avi") == RET_ERROR )
		return RET_ERROR;
	
	InitOpenGL();

	glutInit(&argc, argv);
	glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH);
	glutInitWindowSize(_glb.FRAME_W, _glb.FRAME_H);
	glutCreateWindow("Trabalho 2 - Camera calibration");
	
	glutDisplayFunc(DisplayFunction);
	glutReshapeFunc(ReshapeFunction);
	glutIdleFunc(IdleFunction);
	
	glutMainLoop();

	EndOpenCV();

	return 0;
}