#include "RecordPlugin.h"
#include "ritkManager.h"
#include <fstream>
#include <iostream>
#include <QString>
#include <QDate>
#include <QTime>
#include <QMessageBox>
#include <QFileDialog>

#include <iostream>
#include <itkImageFileWriter.h>
#include <itkRescaleIntensityImageFilter.h>
#include <itkCastImageFilter.h>

RecordPlugin::RecordPlugin()
{
	SetSingleEventInstanceMode(false);
	// Create widget
	m_Widget = new RecordPluginWidget();

	connect(m_Widget, SIGNAL(StartRecording()), this, SLOT(StartRecording()));
	connect(m_Widget, SIGNAL(StopRecording()), this, SLOT(StopRecording()));
	connect(m_Widget, SIGNAL(PreviewChanged(bool)), this, SLOT(SetPreview(bool)));
	connect(m_Widget->m_PushButtonSetOutputDirectory, SIGNAL(clicked()), this, SLOT(PushButtonSetOutputDirectoryClicked()));
	connect(this, SIGNAL(SetFrameCounter(long)), m_Widget, SLOT(SetFrameCounter(long)));
	connect(m_Widget->m_SnapshotButton, SIGNAL(clicked()), this, SLOT(DoSnapshot()));
	connect(m_Widget->m_PNGCheckbox, SIGNAL(toggled(bool)), this, SLOT(WritePNGs(bool)));
	connect(m_Widget->m_TXTCheckbox, SIGNAL(toggled(bool)), this, SLOT(WriteTXTs(bool)));
	connect(m_Widget, SIGNAL(CustomFilenamePrefixChanged(int)), this, SLOT(UpdateFilename(int)));
	connect(m_Widget, SIGNAL(CustomFilenameChanged(QString)),this, SLOT(UpdateRadioButton(QString)));
	//m_Widget->m_rad,ioButtonNoPostfix->
	

	// Create writers
	m_RecordWriter = WriterType::New();
	m_FrameWriterTXT = WriterType::New();
	m_CustomFileName = "";

	// Recorder members
	m_RecordFlag = false;
	m_FirstFilenameChange = true;

	// True by default
	m_PreviewFlag = true;

	m_PNGFlag = false;
	m_TXTFlag = false;

	m_PNGTXTFileName = "";

	this->SetNumberOfPassesToAverageRuntime(10);

	// default Output Directory (executable)
	m_DirectoryName = QApplication::applicationDirPath();
	m_Widget->m_LabelOutputDirectory->setText(m_DirectoryName);

	m_Widget->m_radioButtonNoPostfix->setVisible(false);
	m_Widget->m_radioButtonPostfix->setVisible(false);
	m_Widget->m_radioButtonPrefix->setVisible(false);
	m_Widget->m_radioButtonNoPostfix->setChecked(true);
	m_Widget->m_TXTCheckbox->setChecked(false);
	m_Widget->m_PNGCheckbox->setChecked(false);

}

RecordPlugin::~RecordPlugin()
{
	delete m_Widget;
}


//----------------------------------------------------------------------------
QString
RecordPlugin::GetName()
{
	return tr("RecordPlugin");
}


//----------------------------------------------------------------------------
QWidget*
RecordPlugin::GetPluginGUI()
{
	return m_Widget;
}


//----------------------------------------------------------------------------
void
RecordPlugin::StartRecording()
{
	QMutexLocker locker(&m_MemberMutex);
	{
		m_RecordFlag = true;

		// Create timestamp filename
		QDate Date = QDate::currentDate();
		QTime Time = QTime::currentTime();
		QString Timestamp = QString::number(Date.year()) + "-" + QString::number(Date.month()).rightJustified(2, '0') +  "-" + QString::number(Date.day()).rightJustified(2, '0') + "_" ;
		Timestamp	     += QString::number(Time.hour()).rightJustified(2, '0')+  "-" + QString::number(Time.minute()).rightJustified(2, '0') +  "-" + QString::number(Time.second()).rightJustified(2, '0');
		
		m_CustomFileName = m_Widget->m_FilenameEdit->text();
		m_RecordFileName = m_DirectoryName + "/";
		m_RecordFileName = QDir::toNativeSeparators(m_RecordFileName);

		// If there is no typed filename the timestamp is taken as default filename
		if (m_CustomFileName.isEmpty())
		{
			m_RecordFileName += Timestamp;
		}
		else
		{
			switch(this->m_FilenameState)
			{
					
				// Custom Name with Timestamp as Prefix
			case 1: 
				m_RecordFileName += Timestamp;
				m_RecordFileName += "_";
				m_RecordFileName += m_CustomFileName;
				break;
				
				// Custom Name with Timestamp as Postfix
			case 2:
				m_RecordFileName += m_CustomFileName;
				m_RecordFileName += "_";
				m_RecordFileName += Timestamp;
				break;

				// Pure Custom Name
			default:
				m_RecordFileName += m_CustomFileName;
				break;
			}
		}
		
		m_RecordFileName += ".ris";

		LOG_DEB("FileName=" << m_RecordFileName.toStdString().c_str());

		// Set file name to writer
		m_RecordWriter->SetFileName(m_RecordFileName.toStdString().c_str());

		// Initialize frame counter
		m_RecordFrameCounter = 0;
		emit SetFrameCounter(0);
	}
}

