#include "MainWindow.h"
#include "KLTTrackerProcessor.h"
#include "FASTTrackerProcessor.h"
#include "RuinsReconstructorProcessor.h"
#include <QDir>
#include <QMenu>
#include <QAction>
#include <QComboBox>
#include <QStatusBar>
#include <QDockWidget>
#include <QFileDialog>
#include <QMessageBox>
#include <QInputDialog>
#include <QApplication>
#include <QDesktopWidget>

namespace
{
	const int WINDOW_WIDTH				= 320;
	const int WINDOW_HEIGHT				= 240;
	const int NUM_POINTS				= 32;
	const double LEGO_W					= 1.6;
	const double LEGO_H					= 0.95;
	
	const double MODEL_POINTS[3*NUM_POINTS] =
	{
		3*LEGO_W,1*LEGO_H,6*LEGO_W,
		3*LEGO_W,1*LEGO_H,5*LEGO_W,
		2*LEGO_W,1*LEGO_H,6*LEGO_W,
		2*LEGO_W,1*LEGO_H,5*LEGO_W,
		0*LEGO_W,1*LEGO_H,6*LEGO_W,
		0*LEGO_W,2*LEGO_H,6*LEGO_W,
		0*LEGO_W,3*LEGO_H,6*LEGO_W,
		0*LEGO_W,3*LEGO_H,5*LEGO_W,
		3*LEGO_W,3*LEGO_H,6*LEGO_W,
		3*LEGO_W,3*LEGO_H,5*LEGO_W,
		3*LEGO_W,2*LEGO_H,6*LEGO_W,
		3*LEGO_W,2*LEGO_H,5*LEGO_W,
		1*LEGO_W,3*LEGO_H,4*LEGO_W,
		1*LEGO_W,3*LEGO_H,2*LEGO_W,
		1*LEGO_W,3*LEGO_H,1*LEGO_W,
		1*LEGO_W,2*LEGO_H,1*LEGO_W,
		1*LEGO_W,1*LEGO_H,1*LEGO_W,
		3*LEGO_W,3*LEGO_H,1*LEGO_W,
		3*LEGO_W,3*LEGO_H,0*LEGO_W,
		4*LEGO_W,2*LEGO_H,1*LEGO_W,
		4*LEGO_W,2*LEGO_H,0*LEGO_W,
		4*LEGO_W,1*LEGO_H,1*LEGO_W,
		6*LEGO_W,2*LEGO_H,1*LEGO_W,
		6*LEGO_W,2*LEGO_H,0*LEGO_W,
		7*LEGO_W,1*LEGO_H,1*LEGO_W,
		7*LEGO_W,1*LEGO_H,0*LEGO_W,
		0*LEGO_W,3*LEGO_H,4*LEGO_W,
		0*LEGO_W,3*LEGO_H,2*LEGO_W,
		0*LEGO_W,0*LEGO_H,6*LEGO_W,
		1*LEGO_W,0*LEGO_H,6*LEGO_W,
		5*LEGO_W,0*LEGO_H,6*LEGO_W,
		7*LEGO_W,0*LEGO_H,6*LEGO_W,
	};
}

