#include "TOFSimulatorSource.h"

#include <itkImageRegionIterator.h>

#include <vtkRenderWindow.h>
#include <vtkMath.h>
#include <itkVTKImageImport.h>
#include <QDialog>
#include <QString>

#ifdef _WIN32
	#include <noise/noise.h>
#else
	#include <libnoise/noise.h>
#endif


// Euclidean distance between 2 3D points
#define DIST3D(x,y) sqrt((x[0]-y[0])*(x[0]-y[0]) + (x[1]-y[1])*(x[1]-y[1]) + (x[2]-y[2])*(x[2]-y[2]))

/// Non-linear distance correction
#define REAL_DEPTH( zBufferVal , zNear , zFar ) ( (zNear) * (zFar) / ((zBufferVal) * ((zNear) - (zFar) ) + (zFar) ) )


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
TOFSimulatorSource<TRImage,TRGBRImage>
::TOFSimulatorSource()
{
	// Invalid frustum. Has to be set by the user
	m_FrustumPlanes[0] = 0;
	m_FrustumPlanes[1] = 0;
	m_CameraPosition[0] = 0;
	m_CameraPosition[1] = 0;
	m_CameraPosition[2] = 0;
	m_FocalPoint[0] = 0;
	m_FocalPoint[1] = 0;
	m_ViewAngle = 0;

	// Initialize members related to simulation of specular reflections
	m_SimulateSpecularReflections = false;
	m_ReflectionFrequency = 2;
	m_ReflectionOctave = 9;
	m_ReflectionPersistence = 0.9;
	m_PerlinMaskLoaded = false;
	m_FileReader = ImgReaderType::New();
	m_ResampleFilter = ResampleFilterType::New();

	m_SimulateNoise = true;
	m_IncorporateAmplitudes = true;
	m_NoiseVariance = 10;
	m_RandomNumberGenerator = itk::Statistics::MersenneTwisterRandomVariateGenerator::New();

	// Initialize clock
	m_Clock = itk::RealTimeClock::New();

	// One for the depth and one for the RGB data
	this->SetNumberOfRequiredInputs(2);

	// We require a different output
	ChangeOutputMacro(RGBRImageType);
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
TOFSimulatorSource<TRImage,TRGBRImage>
::~TOFSimulatorSource()
{
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::Lock()
{
	m_Mutex.Lock();
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::Unlock()
{
	m_Mutex.Unlock();
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetViewFrustumParameters(double *CameraPosition, double *FocalPoint, double *FrustumPlanes, double ViewAngle)
{
	// Copy camera positions
	m_CameraPosition[0] = CameraPosition[0];
	m_CameraPosition[1] = CameraPosition[1];
	m_CameraPosition[2] = CameraPosition[2];

	// Copy focal point
	m_FocalPoint[0] = FocalPoint[0];
	m_FocalPoint[1] = FocalPoint[1];
	m_FocalPoint[2] = FocalPoint[2];

	// Copy frustum planes
	m_FrustumPlanes[0] = FrustumPlanes[0];
	m_FrustumPlanes[1] = FrustumPlanes[1];

	// The view angle
	m_ViewAngle = ViewAngle;
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetNoiseSimulation(bool b)
{
	m_SimulateNoise = b;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetNoiseVariance(float sigma)
{
	m_NoiseVariance = sigma;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::IncorporateAmplitudesIntoNoiseComputation(bool b)
{
	m_IncorporateAmplitudes = b;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetSpecularReflectionSimulation(bool b)
{
	m_SimulateSpecularReflections = b;
	this->Modified();
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetSpecularReflectionFrequency(int frequency)
{
	m_ReflectionFrequency = frequency;
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetSpecularReflectionOctave(int octave)
{
	m_ReflectionOctave = octave;
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::SetSpecularReflectionPersistence(double persistence)
{
	m_ReflectionPersistence = persistence;
}

//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
bool
TOFSimulatorSource<TRImage,TRGBRImage>
::LoadPerlinMask(QString FileName)
{
	Lock();
	
	m_PerlinMaskLoaded=false;

	// Found a valid file name
	if ( !FileName.isEmpty() ) 
	{
		//const char* filename = "C:\\Program Files\\RITK\\Data\\PerlinMask.png";
		m_FileReader->SetFileName(FileName.toStdString().c_str());

		try
		{
			m_FileReader->Update();
		}
		catch(itk::ExceptionObject & err )
		{
			std::cerr << "ExceptionObject caught !" << std::endl;
			std::cerr << err << std::endl;
			Unlock();
			return m_PerlinMaskLoaded;
		}
		
		m_ResampleFilter->SetInput(m_FileReader->GetOutput());
		m_ResampleFilter->SetSize(this->GetOutput()->GetRequestedRegion().GetSize());
		PerlinMaskType::SpacingType ResampleSpacing;
		ResampleSpacing[0] = m_FileReader->GetOutput()->GetSpacing()[0] * (static_cast<double>(m_FileReader->GetOutput()->GetRequestedRegion().GetSize(0)) / static_cast<double>(this->GetOutput()->GetRequestedRegion().GetSize(0)));
		ResampleSpacing[1] = m_FileReader->GetOutput()->GetSpacing()[1] * (static_cast<double>(m_FileReader->GetOutput()->GetRequestedRegion().GetSize(1)) / static_cast<double>(this->GetOutput()->GetRequestedRegion().GetSize(1)));
		m_ResampleFilter->SetOutputSpacing(ResampleSpacing);
		try
		{
			m_ResampleFilter->Update();
		}
		catch(itk::ExceptionObject & err )
		{
			std::cerr << "ExceptionObject caught !" << std::endl;
			std::cerr << err << std::endl;
			Unlock();
			return m_PerlinMaskLoaded;
		}

	}

	m_PerlinMaskLoaded=true;

	Unlock();
	return m_PerlinMaskLoaded;

}

//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::UnloadPerlinMask()
{
	Lock();
	m_PerlinMaskLoaded=false;
	Unlock();

}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
bool
TOFSimulatorSource<TRImage,TRGBRImage>
::GetPerlinMaskLoaded()
{
	return m_PerlinMaskLoaded;
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::GenerateOutputInformation()
{
	// Don't copy any information from our input as Superclass::GenerateOutputInformation would do
	this->GetOutput()->SetRegions(GetDepthImageInput()->GetRequestedRegion());

#if 1
	// Compute the focal paramers
	float FocalLength = m_FrustumPlanes[0];
	float AspectRatio = ((float)this->GetOutput()->GetRequestedRegion().GetSize()[0])/((float)this->GetOutput()->GetRequestedRegion().GetSize()[1]);
	float NearHeight = 2.f*tan(vtkMath::RadiansFromDegrees(m_ViewAngle/2.0))*FocalLength;
	float NearWidth = NearHeight*AspectRatio;
	float sx = this->GetOutput()->GetRequestedRegion().GetSize()[0] / NearWidth;
	float sy = this->GetOutput()->GetRequestedRegion().GetSize()[1] / NearHeight;
	float fx = FocalLength*sx;
	float fy = FocalLength*sy;
#else
	float fy = (this->GetOutput()->GetRequestedRegion().GetSize()[1]/2.0)/tan(vtkMath::RadiansFromDegrees(m_ViewAngle/2.0));
	float fx = (this->GetOutput()->GetRequestedRegion().GetSize()[0]/2.0)/tan(vtkMath::RadiansFromDegrees(m_ViewAngle/2.0));
#endif
	

	// Set the parameters
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::FX] = fx;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::FY] = fy;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::CX] = this->GetOutput()->GetRequestedRegion().GetSize()[0]/2.0;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::CY] = this->GetOutput()->GetRequestedRegion().GetSize()[1]/2.0;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::K1] = 0;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::K2] = 0;
	this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::K3] = 0;

	//std::cout << this->GetOutput()->GetCameraParameters() << std::endl;

	// The correct output
	CastOutput(RImageType,RGBRImageType)
	CastedOutput->GetRGBImage()->SetRegions(GetDepthImageInput()->GetRequestedRegion());
}


//----------------------------------------------------------------------------
template<class TRImage, class TRGBRImage>
void
TOFSimulatorSource<TRImage,TRGBRImage>
::GenerateRImage()
{
	// Lock to ensure thread safety
	Lock();

	// Sanity checks
	if ( !GetDepthImageInput() )
		itkExceptionMacro(<<"No depth image set!");
	if ( !GetRGBAImageInput() )
		itkExceptionMacro(<<"No RGBA image set!");
	if ( DIST3D(m_FocalPoint,m_CameraPosition) < 1e-5 || m_FrustumPlanes[0] > m_FrustumPlanes[1] || abs(m_FrustumPlanes[1]-m_FrustumPlanes[0]) < 1e-5 )
		itkExceptionMacro("Invalid view frustum parameters detected!");

	// Check for invalid image regions
	if ( GetDepthImageInput()->GetRequestedRegion() != GetRGBAImageInput()->GetRequestedRegion() )
		itkExceptionMacro(<<"Region mismatch detected beween depth image and RGBA image: " << GetDepthImageInput()->GetRequestedRegion() << " vs. " << GetRGBAImageInput()->GetRequestedRegion());
	if ( this->GetOutput()->GetRequestedRegion() != GetDepthImageInput()->GetRequestedRegion() )
		itkExceptionMacro(<<"Region mismatch detected for depth image input: " << this->GetOutput()->GetRequestedRegion() << " vs. " << GetDepthImageInput()->GetRequestedRegion());

	// The correct output
	CastOutput(RImageType,RGBRImageType)

	// Set up the perlin noise generator
	noise::module::Perlin PerlinModule;
	PerlinModule.SetFrequency(m_ReflectionFrequency);
	PerlinModule.SetOctaveCount(m_ReflectionOctave);
	PerlinModule.SetPersistence(m_ReflectionPersistence);

	// Create the RGB image that holds the amplitudes.
	const RGBAImageType::PixelType *RGBAPtr = GetRGBAImageInput()->GetBufferPointer();
	typename RGBRImageType::RGBImageType::PixelType *RGBPtr = CastedOutput->GetRGBImage()->GetBufferPointer();
	for ( unsigned long l = 0; l < GetRGBAImageInput()->GetBufferedRegion().GetNumberOfPixels(); l++)
	{
		// Conversion RGB -> Amps
		/*RGBAImageType::PixelType Sample = *RGBAPtr;
		float tmp = (11*Sample[0] + 16*Sample[1] + 5*Sample[2])*255.f/32.f;*/
		float tmp = (11*(*RGBAPtr)[0] + 16*(*RGBAPtr)[1] + 5*(*RGBAPtr)[2])*255.f/32.f;
		
		(*RGBPtr)[0] = tmp;
		(*RGBPtr)[1] = tmp;
		(*RGBPtr)[2] = tmp;

		// Update pointer
		RGBAPtr++;
		RGBPtr++;
	}

	// For convenience
	typename DepthImageType::SizeType Size = this->GetOutput()->GetRangeImage()->GetBufferedRegion().GetSize();
	double FX = this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::FX];
	double FX2 = FX*FX;
	double FY = this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::FY];
	double FY2 = FY*FY;
	double DetectorPlaneCenter[2];
	DetectorPlaneCenter[0] = this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::CX];
	DetectorPlaneCenter[1] = this->GetOutput()->GetCameraParameters()[ritk::CameraParameters::CY];

	// Set up our pointers
	const typename DepthImageType::PixelType		*DepthPtr	=	GetDepthImageInput()->GetBufferPointer();
	typename RImageType::RangeImageType::PixelType	*RangePtr	=	this->GetOutput()->GetRangeImage()->GetBufferPointer();
	typename RImageType::FlagImageType::PixelType	*FlagPtr	=	this->GetOutput()->GetFlagImage()->GetBufferPointer();

	// Data array offset for Perlin Mask
	unsigned long offset=0;

	// Loop and convert the orthographic depth data to radial range data and corrupt the data by noise and artifacts
	for ( typename DepthImageType::SizeType::SizeValueType y = 0; y < Size[1]; y++ )
	{
		for ( typename DepthImageType::SizeType::SizeValueType x = 0; x < Size[0]; x++ )
		{
			// Indices relative to the detector plane center
			float XRel = x - DetectorPlaneCenter[0];
			float YRel = y - DetectorPlaneCenter[1];

			// Valid flag by default
			*FlagPtr = RImageType::VALID;

			// The additional noise and artifacts
			float n = 0.f;

			// Generate a normal variate to simulate temporal noise
			if ( m_SimulateNoise )
			{
				// Unfortunately the mersenne twister generator is not thread safe...
				//m_Mutex.Lock();

				// Incorporate amplitudes
				if ( m_IncorporateAmplitudes && false )
				{
					/*float temp1 = m_NoiseVariance;
					float temp2 = itAmp.Value();
					n = (float)m_RandomNumberGenerators[threadId]->GetNormalVariate(0,EffectiveVariance/(itAmp.Value()*itAmp.Value()+1e-5));
					m_Mutex.Unlock();
					++itAmp;*/
				}
				else
				{
					n = (float)m_RandomNumberGenerator->GetNormalVariate(0,m_NoiseVariance) * 100;
				}

				//m_Mutex.Unlock();
			}

			// Perline noise to simulate reflections and amplitude related errors
			if ( m_SimulateSpecularReflections)
			{
				//std::cout<<x<<"/"<<y<<"  --> Value is: "<<*PerlinPtr<<std::endl;
				if ( m_PerlinMaskLoaded )
				{
					if ( m_ResampleFilter->GetOutput()->GetRequestedRegion() == this->GetOutput()->GetRequestedRegion() )
					{
						if ( abs(m_ResampleFilter->GetOutput()->GetBufferPointer()[offset]) > 0 )
						{
							// Noise generator
							double Noise = (1.0 + PerlinModule.GetValue(XRel/(Size[0]*1.0),YRel/(Size[1]*1.0),0.7))/2.0;

							// Heuristic
							if ( Noise > 0.8 )
							{
								// Scaling
								Noise *= 40;

								// Scaling from perlin mask
								Noise *= static_cast<double>(abs(m_ResampleFilter->GetOutput()->GetBufferPointer()[offset]))/255.0;

								// Spikes point towards us
								n -= Noise;

								// Invalid flag!
								*FlagPtr &= ~RImageType::VALID;
							}
						}
					}
					else
					{
						m_PerlinMaskLoaded=false;
						std::cout << std::endl << "Region of loaded perlin mask file differs from RImage regions (Load a different mask)!" << std::endl<< std::endl;
					}
				}
				else 
				{
					if ( abs(XRel) < 20 && abs(YRel) < 20 )
					{
						// Noise generator
						double Noise = (1.0 + PerlinModule.GetValue(XRel/(Size[0]*1.0),YRel/(Size[1]*1.0),0.7))/2.0;

						// Heuristic
						if ( Noise > 0.8 )
						{
							// Scaling
							Noise *= 40;

							// Spikes point towards us
							n -= Noise;

							// Invalid flag!
							*FlagPtr &= ~RImageType::VALID;
						}
					}
				}
			}

			// Assemble the final range value: (1) transform from depth to radial distance (2) add noise
			//*RangePtr = REAL_DEPTH(*DepthPtr,m_FrustumPlanes[0],m_FrustumPlanes[1]) * sqrt(1.0 + (XRel*XRel)/FX2 + (YRel*YRel)/FY2) + n;#
			*RangePtr = REAL_DEPTH(*DepthPtr,m_FrustumPlanes[0],m_FrustumPlanes[1]) + n;

			// Update pointers
			DepthPtr++;
			RangePtr++;
			FlagPtr++;
			offset++;
		}
	}

	// Set the output and ourself to modified
	this->Modified();
	this->GetOutput()->SetTimeStamp(m_Clock->GetTimeStamp());
	this->GetOutput()->Modified();

	// Done here, unlock
	Unlock();
}