//----------------------------------------------------------------------------
void
RecordPlugin::UpdateFilename(int state)
{
	m_FilenameState = state;
}


//----------------------------------------------------------------------------
void
RecordPlugin::StopRecording()
{
	QMutexLocker locker(&m_MemberMutex);

	m_RecordFlag = false;
	m_RecordWriter->Save();
		
	QMessageBox::information(m_Widget,tr("Sequence saved"),tr("Sequence saved under <") + m_RecordFileName + tr(">"));
}


//----------------------------------------------------------------------------
void
RecordPlugin::DoSnapshot()
{
	// Create timestamp filename
	QDate Date = QDate::currentDate();
	QTime Time = QTime::currentTime();
	m_RecordFileName = QApplication::applicationDirPath() + "/";
	m_RecordFileName += QString::number(Date.year()) + QString::number(Date.month()).rightJustified(2, '0') + QString::number(Date.day()).rightJustified(2, '0');
	m_RecordFileName += QString::number(Time.hour()).rightJustified(2, '0') + QString::number(Time.minute()).rightJustified(2, '0') + QString::number(Time.second()).rightJustified(2, '0');
	m_RecordFileName += ".ri";

	LOG_DEB("FileName=" << m_RecordFileName.toStdString().c_str());

	// Set file name to writer
	m_RecordWriter->SetFileName(m_RecordFileName.toStdString().c_str());
	m_RecordWriter->SetInput(m_CurrentFrame);
	m_RecordWriter->Update();
	m_RecordWriter->Save();
	QMessageBox::information(m_Widget,tr("Image saved"),tr("Imaage saved under <") + m_RecordFileName + tr(">"));

	if(m_PNGFlag)
		SavePNG();
	if(m_TXTFlag)
		SaveTXT();
}


//----------------------------------------------------------------------------
void
RecordPlugin::SaveTXT()
{
	if(!m_PNGFlag)
	{
		// Create timestamp filename
		QDate Date = QDate::currentDate();
		QTime Time = QTime::currentTime();
		m_PNGTXTFileName = QApplication::applicationDirPath() + "/";
		m_PNGTXTFileName += QString::number(Date.year()) + QString::number(Date.month()).rightJustified(2, '0') + QString::number(Date.day()).rightJustified(2, '0');
		m_PNGTXTFileName += QString::number(Time.hour()).rightJustified(2, '0') + QString::number(Time.minute()).rightJustified(2, '0') + QString::number(Time.second()).rightJustified(2, '0');
	}

	LOG_DEB("FileName=" << m_PNGTXTFileName.toStdString().c_str());

	// Set all regions and allocate
	ritk::RImageF2::SizeType ImageSize = m_CurrentFrame->GetRangeImage()->GetLargestPossibleRegion().GetSize();
	const ritk::RImageF2::RangeImageType::PixelType *RangeDataBuffer = m_CurrentFrame->GetRangeImage()->GetBufferPointer();

	QString txtName = m_PNGTXTFileName + ".txt";
	std::ofstream outfile (txtName.toStdString().c_str(), std::ofstream::binary);
	for(int i=0; i<ImageSize[1]; ++i)
	{
		for(int j=0; j<ImageSize[0]; ++j, ++RangeDataBuffer)
		{
			outfile << *RangeDataBuffer;
			if(j<ImageSize[0]-1)
				outfile << ",";
		}
		outfile << std::endl;
	}
}