MainWindow::MainWindow( QWidget *parent, Qt::WFlags flags )
	: QMainWindow( parent, flags ), _videoCap( 0 ),
		_rgbInputBuffer( 0 ), _originalBuffer( 0 ), _processedBuffer( 0 )
{
	_width = 0;
	_height = 0;
	_bufferSize = 0;

	// build processors list
	_processorCombo = new QComboBox( this );
	_processorCombo->addItem( "Kanade-Lucas-Tomasi Tracker" );
	_processorCombo->addItem( "FAST Tracker" );
	_processorCombo->addItem( "Ruins Reconstructor" );
	_processorCombo->setObjectName( "processorCombo" );

	// add toolbar items
	ui.setupUi( this );
	ui.mainToolBar->addWidget( new QLabel( "  Processor: ", ui.mainToolBar ) );
	_processorCombo->setParent( ui.mainToolBar );
	ui.mainToolBar->addWidget( _processorCombo );
	ui.mainToolBar->addAction( ui.actionOpenProcessorProperties );
	ui.mainToolBar->addAction( ui.actionReprocess );

	// build actual processors
	_processor = 0;
	_processorList.push_back( new KLTTrackerProcessor() );
	_processorList.push_back( new FASTTrackerProcessor() );
	_processorList.push_back( new RuinsReconstructorProcessor( ui.leftCanvas ) );

	// build status bar
	_frameRateLabel = new QLabel( statusBar() );
	_frameRateLabel->setText( "   Frame Rate: 0 Hz   " );
	statusBar()->addPermanentWidget( _frameRateLabel );

	// setup canvas
	ui.leftCanvas->setNumber( 0 );
	ui.rightCanvas->setNumber( 1 );
	ui.leftCanvas->setProcessor( new ProxyProcessor( ui.rightCanvas ) );

	// call methods that should trigger events
	_processorCombo->setCurrentIndex( 1 );
	//on_processorCombo_currentIndexChanged( 0 );
}

MainWindow::~MainWindow()
{
	delete []_rgbInputBuffer;
	delete []_originalBuffer;
	delete []_processedBuffer;

	if( _videoCap )
	{
		imVideoCaptureDestroy( _videoCap );
	}
}

bool MainWindow::initialize()
{
	_videoCap = imVideoCaptureCreate();
	if( !_videoCap )
	{
		QMessageBox::warning( this, "Warning", "No Capture Device available!" );
		ui.actionToggleConnection->setEnabled( false );
		return false;
	}

	return true;
}

void MainWindow::on_actionOpenImage_triggered()
{
	QString filename = QFileDialog::getOpenFileName( this, "Open Image File", QString(), "Images (*.png *.jpg *.gif)" );
	if( filename.isNull() )
		return;

	_image.load( filename );
	if( _image.isNull() )
	{
		QMessageBox::critical( this, "Error", "Could not load the specified file! Unsupported format?" );
		return;
	}

	resetFormat();

	// load image into _originalBuffer, converting as necessary
	unsigned char* luminance = _originalBuffer;
	unsigned char* rgb = _rgbInputBuffer;
	for( int j = _height - 1; j >= 0; --j )
	{
		for( int i = 0; i < _width; ++i )
		{
			QRgb color = _image.pixel( i, j );
			int r = qRed( color );
			int g = qGreen( color );
			int b = qBlue( color );
			*rgb++ = r;
			*rgb++ = g;
			*rgb++ = b;
			*luminance++ = ( 0.2999 * r + 0.587 * g + 0.114 * b );
		}
	}

	// process
	_processor->process( _originalBuffer, _processedBuffer );

	// update both canvas
	ui.leftCanvas->updateGL();
	ui.rightCanvas->updateGL();
}

void MainWindow::on_actionSaveImage_triggered()
{
	QString filename = QFileDialog::getSaveFileName( this, "Choose a filename to save the Image under",
		QString(), "PNG Portable Network Graphics (*.png)" );
	
	if( filename.isNull() )
		return;

	if( !filename.endsWith( ".png" ) )
		filename.append( ".png" );

	// create the image
	QImage image( _width, _height, QImage::Format_RGB32 );
	unsigned char* buffer = _originalBuffer;
	for( int j = _height - 1; j >= 0; --j )
	{
		for( int i = 0; i < _width; ++i )
		{
			unsigned char lum = *buffer++;
			image.setPixel( i, j, qRgb( lum, lum, lum ) );
		}
	}

	image.save( filename, "PNG" );
}

