/**
* INF2064 - TRABALHO 1
*
* Ricardo Costa 0311606
*/

#include <omp.h>
#include <Windows.h>
#include <cv.h>
#include <cvaux.h>
#include <highgui.h>
#include <stdio.h>
#include <math.h>

#include <queue>
#include <set>

#define ROUND(x)	(int)((x) + 0.5f)

static bool GetFileName(char* file);

/** 
Main
*/
int main( int argc, char** argv )
{
	CvCapture* capture = NULL;
	IplImage* img_out = NULL;
	IplImage* img_out_flip = NULL;
	IplImage* img_out_color = NULL;
	IplImage* img_in = NULL;
	int delay = 10;
	int baseDelay = 10;
	int threshold = 30;
	int treino = 100;
	int width, height;
	int frameCount = 0;
	bool mediana = false;
	bool doUpdate = true;
	float* average = NULL;

	std::multiset<uchar> * lastframes = NULL; // lista ordenada por valor dos últimos pixels
	std::queue<uchar> * lastframesQueue = NULL; // fila dos últimos pixels inseridos

// Cria as janelas e objetos para visualizar o vídeo
	cvNamedWindow("output", CV_WINDOW_AUTOSIZE);
	cvNamedWindow("input", CV_WINDOW_AUTOSIZE);

	CvMemStorage* storage = cvCreateMemStorage();
	IplConvKernel* openingKernel = cvCreateStructuringElementEx(2, 2, 0, 0, CV_SHAPE_RECT);

	IplImage* img_intro = cvLoadImage("intro.png");
	cvShowImage("input", img_intro);
	IplImage* img_treino = cvLoadImage("treino.png");

	cvCreateTrackbar("threshold", "input", &threshold, 255, NULL );
	cvCreateTrackbar("treino", "input", &treino, 200, NULL );

// Loop Principal
	for ( ; ; )
	{
		int key = cvWaitKey(delay);
		if ( key == 'q' || key == 27 )
		{
			break;
		}
		DWORD startTime = GetTickCount();
		switch ( key )
		{
		case 'm':
			{
				mediana = !mediana;
				printf("Treinamento com %s\n", ( mediana ? "mediana" : "media" ));
			}
			//fall through
		case 'o':
			{
				char file[MAX_PATH];
				if ( !GetFileName(file) )
				{
					break;
				}
				if ( capture )
				{
					cvReleaseCapture(&capture);
				}
				capture = cvCaptureFromAVI(file);
				if ( !capture )
				{
					printf("Erro ao carregar video de entrada\n");
					continue;
				}
				double fps = cvGetCaptureProperty(capture, CV_CAP_PROP_FPS);
				baseDelay = delay = (int)(1000. / fps + 0.5);
				frameCount = 0;
			}
			break;
		case 'u':
			{
				doUpdate = !doUpdate;
				printf("Ajuste %s\n", ( doUpdate ? "ativado" : "desativado" ));
			}
			break;
		case 'w':
			{
				if ( capture )
				{
					cvReleaseCapture(&capture);
				}
				capture = cvCaptureFromCAM(-1);
				if ( !capture )
				{
					printf("Erro ao capturar video da webcam\n");
					continue;
				}
				baseDelay = delay = 10;
				frameCount = 0;
			}
			break;
		}

		if ( !capture )
		{
			continue;
		}

	// Obtém o próximo quadro do vídeo de entrada
		IplImage* frame = cvQueryFrame(capture);
		if ( !frame )
		{
			continue;
		}
		frameCount++;

	// Cria as imagens e estruturas necessárias para o vídeo
		if ( frameCount == 1 )
		{
			if ( img_in )
			{
				cvReleaseImage(&img_in);
			}
			img_in = cvCreateImage(cvSize(frame->width, frame->height), 8, 1);
			if ( img_out )
			{
				cvReleaseImage(&img_out);
			}
			img_out = cvCreateImage(cvSize(img_in->width, img_in->height), 8, 1);
			if ( img_out_flip )
			{
				cvReleaseImage(&img_out_flip);
			}
			img_out_flip = cvCreateImage(cvSize(img_in->width, img_in->height), 8, 1);
			if ( img_out_color )
			{
				cvReleaseImage(&img_out_color);
			}
			img_out_color = cvCreateImage(cvSize(img_out->width, img_out->height), 8, 3);
			if ( average )
			{
				delete[] average;
			}
			if ( lastframes )
			{
				delete[] lastframes;
			}
			if ( lastframesQueue )
			{
				delete[] lastframesQueue;
			}
			int size = img_in->width * img_in->height;
			average = new float[size];
			ZeroMemory(average, size * sizeof(float));
			if ( mediana )
			{
				lastframes = new std::multiset<uchar>[size];
				lastframesQueue = new std::queue<uchar>[size];
			}				
			cvShowImage("output", img_treino);
		}

	// Converte imagem de entrada para grayscale e aplica um filtro de ruído (Gaussian blur 5x5)
		cvConvertImage(frame, img_in, 1);
		cvSmooth(img_in, img_in, CV_GAUSSIAN, 5, 5, 1, 1);

	// Fase de Treinamento - acumula valores dos primeiros frames do vídeo
		if ( frameCount < treino )
		{
			cvShowImage("input", frame);
			#pragma omp parallel for
			for ( int h = 0; h < img_in->height; h++ )
			{
				uchar* imgdata = (uchar*)(img_in->imageData + h * img_in->width);
				for ( int w = 0; w < img_in->width; w++ )
				{
					int pixel = h * img_in->width + w;
					if ( mediana )
					{
						lastframes[pixel].insert(imgdata[w]);
						lastframesQueue[pixel].push(imgdata[w]);
					}
					else
					{
						average[pixel] += imgdata[w];
					}
				}
			}
		}
	// Calcula a média ou mediana dos valores obtidos no treinamento
		else if ( frameCount == treino )
		{
			#pragma omp parallel for
			for ( int pixel = 0; pixel < img_in->width * img_in->height; pixel++ )
			{
				if ( mediana )
				{
					std::multiset<uchar>::iterator it = lastframes[pixel].begin();
					int size = lastframes[pixel].size() / 2;
					for ( int i = 0; i < size; i++ )
					{
						it++;
					}
					average[pixel] = *it;
				}
				else
				{
					average[pixel] /= treino;
				}
			}
		}
	// Faz a segmentação do fundo, ajustando a média ou mediana periodicamente
		else
		{
			bool calcAverage = ( (frameCount - treino) % ( mediana ? 6 : 3 ) == 0 );
			#pragma omp parallel for
			for ( int h = 0; h < img_in->height; h++ )
			{
				uchar* imgdatain = (uchar*)(img_in->imageData + h * img_in->width);
				uchar* imgdataout = (uchar*)(img_out->imageData + h * img_out->width);
				for ( int w = 0; w < img_in->width; w++ )
				{
					int pixel = h * img_in->width + w;
					if ( calcAverage && doUpdate )
					{
						if ( mediana )
						{
							uchar value = lastframesQueue[pixel].front();
							lastframesQueue[pixel].pop();
							lastframes[pixel].erase(lastframes[pixel].find(value));
							lastframes[pixel].insert(imgdatain[w]);
							lastframesQueue[pixel].push(imgdatain[w]);
							std::multiset<uchar>::iterator it = lastframes[pixel].begin();
							int size = lastframes[pixel].size() / 2;
							for ( int i = 0; i < size; i++ )
							{
								it++;
							}
							average[pixel] = *it;
						}
						else
						{
							average[pixel] = average[pixel] * 0.99f + imgdatain[w] * 0.01f;
						}
					}
					if ( abs(ROUND(average[pixel]) - imgdatain[w]) > threshold )
					{
						imgdataout[w] = 255;
					}
					else
					{
						imgdataout[w] = 0;
					}
				}
			}

		// Aplica o filtro morfológico Open (Erode + Dilate) para remover pontos de ruído
			cvMorphologyEx(img_out, img_out, NULL, openingKernel, CV_MOP_OPEN);

		// Aplica o filtro morfológico Dilate seguido de Smooth na imagem que será usada para 
		// calcular os contornos com o objetivo de suavizá-los
			cvConvertImage(img_out, img_out_color);
			cvConvertImage(img_out_color, img_out_flip, 1);
			cvDilate(img_out, img_out, NULL, 3);
			cvSmooth(img_out, img_out);

		// Encontra os contornos e os desenha na imagem de output
			CvSeq* contours = NULL;
			cvFindContours(img_out, storage, &contours, sizeof(CvContour), CV_RETR_EXTERNAL);
			for ( int index = 0; contours; contours = contours->h_next, index++ )
			{
				cvDrawContours(img_out_color, contours, CV_RGB(255,255,0), CV_RGB(0,0,0), 0, 1, CV_AA);
			}

		// Desenha os contornos também na imagem de input
			cvDilate(img_out_flip, img_out_flip, NULL, 3);
			cvSmooth(img_out_flip, img_out_flip);

			contours = NULL;
			cvFindContours(img_out_flip, storage, &contours, sizeof(CvContour), CV_RETR_EXTERNAL);
			for ( int index = 0; contours; contours = contours->h_next, index++ )
			{
				cvDrawContours(frame, contours, CV_RGB(255,255,0), CV_RGB(0,0,0), 0, 1, CV_AA);
			}

		// Exibe as imagens de input e output
			cvShowImage("input", frame);
			cvShowImage("output", img_out_color);

		// Calcula o tempo decorrido para ajustar o frame rate do vídeo
			DWORD elapsed = GetTickCount() - startTime;
			delay = baseDelay - elapsed;
			if ( delay < 1 )
			{
				delay = 1;
			}
		}
	}

// Libera as estruturas alocadas e sai do programa
	cvReleaseStructuringElement(&openingKernel);
	cvReleaseImage(&img_intro);
	cvReleaseImage(&img_out);
	cvReleaseImage(&img_out_flip);
	cvReleaseImage(&img_out_color);
	cvReleaseCapture(&capture);
	cvDestroyWindow("input");
	cvDestroyWindow("output");

	return 0;
}

/** 
GetFileName - abre um diálogo para obter o nome do arquivo de vídeo
*/
bool GetFileName(char* file)
{
	OPENFILENAME ofn;       // common dialog box structure
	HWND hwnd;              // owner window
	HANDLE hf;              // file handle

	ZeroMemory(&ofn, sizeof(ofn));
	ofn.lStructSize = sizeof(ofn);
	ofn.hwndOwner = NULL;
	ofn.lpstrFile = file;
	ofn.lpstrFile[0] = '\0';
	ofn.nMaxFile = MAX_PATH;
	ofn.lpstrFilter = "All\0*.*\0Videos\0*.AVI;*.MPG\0";
	ofn.nFilterIndex = 2;
	ofn.lpstrFileTitle = NULL;
	ofn.nMaxFileTitle = 0;
	ofn.lpstrInitialDir = NULL;
	ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;

	return ( GetOpenFileName(&ofn) == TRUE );
}