//----------------------------------------------------------------------------
void
RecordPlugin::SavePNG()
{
	// Create timestamp filename
	QDate Date = QDate::currentDate();
	QTime Time = QTime::currentTime();
	m_PNGTXTFileName = QApplication::applicationDirPath() + "/";
	m_PNGTXTFileName += QString::number(Date.year()) + QString::number(Date.month()).rightJustified(2, '0') + QString::number(Date.day()).rightJustified(2, '0');
	m_PNGTXTFileName += QString::number(Time.hour()).rightJustified(2, '0') + QString::number(Time.minute()).rightJustified(2, '0') + QString::number(Time.second()).rightJustified(2, '0');

	LOG_DEB("FileName=" << m_PNGTXTFileName.toStdString().c_str());

	for ( unsigned int i = 0; i < m_CurrentFrame->GetNumberOfPayloads(); i++ )
	{
		// Get the payload container and its actual payload object.
		ritk::GenericPayloadContainer::ConstPointer PayloadContainer = m_CurrentFrame->GetPayloadContainer(i);
		const ritk::GenericPayloadContainer::PayloadType *Payload = (const ritk::GenericPayloadContainer::PayloadType*)PayloadContainer->GetPayload();
		
		// Cast the payload an unsigned char image. If successful, add the payload to the available payloads
		typedef itk::Image<unsigned char,2> ImageTypeUC;
		ImageTypeUC::ConstPointer ImageUC = dynamic_cast<const ImageTypeUC*>(Payload);
		if ( ImageUC )
		{
			typedef itk::RescaleIntensityImageFilter<ImageTypeUC> FilterType;
			FilterType::Pointer Filter = FilterType::New();
			Filter->SetInput(ImageUC);
			Filter->SetOutputMinimum(0);
			Filter->SetOutputMaximum(255);
			Filter->Update();

			itk::ImageFileWriter<ImageTypeUC>::Pointer writer = itk::ImageFileWriter<ImageTypeUC>::New();
			writer->SetFileName(m_PNGTXTFileName.toStdString() + "_" + PayloadContainer->GetIdentifier() + ".png");
			writer->SetInput(Filter->GetOutput());
			writer->Update();

			// continue as this is the only possible cast
			continue;
		}

		typedef itk::Image<itk::Vector<unsigned char,3>,2> ImageTypeUC3;
		ImageTypeUC3::ConstPointer ImageUC3 = dynamic_cast<const ImageTypeUC3*>(Payload);
		if ( ImageUC3 )
		{
			typedef itk::RescaleIntensityImageFilter<ImageTypeUC3> FilterType;
			itk::ImageFileWriter<ImageTypeUC3>::Pointer writer = itk::ImageFileWriter<ImageTypeUC3>::New();
			writer->SetFileName(m_PNGTXTFileName.toStdString() + "_" + PayloadContainer->GetIdentifier() + ".png");
			writer->SetInput(ImageUC3);
			writer->Update();

			// continue as this is the only possible cast
			continue;
		}

		// Cast the payload to a float image. If successful, add the payload to the available payloads
		typedef itk::Image<float,2> ImageTypeF;
		ImageTypeF::ConstPointer ImageF = dynamic_cast<const ImageTypeF*>(Payload);
		if ( ImageF )
		{
			itk::ImageFileWriter<ImageTypeUC>::Pointer writer = itk::ImageFileWriter<ImageTypeUC>::New();
			writer->SetFileName(m_PNGTXTFileName.toStdString() + "_" + PayloadContainer->GetIdentifier() + ".png");

			if(i==0) // Range
			{
				typedef itk::CastImageFilter< ImageTypeF, ImageTypeUC > CastFilterType;
				CastFilterType::Pointer castFilter = CastFilterType::New();
				castFilter->SetInput(ImageF);
				castFilter->Update();

				writer->SetInput(castFilter->GetOutput());
			}
			else // all others like amplitudes ...
			{
				typedef itk::RescaleIntensityImageFilter<ImageTypeF,ImageTypeUC> FilterType;
				FilterType::Pointer rescaleFilter = FilterType::New();
				rescaleFilter->SetInput(ImageF);
				rescaleFilter->SetOutputMinimum(0);
				rescaleFilter->SetOutputMaximum(255);
				rescaleFilter->Update();

				writer->SetInput(rescaleFilter->GetOutput());
			}
			writer->Update();

			// continue as this is the only possible cast
			continue;
		}
	}
}


//----------------------------------------------------------------------------
void
RecordPlugin::WriteTXTs(bool value)
{
	m_TXTFlag = value;
}


//----------------------------------------------------------------------------
void
RecordPlugin::WritePNGs(bool value)
{
	m_PNGFlag = value;
}


//----------------------------------------------------------------------------
void
RecordPlugin::SetPreview(bool Enable)
{
	QMutexLocker locker(&m_MemberMutex);

	m_PreviewFlag = Enable;
}