void MainWindow::on_actionToggleConnection_toggled( bool on )
{
	if( on )
	{
		// reset any open image
		_image = QImage();

		// connect to a capture device
		int device = 0;
		int deviceCount = imVideoCaptureDeviceCount();
		if( deviceCount > 1 )
		{
			QStringList deviceDescs;
			for( int i = 0; i < deviceCount; i++ )
			{
				deviceDescs << imVideoCaptureDeviceDesc( i );
			}

			bool ok;
			QString item = QInputDialog::getItem( this, "Capture Device Selection",
				"Select the capture device you wish to connect to:", deviceDescs, 0, false, &ok );
			
			if( !ok )
			{
				return; // cancel connect
			}

			device = deviceDescs.indexOf( item );

			if( device < 0 )
				qFatal( "Internal error: couldn't recover the selected device!" );
		}

		if( !imVideoCaptureConnect( _videoCap, device ) ) 
		{
			QMessageBox::critical( this, "Error", "Could not connect to the capture device!" );
			return;
		}

		// defaults to the following resolution:
		imVideoCaptureSetImageSize( _videoCap, WINDOW_WIDTH, WINDOW_HEIGHT );

		// connection was successful
		resetFormat();

		ui.actionOpenImage->setEnabled( false );
		ui.actionDeviceOptions->setEnabled( true );
		ui.actionToggleCapture->setEnabled( true );
		ui.actionToggleCapture->setChecked( true );
	}
	else
	{
		// pause capture
		ui.actionToggleCapture->setChecked( false );
		ui.actionToggleCapture->setEnabled( false );

		// disconnect from the capture device
		imVideoCaptureDisconnect( _videoCap );

		ui.actionOpenImage->setEnabled( true );
		ui.actionDeviceOptions->setEnabled( false );
	}
}

void MainWindow::on_actionDeviceOptions_triggered()
{
	int dialogId = 0;
	
	// Build a menu listing the available configuration dialogs...
	int dlgCount = imVideoCaptureDialogCount( _videoCap );
	if( dlgCount > 1 )
	{
		QMenu menu;
		for( int i = 0; i < dlgCount; i++)
		{
			QAction *action = menu.addAction( imVideoCaptureDialogDesc( _videoCap, i ) );
			action->setData( i );
		}

		QAction *choosen = menu.exec( QCursor::pos() );
		if( !choosen )
			return;

		dialogId = choosen->data().toInt();
	}

	int live = imVideoCaptureLive( _videoCap, -1 );
	if( live )
	{
		// deactivate the capture before calling the dialog
		imVideoCaptureLive( _videoCap, 0 );
	}

	if( !imVideoCaptureShowDialog( _videoCap, dialogId, NULL ) ) 
	{
		QMessageBox::critical( this, "Error", "Could not set configuration." );
		return;
	}

	// the dialogs can change the image size...
	resetFormat();

	if( live )
	{
		// restore the live state
		imVideoCaptureLive( _videoCap, 1 );
	}
}
void MainWindow::on_actionDisplayInColor_toggled( bool on )
{
	if( on )
	{
		ui.rightCanvas->setFrameBuffer( _width, _height, GL_RGB, _rgbInputBuffer );
	}
	else
	{
		ui.rightCanvas->setFrameBuffer( _width, _height, GL_LUMINANCE, _processedBuffer );
	}
	ui.rightCanvas->updateGL();
}

void MainWindow::on_actionToggleCapture_toggled( bool on )
{
	// invert live state
	int live = !imVideoCaptureLive( _videoCap, -1 );

	if( !imVideoCaptureLive( _videoCap, live ) ) 
	{
		QMessageBox::critical( this, "Error", "Could not change live state!" );
		return;
	}

	// start or kill the capture timer
	if( live )
	{
		_frameCount = 0;
		_frameTime.start();
		_captureTimer = startTimer( 0 );
	}
	else
	{
		killTimer( _captureTimer );
	}
}

void MainWindow::on_actionOpenProcessorProperties_triggered()
{
	QWidget* widget = _processor->getPropertiesWidget( this );
	if( !widget )
	{
		QMessageBox::warning( this, "Not Available", "The current processor does not have a Properties Widget." );
		return;
	}
	widget->move( frameGeometry().topRight() );
	widget->show();
}

