/*======================================================================================

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkRGBMappingImageFilter.txx
	Version:		$Revision: 2774 $
	Date:			$Date: 2012-03-13 10:52:35 +0100 (Di, 13 Mrz 2012) $

	Copyright (c) 2010-2011
	Pattern Recognition Lab, 
    University Erlangen-Nuremberg, Germany
    All rights reserved.

	This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. 
	See RITKCopyright.txt for more information.

	http://www5.informatik.uni-erlangen.de/research/software/ritk

========================================================================================*/


#include "ritkRGBMappingImageFilter.h"
#include <itkImageRegionIterator.h>

namespace ritk
{

template<class TWorldCoordImage, class TRGBImage>
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::RGBMappingImageFilter()
{
	// Translation vector
	m_TransVector.fill(0);

	// Rotation matrix
	m_RotMatrix.set_identity();

	// By default the mapping is not done with interpolation
	m_ExactMapping = false;

	// Invalid parameters by default
	for ( int i = 0; i < CameraParameters::NumCameraParameters; i++ )
		m_CameraParametersRGB[i] = 0;
}


template<class TWorldCoordImage, class TRGBImage>
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::~RGBMappingImageFilter()
{
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::SetCameraParametersRGB(CameraParameters Params)
{
	m_CameraParametersRGB = Params;
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::SetCameraTranslation(TransVectorType Trans)
{
	m_TransVector = Trans;

	// recalculate the transformation parameters
	RotMatrixType rotTransp = m_RotMatrix.transpose();
	TransVectorType newTranslation = rotTransp * m_TransVector;
	m_TransVector = -newTranslation;
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::SetCameraRotation(RotMatrixType Rot)
{
	m_RotMatrix = Rot;

	// recalculate the transformation parameters
	RotMatrixType rotTransp = m_RotMatrix.transpose();
	TransVectorType newTranslation = rotTransp * m_TransVector;
	m_TransVector = -newTranslation;
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::GenerateInputRequestedRegion()
{
	//Superclass::GenerateInputRequestedRegion();

	typedef ImageBase<Superclass::InputImageDimension> ImageBaseType;
	typename ImageBaseType::ConstPointer constInput = dynamic_cast< ImageBaseType const *>( this->ProcessObject::GetInput(0) );
	InputImagePointer input = const_cast< TWorldCoordImage * > ( this->GetInput(0) );
	typename InputImageType::RegionType inputRegion;
	this->CallCopyOutputRegionToInputRegion(inputRegion, this->GetOutput()->GetRequestedRegion());
	input->SetRequestedRegion( inputRegion );

	//typename ImageBaseType::ConstPointer constInputRGB = dynamic_cast< ImageBaseType const *>( this->ProcessObject::GetInput(1) );
	//InputImagePointer input1 = const_cast< TRGBImage * > ( this->GetInput(1) );
	//InputImageRegionType inputRegion1;

	//		SizeType ImageSizeRGB;
	//	ImageSizeRGB[0] = 640;
	//	ImageSizeRGB[1] = 480;
	//	RegionType RegionRGB;
	//	IndexType Index;
	//	Index[0] = 0;
	//	Index[1] = 0;
	//	RegionRGB.SetIndex(Index);
	//	RegionRGB.SetSize(ImageSizeRGB);

	//this->CallCopyOutputRegionToInputRegion(inputRegion1, this->GetOutput()->GetRequestedRegion());
	//input1->SetRequestedRegion( RegionRGB );
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::BeforeThreadedGenerateData()
{
	// Sanity check
	if ( m_CameraParametersRGB[CameraParameters::FX]<0.0001 || m_CameraParametersRGB[CameraParameters::FY]<0.0001 )
		itkExceptionMacro(<<"Invalid intrinsic parameters!");
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::SetUseUndistortion(bool b)
{
	m_UseUndistortion = b;
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::SetExactMapping(bool exact)
{
	m_ExactMapping = exact;
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::ThreadedGenerateData(const typename OutputImageType::RegionType &OutputRegionForThread,int threadId)
{
	typedef itk::ImageRegionConstIterator<InputImageType> WorldCoordIteratorConstType;
	typedef itk::ImageRegionIterator<OutputImageType> RGBIteratorType;
	WorldCoordIteratorConstType	itWC(this->GetInput(),OutputRegionForThread);
	const typename OutputImageType::PixelType* ptrRGBIn = this->GetRGBImageInput()->GetBufferPointer();
	RGBIteratorType	itRGBOut(this->GetOutput(),OutputRegionForThread);

	double fx = m_CameraParametersRGB[CameraParameters::FX];
	double fy = m_CameraParametersRGB[CameraParameters::FY];
	double cx = m_CameraParametersRGB[CameraParameters::CX];
	double cy = m_CameraParametersRGB[CameraParameters::CY];
	double k1 = m_CameraParametersRGB[CameraParameters::K1];
	double k2 = m_CameraParametersRGB[CameraParameters::K2];
	double k3 = m_CameraParametersRGB[CameraParameters::K3];
	double p1 = m_CameraParametersRGB[CameraParameters::P1];
	double p2 = m_CameraParametersRGB[CameraParameters::P2];
	int widthC = m_CameraParametersRGB[CameraParameters::SX];
	int heightC = m_CameraParametersRGB[CameraParameters::SY];

	for ( itWC.GoToBegin(),itRGBOut.GoToBegin(); !itWC.IsAtEnd(); ++itWC,++itRGBOut )
	{
        double X = itWC.Value()[0], Y = itWC.Value()[1], Z = -itWC.Value()[2];

		double x = m_RotMatrix(0,0)*X + m_RotMatrix(0,1)*Y + m_RotMatrix(0,2)*Z + m_TransVector[0];
        double y = m_RotMatrix(1,0)*X + m_RotMatrix(1,1)*Y + m_RotMatrix(1,2)*Z + m_TransVector[1];
        double z = m_RotMatrix(2,0)*X + m_RotMatrix(2,1)*Y + m_RotMatrix(2,2)*Z + m_TransVector[2];

        double r2, r4, r6, a1, a2, a3, cdist;//, icdist2;
        double xd, yd;

        z = z ? 1./z : 1;
        x *= z; y *= z;

        r2 = x*x + y*y;
        r4 = r2*r2;
        r6 = r4*r2;
        a1 = 2*x*y;
        a2 = r2 + 2*x*x;
        a3 = r2 + 2*y*y;
        cdist = 1 + k1*r2 + k2*r4 + k3*r6;
        xd = x*cdist + p1*a1 + p2*a2;
        yd = y*cdist + p1*a3 + p2*a1;

        double uRGBf = xd*fx + cx;
        double vRGBf = yd*fy + cy;

		int u = uRGBf;
		int v = vRGBf;

		if(!m_ExactMapping)
		{
			if(u>=0 && u<widthC && v>=0 && v<heightC)
			{
				itRGBOut.Value()[0] = ptrRGBIn[(int)u+(int)v*widthC][2];
				itRGBOut.Value()[1] = ptrRGBIn[(int)u+(int)v*widthC][1];
				itRGBOut.Value()[2] = ptrRGBIn[(int)u+(int)v*widthC][0];
			}
			else
			{
				itRGBOut.Value()[0] = 0;
				itRGBOut.Value()[1] = 0;
				itRGBOut.Value()[2] = 0;
			}
		}
		else
		{
			if(u>=0 && u<widthC && v>=0 && v<heightC)
			{
				double rLT = 0.0;
				double gLT = 0.0;
				double bLT = 0.0;
				double rRT = 0.0;
				double gRT = 0.0;
				double bRT = 0.0;
				double rLB = 0.0;
				double gLB = 0.0;
				double bLB = 0.0;
				double rRB = 0.0;
				double gRB = 0.0;
				double bRB = 0.0;

				double distLT = 0.0;
				double distRT = 0.0;
				double distLB = 0.0;
				double distRB = 0.0;

				int uLT = int(uRGBf-0.5);
				int vLT = int(vRGBf-0.5);
				int uRT = int(uRGBf+0.5);
				int vRT = int(vRGBf-0.5);
				int uLB = int(uRGBf-0.5);
				int vLB = int(vRGBf+0.5);
				int uRB = int(uRGBf+0.5);
				int vRB = int(vRGBf+0.5);

				if(uLT>0 && vLT>0)
					distLT	= sqrt((uLT+0.5-uRGBf)*(uLT+0.5-uRGBf) + (vLT+0.5-vRGBf)*(vLT+0.5-vRGBf));
				if(uRT<widthC-1	&& vRT>0)
					distRT	= sqrt((uRT+0.5-uRGBf)*(uRT+0.5-uRGBf) + (vRT+0.5-vRGBf)*(vRT+0.5-vRGBf));
				if(uLB>0 && vLB<heightC-1)
					distLB	= sqrt((uLB+0.5-uRGBf)*(uLB+0.5-uRGBf) + (vLB+0.5-vRGBf)*(vLB+0.5-vRGBf));
				if(uRB<widthC-1 && vRB<heightC-1)
					distRB	= sqrt((uRB+0.5-uRGBf)*(uRB+0.5-uRGBf) + (vRB+0.5-vRGBf)*(vRB+0.5-vRGBf));

				double distsSum = distLT + distRT + distLB + distRB;

				if(uLT>=0 && vLT>=0)
				{
					rLT = ptrRGBIn[(vLT*widthC+uLT)][0]*distLT/distsSum;
					gLT = ptrRGBIn[(vLT*widthC+uLT)][1]*distLT/distsSum;
					bLT = ptrRGBIn[(vLT*widthC+uLT)][2]*distLT/distsSum;
				}
				if(uRT<widthC-1 && vRT>=0)
				{
					rRT = ptrRGBIn[(vRT*widthC+uRT)][0]*distRT/distsSum;
					gRT = ptrRGBIn[(vRT*widthC+uRT)][1]*distRT/distsSum;
					bRT = ptrRGBIn[(vRT*widthC+uRT)][2]*distRT/distsSum;
				}
				if(uLB>=0 && vLB<heightC-1)
				{
					rLB = ptrRGBIn[(vLB*widthC+uLB)][0]*distLB/distsSum;
					gLB = ptrRGBIn[(vLB*widthC+uLB)][1]*distLB/distsSum;
					bLB = ptrRGBIn[(vLB*widthC+uLB)][2]*distLB/distsSum;
				}
				if(uRB<widthC-1 && vRB<heightC-1)
				{
					rRB = ptrRGBIn[(vRB*widthC+uRB)][0]*distRB/distsSum;
					gRB = ptrRGBIn[(vRB*widthC+uRB)][1]*distRB/distsSum;
					bRB = ptrRGBIn[(vRB*widthC+uRB)][2]*distRB/distsSum;
				}
				double c1 = rLT + rRT + rLB + rRB;
				double c2 = gLT + gRT + gLB + gRB;
				double c3 = bLT + bRT + bLB + bRB;

				itRGBOut.Value()[0] = c1+0.5;
				itRGBOut.Value()[1] = c2+0.5;
				itRGBOut.Value()[2] = c3+0.5;
			}
			else
			{
				itRGBOut.Value()[0] = 0;
				itRGBOut.Value()[1] = 0;
				itRGBOut.Value()[2] = 0;
			}
		}
	}
}


//----------------------------------------------------------------------------
template<class TWorldCoordImage, class TRGBImage>
void
RGBMappingImageFilter<TWorldCoordImage,TRGBImage>
::AfterThreadedGenerateData()
{
}


} // end namespace ritk