//----------------------------------------------------------------------------
void
RecordPlugin::ProcessEvent(ritk::Event::Pointer EventP)
{
	// New frame event
	if ( EventP->type() == ritk::NewFrameEvent::EventType )
	{
		// Cast the event
		ritk::NewFrameEvent::Pointer NewFrameEventP = qSharedPointerDynamicCast<ritk::NewFrameEvent,ritk::Event>(EventP);
		if ( !NewFrameEventP )
		{
			LOG_DEB("Event mismatch detected: Type=" << EventP->type());
			return;
		}

		QMutexLocker locker(&m_MemberMutex);

		this->StartTimer();
		
		LockPlugin();
		m_CurrentFrame = NewFrameEventP->RImage;
		UnlockPlugin();

		// Update preview widget
		if ( m_PreviewFlag )
		{
		  LockPlugin();
		  m_Widget->SetRangeData(m_CurrentFrame);
		 UnlockPlugin();
		}

		// Update writer
		if ( m_RecordFlag )
		{
			try
			{
				m_RecordWriter->SetInput(m_CurrentFrame);
				m_RecordWriter->Update();
			}
			catch ( itk::ExceptionObject &e )
			{
				LOG_DEB("Error while writing: " << e);
			}

			m_RecordFrameCounter++;
			emit SetFrameCounter(m_RecordFrameCounter);
			
			if(m_PNGFlag)
				SavePNG();
			if(m_TXTFlag)
				SaveTXT();
		}

		this->StopTimer();
		this->GetElapsedTime(&m_Runtime[0],&m_Runtime[1]);
	}

}

//----------------------------------------------------------------------------
void
RecordPlugin::SaveFrameTXT()
{
#if 0
	// Create timestamp filename
	QDate Date = QDate::currentDate();
	QTime Time = QTime::currentTime();

	if(m_OutputDirectory.isEmpty())
	{
		m_RecordFileName = QApplication::applicationDirPath() + "/";
	}
	else
	{
		m_RecordFileName = m_OutputDirectory + "/";
	}

	m_RecordFileName += QString::number(Date.year()) + QString::number(Date.month()) + QString::number(Date.day());
	m_RecordFileName += QString::number(Time.hour()) + QString::number(Time.minute()) + QString::number(Time.second());
	m_RecordFileName += ".txt";

	LOG_DEB("FileName=" << m_RecordFileName.toStdString());

	// Set file name to writer
	m_FrameWriterTXT->SetFileName(m_RecordFileName.toStdString().c_str());

	QMutexLocker locker(&m_CurrentFrameMutex);

	try
	{
		m_FrameWriterTXT->SetOutputFormat(RImageFileWriterF2::TXTFormat);
		m_FrameWriterTXT->SetInput(m_CurrentFrame);
		m_FrameWriterTXT->Update();				
	}
	catch ( itk::ExceptionObject &e )
	{
		LOG_DEB("Error while writing: " << e);
	}			

	m_FrameWriterTXT->Save();

	QMessageBox::information(m_Widget,tr("Frame saved as *.txt"),tr("Frame saved under <") + m_RecordFileName + tr(">"));
	//std::cout << "Frame saved as *.txt under <" << m_RecordFileName.toStdString() << ">" << std::endl;
#endif
}

//----------------------------------------------------------------------------
void
RecordPlugin::PushButtonSetOutputDirectoryClicked()
{
	m_DirectoryName = QFileDialog::getExistingDirectory(m_Widget, "Set File Path", "/");		

	if (m_DirectoryName.isEmpty())
		return;

	m_OutputDirectory = m_DirectoryName;
	m_Widget->m_LabelOutputDirectory->setText(m_OutputDirectory);
}
	

//----------------------------------------------------------------------------
void
RecordPlugin::UpdateRadioButton(QString filename)
{
	if (!filename.isEmpty())
	{
		m_Widget->m_radioButtonNoPostfix->setVisible(true);
		m_Widget->m_radioButtonPostfix->setVisible(true);
		m_Widget->m_radioButtonPrefix->setVisible(true);
		m_Widget->m_LabelFilename1->setVisible(false);
	}
	else
	{
		m_Widget->m_radioButtonNoPostfix->setVisible(false);
		m_Widget->m_radioButtonPostfix->setVisible(false);
		m_Widget->m_radioButtonPrefix->setVisible(false);
		m_Widget->m_LabelFilename1->setVisible(true);
	}
}