void MainWindow::on_processorCombo_currentIndexChanged( int index )
{
	Processor *newProcessor = _processorList[index];
	
	ui.processorToolBar->clear();
	newProcessor->setupToolBar( ui.processorToolBar );

	if( _bufferSize )
	{
		newProcessor->resize( _width, _height );
	}

	ui.rightCanvas->setProcessor( newProcessor );

	if( _bufferSize )
	{
		newProcessor->process( _originalBuffer, _processedBuffer );
		ui.rightCanvas->updateGL();
	}

	_processor = newProcessor;
}

void MainWindow::on_actionReprocess_triggered()
{
	if( _bufferSize )
	{
		_processor->process( _originalBuffer, _processedBuffer );
		ui.leftCanvas->updateGL();
		ui.rightCanvas->updateGL();
	}
}

void MainWindow::timerEvent( QTimerEvent *e )
{
	imVideoCaptureFrame( _videoCap, _rgbInputBuffer, IM_RGB|IM_PACKED, 1000 );

	// extract and invert luminance
	uchar* luminance = _originalBuffer;
	uchar* rgb = _rgbInputBuffer;
	for( int i = 0; i < _bufferSize; ++i )
	{
		*luminance++ = ( 0.2999 * (*rgb++) + 0.587 * (*rgb++) + 0.114 * (*rgb++) );
	}

	// process
	_processor->process( _originalBuffer, _processedBuffer );

	// update both canvas
	ui.leftCanvas->updateGL();
	ui.rightCanvas->updateGL();

	// estimate frame-count
	++_frameCount;
	if( _frameTime.elapsed() >= 1000 )
	{
		double fps = _frameCount * ( 1000.0 / _frameTime.restart() );
		_frameRateLabel->setText( QString( "   Frame Rate: %1 Hz   " ).arg( fps, 0, 'f', 2 ) );
		_frameCount = 0;
	}
}

void MainWindow::resetFormat()
{
	delete []_rgbInputBuffer;
	delete []_originalBuffer;
	delete []_processedBuffer;

	if( !_image.isNull() || ui.actionToggleConnection->isChecked() )
	{
		if( !_image.isNull() )
		{
			_width = _image.width();
			_height = _image.height();
		}
		else
		{
			imVideoCaptureGetImageSize( _videoCap, &_width, &_height );
			if( _width == 0 || _height == 0 )
			{
				QMessageBox::critical( this, "Error", "Invalid image size reported for the capture device!" );
			}
		}

		_bufferSize = _width * _height;
		_rgbInputBuffer = new unsigned char[_bufferSize * 3];
		_originalBuffer = new unsigned char[_bufferSize];
		_processedBuffer = new unsigned char[_bufferSize];

		ui.leftCanvas->setFrameBuffer( _width, _height, GL_LUMINANCE, _originalBuffer );
		on_actionDisplayInColor_toggled( ui.actionDisplayInColor->isChecked() );
	}
	else
	{
		// blank screens
		_width = WINDOW_WIDTH;
		_height = WINDOW_HEIGHT;
		_bufferSize = 0;
		_rgbInputBuffer = 0;
		_originalBuffer = 0;
		_processedBuffer = 0;

		ui.leftCanvas->setFrameBuffer( _width, _height, 0, 0 );
		ui.rightCanvas->setFrameBuffer( _width, _height, 0, 0 );
	}

	// notify the processor about any resolution change
	_processor->resize( _width, _height );

	// calculate by how much we need grow or shrink the main window
	float aspect = static_cast<float>( _width ) / _height;
	int dw = _width - ui.leftCanvas->width();
	int dh = _height - ui.leftCanvas->height();
	QRect screenRect = QApplication::desktop()->screenGeometry();
	int newWidth = width() + dw * 2;
	int newHeight = height() + dh;
	int maxWidth = screenRect.width() * 0.95;
	if( newWidth > maxWidth )
	{
		float scale = maxWidth / static_cast<float>( newWidth );
		newWidth *= scale;
		newHeight *= scale;
	}
	resize( newWidth, newHeight );
}
