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

Program:   ORFEO Toolbox
Language:  C++
Date:      $Date$
Version:   $Revision$


Copyright (c) Centre National d'Etudes Spatiales. All rights reserved.
See OTBCopyright.txt for details.


This software is distributed WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.  See the above copyright notices for more information.

=========================================================================*/
#ifndef __otbExtractPersistentFilter_h
#define __otbExtractPersistentFilter_h

#include "otbPersistentImageFilter.h"
#include <vnl/vnl_matrix.h>
#include <map>
#include <itkImageRegionConstIteratorWithIndex.h>
#include <itkProgressReporter.h>

namespace otb
{

	//template<class TInputImage, class TVInputImage > 
	template<class TInputImage > 
	class ITK_EXPORT ExtractPersistentFilter : 
		public otb::PersistentImageFilter<TInputImage, TInputImage>
	{


	private: 
		// Temporary results container 
		std::vector< std::vector < vnl_matrix<double> > >  m_temporaryTables; 
		std::vector< vnl_vector<unsigned long int> >   m_tempCount; 

		// Final result member 
		std::vector<vnl_matrix<double> > m_signatures;
		vnl_vector<unsigned long int>  m_countTrain, m_countSample;

		int m_numB;
		std::map< int,  int> m_Map;	
		int m_size, m_nirIndex;
		unsigned int m_maxSampleSize;
		float m_minRange, m_maxRange, m_waterThreshold;

	public:

		/** Standard typedefs */
		typedef ExtractPersistentFilter								Self;
		typedef otb::PersistentImageFilter<TInputImage, TInputImage>	Superclass;
		typedef itk::SmartPointer<Self>								Pointer;
		typedef itk::SmartPointer<const Self>							ConstPointer;

		/** Type macro */
		itkNewMacro(Self);

		/** Creation through object factory macro */
		itkTypeMacro(ExtractPersistentFilter, otb::PersistentImageFilter);

		/** Template parameters typedefs */
		typedef TInputImage							ImageType;
		typedef typename ImageType::ConstPointer		ConstPointerType;
		typedef typename ImageType::RegionType		RegionType;
		typedef typename ImageType::PixelType			PixelType;
		typedef typename ImageType::InternalPixelType			InternalPixelType;

		itkSetMacro(size,int);
		itkSetMacro(numB,int);


		void SetMap(std::map< int, int> Map)
		{
			m_Map= Map;
			if (!(m_size > 0)) m_size = m_Map.size();
		} 

		void SetMinRange(float minRange)
		{
			m_minRange= minRange;
		} 

		void SetMaxRange(float maxRange)
		{
			m_maxRange= maxRange;
		} 

		void SetWaterThreshold(float waterThreshold)
		{
			m_waterThreshold= waterThreshold;
		} 

		void SetNIRIndex(int nirIndex)
		{
			m_nirIndex= nirIndex;
		} 

		void SetCountTrain(vnl_vector<unsigned long int> countTrain)
		{
			m_countTrain= countTrain;
		} 

		std::vector<vnl_matrix<double> > GetSampledSignatures()
		{
			return m_signatures;
		} 

		vnl_vector<unsigned long int>  GetCountSample()
		{
			return m_countSample;
		} 

		void SetMaxSampleSize(unsigned int maxSampleSize)
		{
			m_maxSampleSize = maxSampleSize;
		} 


		unsigned int  GetMaxSampleSize()
		{
			return m_maxSampleSize;
		} 


		virtual void Reset() 
		{ 
			// Retrieve the number of threads 
			unsigned int numberOfThreads = this->GetNumberOfThreads(); 
			// Reset the temporary results container
			m_tempCount = std::vector<vnl_vector<unsigned long int> >(numberOfThreads, vnl_vector<unsigned long int>(m_size, 0));
			m_temporaryTables = std::vector< std::vector<vnl_matrix<double> > >(numberOfThreads);

			for (int i=0;i<numberOfThreads;i++)
			{
				m_temporaryTables[i] = std::vector<vnl_matrix<double> >(m_size, vnl_matrix<double>(m_maxSampleSize,m_numB,0 ));
			}

			// Reset the final result 
			m_signatures = std::vector<vnl_matrix<double> >(m_size); 
			m_countSample =  vnl_vector<unsigned long int>(m_size, 0);

		}

		virtual void Synthetize() 
		{ 
			// compute the total number of samples in each class
			for(unsigned int threadId = 0; threadId < this->GetNumberOfThreads();++threadId) 
			{ 
				m_countSample+=m_tempCount[threadId]; 
			}

			//aggregate the values of the matrices

			for(int cl=0;cl<m_size;cl++)
			{
				m_signatures[cl]=vnl_matrix<double>(m_countSample[cl],m_numB);
				int rowcounter(-1);

				for(unsigned int threadId = 0; threadId < this->GetNumberOfThreads();++threadId) 
				{ 
					for(int i=0; i<m_tempCount[threadId][cl];i++)
					{
						rowcounter++;
						for(int j=0;j<m_numB;j++)
						{
							m_signatures[cl].put(rowcounter, j,(double)m_temporaryTables[threadId][cl].get(i,j));
						}
					}
				}	
			}
		}

	protected: 

		virtual void ThreadedGenerateData(const RegionType& outputRegionForThread, //int threadId) 
			itk::ThreadIdType threadId)
		{ 
			ConstPointerType inputClassPtr = this->GetInput(0);
			ConstPointerType inputImPtr = this->GetInput(1);

			// Declare an iterator on the region 
			itk::ImageRegionConstIteratorWithIndex<ImageType> itClass(inputClassPtr, outputRegionForThread); 
			itk::ImageRegionConstIteratorWithIndex<ImageType> itIm(inputImPtr, outputRegionForThread); 

			int classeMask, iter;

			// Walk the region of the image with the iterator 
			for (itClass.GoToBegin(),itIm.GoToBegin(); !itClass.IsAtEnd(),!itIm.IsAtEnd(); ++itClass,++itIm) 
			{ 
				PixelType imagePixel =		itIm.Get();
				InternalPixelType labelPixel =		itClass.Get()[0];

				if (imagePixel[m_nirIndex]>=m_waterThreshold)
				{
					bool bOK(true);
					vnl_vector<double> pixVal(m_numB);
					for (unsigned int i = 0; i < m_numB ; i++)
					{
						if ((imagePixel[i]>=m_minRange) && (imagePixel[i]<=m_maxRange) ) pixVal.put(i,imagePixel[i]);
						else bOK=false;
					}

					

					//chose where to write the string containing pixel values, based on the values of the labelled image (mask). Also increment the number of pixels per class.
					std::map<int, int>::iterator iter = m_Map.find((int) labelPixel);
					if (( iter != m_Map.end()) && bOK ) 
					{

						classeMask = iter->second;

						if ( ((rand()/(double)RAND_MAX) < m_maxSampleSize/(double)(m_countTrain[classeMask]+1) ) || (classeMask==0) ) // random selection for approximately up to 60000 samples points
						{
							if (m_tempCount[threadId](classeMask) < m_maxSampleSize/(this->GetNumberOfThreads()))//stop at m_maxSampleSize to avoid overflow
							{
								m_temporaryTables[threadId][classeMask].set_row(m_tempCount[threadId](classeMask), pixVal);
								m_tempCount[threadId](classeMask)++; 
							}
						}
					}
				}
			}
		}
	};
}
#endif