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

	Program:		Range Imaging Toolkit (RITK)
	Module:			ritkRangeImageToWorldCoordImageFilter.txx
	Version:		$Revision: 3113 $
	Date:			$Date: 2012-04-19 09:52:59 +0200 (Do, 19 Apr 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 "ritkRangeImageToWorldCoordImageFilter.h"
#include <itkImageRegionIterator.h>

namespace ritk
{

template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
RangeImageToWorldCoordImageFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::RangeImageToWorldCoordImageFilter()
{
	// No undistortion by default
	m_UseUndistortion = false;

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


template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
RangeImageToWorldCoordImageFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::~RangeImageToWorldCoordImageFilter()
{
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
RangeImageToWorldCoordImageFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::SetCameraParameters(CameraParameters Params)
{
	m_CameraParameters = Params;
}


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


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


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
RangeImageToWorldCoordImageFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::ThreadedGenerateData(const typename OutputImageType::RegionType &OutputRegionForThread,int threadId)
{
	// Define variables
	double fx = m_CameraParameters[CameraParameters::FX];
	double fy = m_CameraParameters[CameraParameters::FY];
	double cx = m_CameraParameters[CameraParameters::CX];
	double cy = m_CameraParameters[CameraParameters::CY];
	double k1 = m_CameraParameters[CameraParameters::K1];
	double k2 = m_CameraParameters[CameraParameters::K2];
	double p1 = m_CameraParameters[CameraParameters::P1];
	double p2 = m_CameraParameters[CameraParameters::P2];

	// Helper
	double _fx = 1./fx;
	double _fy = 1./fy;
	double accuracy = 1e-7;

	// Undistortion as in OpenCV
	typedef itk::ImageRegionConstIterator<InputImageType> RangeIteratorType;
	typedef itk::ImageRegionConstIterator<FlagImageType> FlagIteratorType;
	typedef itk::ImageRegionIterator<OutputImageType> WorldCoordIteratorType;
	RangeIteratorType itR(this->GetInput(),OutputRegionForThread);
	FlagIteratorType itM(this->GetFlagImageInput(),OutputRegionForThread);
	WorldCoordIteratorType itWC(this->GetOutput(),OutputRegionForThread);
	for ( itWC.GoToBegin(),itR.GoToBegin(),itM.GoToBegin(); !itWC.IsAtEnd(); ++itWC,++itR,++itM )
	{
		// Get index
		typename OutputImageType::IndexType Index = itWC.GetIndex();

		// Get uncorrected coords
		double xUncorrected = _fx*(Index[0]-cx);
		double yUncorrected = _fy*(Index[1]-cy);

		double x = xUncorrected; //initialize with x'' values
		double y = yUncorrected;

		// Undistort coords as requested
		if ( m_UseUndistortion )
		{
			double xFac1 = 1;
			double xFac2 = 1;
			double yFac1 = 1;
			double yFac2 = 1;
			double step = 1.0;

			double gradxOld = 0.0;
			double gradyOld = 0.0;
			double gradx = 0.0;
			double grady = 0.0;
			
			while ( fabs(xFac1) > accuracy || fabs(yFac1) > accuracy )
			{
				x = x - gradx*step;
				y = y - grady*step; 
				double x2 = x*x;
				double y2 = y*y;
				double xy = x*y;
				double r2 = x2 + y2;
				double _1k1r2k2r4 = 1 + k1*r2 + k2*r2*r2;
				double k12k2r2 = k1 + 2*k2*r2;

				xFac1 = x * _1k1r2k2r4 + p2 * (r2+2*x2) + 2*p1*xy - xUncorrected;
				xFac2 = _1k1r2k2r4 + 2 * x2 * k12k2r2 + 2*p1*y + 6*p2*x;

				yFac1 = y * _1k1r2k2r4 + p1 * (r2+2*y2) + 2*p2*xy - yUncorrected;
				yFac2 = _1k1r2k2r4 + 2 * y2 * k12k2r2 + 2*p2*x + 6*p1*y;

				gradx = 2*xFac1*xFac2;
				grady = 2*yFac1*yFac2;

				if ( gradx*gradxOld < 0.0 || grady*gradyOld < 0.0 ) 
					step *= 0.5;

				gradxOld = gradx;
				gradyOld = grady;
			}
		}

		// Save coords
		double _norm = 1.0/sqrt(x*x + y*y + 1.0);
		/*if ( itR.Value() > 1e5 )
		{
			itWC.Value()[0] = x*_norm*1.430;
			itWC.Value()[1] = y*_norm*1.430;
			itWC.Value()[2] = 1000*1.43;
		}
		else*/
		{
			itWC.Value()[0] = x*_norm*itR.Value();
			itWC.Value()[1] = y*_norm*itR.Value();
			itWC.Value()[2] = _norm*itR.Value();
		}
	}
}


//----------------------------------------------------------------------------
template<class TRangeImage, class TFlagImage, class TWorldCoordImage>
void
RangeImageToWorldCoordImageFilter<TRangeImage,TFlagImage,TWorldCoordImage>
::AfterThreadedGenerateData()
{
}


} // end namespace ritk
