#include "StatisticalModel.h"

#include "MathTools.h"
#include <itkSymmetricEigenAnalysis.h>

#include <fstream>

//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
StatisticalModel<TVector,TMatrix>
::StatisticalModel()
{
	// First, create the output for the mean vector
	m_MeanVectorOutputIdx = ProcessObject::GetNumberOfOutputs();
	VectorDecoratorPointer MeanVectorOutput = static_cast<VectorDecoratorType*>(MakeOutput(m_MeanVectorOutputIdx).GetPointer());
	ProcessObject::SetNumberOfRequiredOutputs(ProcessObject::GetNumberOfOutputs()+1); 
	ProcessObject::SetNthOutput(m_MeanVectorOutputIdx, MeanVectorOutput.GetPointer());

	// Second, create the output for the modes of variation
	m_MoVOutputIdx = ProcessObject::GetNumberOfOutputs();
	VectorContainerDecoratorPointer MoVOutput = static_cast<VectorContainerDecoratorType*>(MakeOutput(m_MoVOutputIdx).GetPointer());
	ProcessObject::SetNumberOfRequiredOutputs(ProcessObject::GetNumberOfOutputs()+1); 
	ProcessObject::SetNthOutput(m_MoVOutputIdx, MoVOutput.GetPointer());

	// Third, create the output for the modes of variation weights
	m_MoVWeightsOutputIdx = ProcessObject::GetNumberOfOutputs();
	VectorDecoratorPointer MoVWeightsOutput = static_cast<VectorDecoratorType*>(MakeOutput(m_MoVWeightsOutputIdx).GetPointer());
	ProcessObject::SetNumberOfRequiredOutputs(ProcessObject::GetNumberOfOutputs()+1); 
	ProcessObject::SetNthOutput(m_MoVWeightsOutputIdx, MoVWeightsOutput.GetPointer());

	// As the last step, create the output for the correlation coefficient matrix
	m_CorrCoeffMatrixOutputIdx = ProcessObject::GetNumberOfOutputs();
	MatrixDecoratorPointer CorrCoeffMatrixOutput = static_cast<MatrixDecoratorType*>(MakeOutput(m_CorrCoeffMatrixOutputIdx).GetPointer());
	ProcessObject::SetNumberOfRequiredOutputs(ProcessObject::GetNumberOfOutputs()+1); 
	ProcessObject::SetNthOutput(m_CorrCoeffMatrixOutputIdx, CorrCoeffMatrixOutput.GetPointer());

	// Default parameter
	m_MaxNumberOfModes = 2;
	m_MaxVariance = 0.9999;
	m_UseOrthomax = false;
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
StatisticalModel<TVector,TMatrix>
::~StatisticalModel()
{
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::DataObjectPointer
StatisticalModel<TVector,TMatrix>
::MakeOutput(unsigned int idx)
{
	// Create the appropriate output
	if ( idx == m_MeanVectorOutputIdx )
		return static_cast<itk::DataObject*>(VectorDecoratorType::New().GetPointer());
	if ( idx == m_MoVOutputIdx )
		return static_cast<itk::DataObject*>(VectorContainerDecoratorType::New().GetPointer());
	else if ( idx == m_MoVWeightsOutputIdx )
		return static_cast<itk::DataObject*>(VectorDecoratorType::New().GetPointer());
	else if ( idx == m_CorrCoeffMatrixOutputIdx )
		return static_cast<itk::DataObject*>(MatrixDecoratorType::New().GetPointer());
	else
		return Superclass::MakeOutput(idx);
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::VectorDecoratorPointer 
StatisticalModel<TVector,TMatrix>
::GetMeanVectorOutput()
{
	return static_cast<VectorDecoratorType*>(ProcessObject::GetOutput(m_MeanVectorOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::VectorDecoratorConstPointer 
StatisticalModel<TVector,TMatrix>
::GetMeanVectorOutput() const
{
	return static_cast<const VectorDecoratorType*>(ProcessObject::GetOutput(m_MeanVectorOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::VectorContainerDecoratorPointer 
StatisticalModel<TVector,TMatrix>
::GetModesOfVariationOutput()
{
	return static_cast<VectorContainerDecoratorType*>(ProcessObject::GetOutput(m_MoVOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::VectorContainerDecoratorConstPointer 
StatisticalModel<TVector,TMatrix>
::GetModesOfVariationOutput() const
{
	return static_cast<const VectorContainerDecoratorType*>(ProcessObject::GetOutput(m_MoVOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::VectorDecoratorPointer 
StatisticalModel<TVector,TMatrix>
::GetModesOfVariationWeightsOutput()
{
	return static_cast<VectorDecoratorType*>(ProcessObject::GetOutput(m_MoVWeightsOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::VectorDecoratorConstPointer 
StatisticalModel<TVector,TMatrix>
::GetModesOfVariationWeightsOutput() const
{
	return static_cast<const VectorDecoratorType*>(ProcessObject::GetOutput(m_MoVWeightsOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::MatrixDecoratorPointer 
StatisticalModel<TVector,TMatrix>
::GetCorrelationCoeffMatrixOutput()
{
	return static_cast<MatrixDecoratorType*>(ProcessObject::GetOutput(m_CorrCoeffMatrixOutputIdx));
}



//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
typename StatisticalModel<TVector,TMatrix>::MatrixDecoratorConstPointer 
StatisticalModel<TVector,TMatrix>
::GetCorrelationCoeffMatrixOutput() const
{
	return static_cast<const MatrixDecoratorType*>(ProcessObject::GetOutput(m_CorrCoeffMatrixOutputIdx));
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
void
StatisticalModel<TVector,TMatrix>
::GenerateData()
{
	// Sanity check
	if ( !GetTrainingDataInput() || GetTrainingDataInput()->Get().size() == 0  )
	{
		if ( GetFileName().empty() )
		{
			itkExceptionMacro(<<"No training data nor file name provided!");
		}
		else
		{
			Load(GetFileName());
			return;
		}
	}

	// For convenience
	const VectorContainerType &TrainingData = GetTrainingDataInput()->Get();

	// Check for invalid dimensions
	for ( unsigned int i = 1; i < TrainingData.size(); i++ )
	{
		if ( TrainingData[0].Size() != TrainingData[i].Size() )
			itkExceptionMacro(<<"Invalid dimension detected: #" << 0 << "=" << TrainingData[0].Size() << " vs. #" << i << "=" << TrainingData[i].Size());
	}

	// For convenience
	unsigned int NumSamples = TrainingData.size();
	unsigned int Dimension = TrainingData[0].Size();

	// Compute the mean vector
	std::cout << "	Computing mean vector..."; std::cout.flush();
	InternalVectorType Mean(Dimension);
	Mean.fill(0);
	for ( unsigned int n = 0; n < NumSamples; n++ )
	{
		for ( unsigned int dim = 0; dim < Dimension; dim++ )
		{
			Mean[dim] += TrainingData[n][dim];
		}
	}
	Mean /= NumSamples;
	std::cout << "[OK]" << std::endl;

	// Generate our matrix D storing the aligned samples column-wise
	std::cout << "	Generating aligned measurement matrix..."; std::cout.flush();
	InternalMatrixType D(Dimension, NumSamples);
	D.fill(0);
	for ( unsigned int n = 0; n < NumSamples; n++ )
	{
		for ( unsigned int dim = 0; dim < Dimension; dim++ )
		{
			D[dim][n] = TrainingData[n][dim] - Mean[dim];
		}
	}
	std::cout << "[OK]" << std::endl;

	// Generate our matrix T=D^t*D
	std::cout << "	Generating implicit covariance matrix..."; std::cout.flush();
	InternalMatrixType T = (D.transpose() * D)/(NumSamples-1);
	std::cout << "[OK]" << std::endl;

	// Eigen-decomposition of T
	std::cout << "	Performing eigen-decomposition..."; std::cout.flush();
	InternalVectorType EigenValues(NumSamples);
	InternalMatrixType EigenVectors(NumSamples, NumSamples);
	typedef itk::SymmetricEigenAnalysis<InternalMatrixType, InternalVectorType> AnalysisType;
	AnalysisType Analysis;
	Analysis.SetDimension(NumSamples);
	Analysis.SetOrderEigenMagnitudes(true);
	Analysis.ComputeEigenValuesAndVectors(T, EigenValues, EigenVectors);
	std::cout << "[OK]" << std::endl;

	// Now compute the eigenvectors of the corvariance matrix. Note that we have to do a normalization and re-ordering 
	std::cout << "	Computing eigen-vectors..."; std::cout.flush();
	EigenVectors = D*EigenVectors.transpose();
	InternalMatrixType SortedEigenVectors(EigenVectors.rows(), EigenVectors.cols());
	InternalVectorType SortedEigenValues(EigenValues.size());
	for ( unsigned int i = 0; i < EigenVectors.cols(); i++ )
	{
		// Flip eigen value
		SortedEigenValues[i] = EigenValues[EigenValues.size()-i-1];

		// Normalize the eigenvector
		InternalVectorType tmp(EigenVectors.rows());
		for ( unsigned int dim = 0; dim < EigenVectors.rows(); dim++ )
		{
			tmp[dim] = EigenVectors[dim][EigenVectors.cols()-i-1];
		}
		tmp = tmp.normalize();

		// Flip eigen vector
		for ( unsigned int dim = 0; dim < EigenVectors.rows(); dim++ )
		{
			SortedEigenVectors[dim][i] = tmp[dim];
		}
	}
	std::cout << "[OK]" << std::endl;
	
	// PCA stats
	unsigned int NumModes = 0;
	float TotalVariation = 0;
	InternalVectorType CumulativeVariation(SortedEigenValues.size());
	CumulativeVariation.fill(0);
	for ( unsigned int i = 0; i < SortedEigenValues.size(); i++ )
	{
		TotalVariation += SortedEigenValues[i];
		CumulativeVariation[i] = TotalVariation;
	}
	for ( unsigned int i = 0; i < SortedEigenValues.size(); i++ )
	{
		//if ( (CumulativeVariation[i]/TotalVariation) < m_MaxVariance && i < m_MaxNumberOfModes && i < NumSamples-1 )
			NumModes++;
		//else
		//	break;
		std::cout << "		MoV #" << i << ": " << SortedEigenValues[i] << " (" << ((SortedEigenValues[i]/TotalVariation)*100) << "/" << ((CumulativeVariation[i]/TotalVariation)*100) << "%)" << std::endl;
	}
	NumModes = std::min<unsigned int>(m_MaxNumberOfModes, NumModes);
	std::cout << "	Number of modes after dimensionality reduction: " << NumModes << std::endl;

	// Dimensionality reduction.
	InternalMatrixType ReducedBasis(SortedEigenVectors.rows(), NumModes);
	VectorType ReducedWeights(NumModes);
	for ( unsigned int i = 0; i < NumModes; i++ )
	{
		ReducedWeights[i] = SortedEigenValues[i];
		for ( unsigned int row = 0; row < ReducedBasis.rows(); row++ )
		{
			ReducedBasis[row][i] = SortedEigenVectors[row][i];
		}
	}
	InternalMatrixType Phi = ReducedBasis;

	// Orthomax rotation
	if ( m_UseOrthomax )
	{
		std::cout << "	Orthomax rotation...";std::cout.flush();
		typedef InternalVectorType::ValueType RealType;
		typedef vnl_svd<RealType> SVDType;
		RealType d = 0;
		InternalMatrixType Lambda;
		InternalMatrixType R(Phi.cols(), Phi.cols());
		R.set_identity();
		RealType p = Phi.rows();
		for ( int i = 0; i < 100; i++ )
		{
			RealType d_old = d;
			Lambda = Phi*R;
			/*double test = 0;
			for ( unsigned int j = 0; j < Lambda.rows(); j++ )
				test += Lambda[j][0]*Lambda[j][1];
			std::cout << "Iteration #" << i << std::endl;
			std::cout << "	Orthogonality: " << test << std::endl;
			InternalMatrixType RotCheck = R.transpose()*R;
			SVDType CheckSVD(R);
			std::cout << "	RotCheck: " << RotCheck.is_identity(1e-5) << " " << CheckSVD.determinant_magnitude() << std::endl;*/
			InternalMatrixType H = HadamardProduct<RealType>(HadamardProduct<RealType>(Lambda, Lambda), Lambda);
			InternalMatrixType D = Diag<RealType>(Lambda.transpose()*Lambda);
			SVDType SVD(Phi.transpose() * (H - (RealType(1.0)/p)*Lambda*D));

			R = SVD.U()*SVD.V().transpose();
			d = Trace<RealType>(SVD.W());

		

			if ( d/d_old <  1.000001 )
			{
				break;
			}
			else
			{
				//std::cout << "i=" << i << " d/d_old=" << (d/d_old) << std::endl;
			}
		}
		Phi = Lambda;
		std::cout << "[OK]" << std::endl;

		// Compute the weighting factors (variance of the projected training data
		//std::cout << "	Computing weights..." << std::endl;
		InternalMatrixType ProjectedData = Phi.transpose()*D;
		/*std::cout << "	PhiSize=" << Phi.rows() << "x" << Phi.cols() << std::endl;
		std::cout << "	TrainingDataSize=" << D.rows() << "x" << D.cols() << std::endl;
		std::cout << "	ProjectedDataSize=" << ProjectedData.rows() << "x" << ProjectedData.cols() << std::endl;*/
		for ( unsigned int i = 0; i < ProjectedData.rows(); i++ )
		{
			double x = 0;
			double xx = 0;
			for ( unsigned int j = 0; j < ProjectedData.cols(); j++ )
			{
				x += ProjectedData[i][j];
				xx += ProjectedData[i][j]*ProjectedData[i][j];
			}
			x /= ProjectedData.rows();
			xx /= ProjectedData.rows();
			ReducedWeights[i] = xx - x*x;
			//std::cout << "		#" << i << ": " << ReducedWeights[i] << std::endl;	
		}
	}


	// Compute the weighting factors (variance of the projected training data
	std::cout << "	Computing weights..." << std::endl;
	InternalMatrixType ProjectedData = Phi.transpose()*D;
	std::cout << "	PhiSize=" << Phi.rows() << "x" << Phi.cols() << std::endl;
	std::cout << "	TrainingDataSize=" << D.rows() << "x" << D.cols() << std::endl;
	std::cout << "	ProjectedDataSize=" << ProjectedData.rows() << "x" << ProjectedData.cols() << std::endl;
	for ( unsigned int i = 0; i < ProjectedData.rows(); i++ )
	{
		double x = 0;
		double xx = 0;
		for ( unsigned int j = 0; j < ProjectedData.cols(); j++ )
		{
			x += ProjectedData[i][j];
			xx += ProjectedData[i][j]*ProjectedData[i][j];
			}
		x /= ProjectedData.rows();
		xx /= ProjectedData.rows();
		ReducedWeights[i] = xx - x*x;
		std::cout << "		#" << i << ": " << ReducedWeights[i] << std::endl;	
	}

	


	// Copy to the output
	VectorType OutputMean(Mean.size());
	for ( unsigned int i = 0; i < Mean.size(); i++ )
		OutputMean[i] = Mean[i];

	// Copy to the output
	VectorContainerType MoVOutput;
	for ( unsigned int i = 0; i < NumModes; i++ )
	{
		VectorType MoV(ReducedBasis.rows());
		for ( unsigned int row = 0; row < ReducedBasis.rows(); row++ )
		{
			MoV[row] = Phi[row][i];
		}
		MoV = MoV.normalize();
		MoVOutput.push_back(MoV);
	}

	// Update the output
	this->GetMeanVectorOutput()->Set(OutputMean);
	this->GetModesOfVariationWeightsOutput()->Set(ReducedWeights);
	this->GetModesOfVariationOutput()->Set(MoVOutput);



	// ============================== //
	// Correlation Coefficient Matrix //
	// ============================== //
	//GenerateCorrelationCoeffMatrix();
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
void
StatisticalModel<TVector,TMatrix>::GenerateCorrelationCoeffMatrix()
{
	// For convenience
	const VectorContainerType &TrainingDataContainer = GetTrainingDataInput()->Get();
	unsigned int NumSamples = TrainingDataContainer.size();
	unsigned int Dimension = TrainingDataContainer[0].Size();

	// Sanity check
	if ( !GetTrainingDataInput() || GetTrainingDataInput()->Get().size() == 0  )	
		itkExceptionMacro(<<"No training data provided!");

	// Check for invalid dimensions
	for ( unsigned int i = 1; i < TrainingDataContainer.size(); i++ )
	{
		if ( TrainingDataContainer[0].Size() != TrainingDataContainer[i].Size() )
			itkExceptionMacro(<<"Invalid dimension detected: #" << 0 << "=" << TrainingDataContainer[0].Size() << " vs. #" << i << "=" << TrainingDataContainer[i].Size());
	}

	// Store the training data in a matrix
	MatrixType TrainingData(Dimension, NumSamples);
	for ( unsigned int i = 0; i < Dimension; i++ )
		for ( unsigned int j = 0; j < NumSamples; j++ )
			TrainingData[i][j] = TrainingDataContainer[j][i];

	// Set up the correlation coefficient matrix
	this->GetCorrelationCoeffMatrixOutput()->Set(MatrixType(100,Dimension));
	MatrixType &CCM = this->GetCorrelationCoeffMatrixOutput()->Get();

	return;


	std::cout << "	Dimension: " << Dimension << std::endl;

	std::cout << "	Computing correlation coeff...";

	// Loop over the number of samples
	unsigned int Pos = 0;
	for ( unsigned int i = 0; i < Dimension; i++ )
	{
		if ( (i%100) == 0 )
			std::cout << "\r	Computing correlation coeff " << i << "/" << (Dimension-1);std::cout.flush();

		if ( i != 2590 && i != 7190 )
			continue;

		// Compute the correlation coefficient to all training samples
		for ( unsigned int j = 0; j < Dimension; j++ )
		{
			// Initialize
			double E_X = 0;
			double E_Y = 0;
			double E_XY = 0;
			double E_XX = 0;
			double E_YY = 0;

			// Pointer to the samples
			typename MatrixType::ValueType *X = &TrainingData[i][0];
			typename MatrixType::ValueType *Y = &TrainingData[j][0];

			// Loop
			for ( unsigned int n = 0; n < NumSamples; n++ )
			{
				E_X += (*X);
				E_Y += (*Y);
				E_XY += (*X)*(*Y);
				E_XX += (*X)*(*X);
				E_YY += (*Y)*(*Y);

				X++;
				Y++;
			}

			// Normalize
			E_X /= NumSamples;
			E_Y /= NumSamples;
			E_XY /= NumSamples;
			E_XX /= NumSamples;
			E_YY /= NumSamples;

			// Pearson correlation
			double nom = E_XY - E_X*E_Y;
			double denom = sqrt(E_XX - E_X*E_X) * sqrt(E_YY - E_Y*E_Y);
			double p = 0;
			if ( p == p && abs(denom) > 1e-6 )
				p = nom / denom;

			// Update the matrix
			CCM[Pos][j] = p;
		}
		Pos++;
	}
	std::cout << "[OK]" << std::endl;

	// Set the output
	//this->GetCorrelationCoeffMatrixOutput()->Set(CCM);
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
void
StatisticalModel<TVector,TMatrix>::Save(std::string s)
{
	std::ofstream out;
	out.open(s.c_str(),std::ios::out|std::ios::trunc|std::ios::binary);

	// Write the size of each element, the dimension and number of modes
	int ElemSize = sizeof(typename VectorType::ValueType);
	int Dimension = this->GetMeanVectorOutput()->Get().size();
	int NumModes = this->GetModesOfVariationOutput()->Get().size();
	out << ElemSize << std::endl;
	out << Dimension << std::endl;
	out << NumModes << std::endl;

	// Write the mean vector
	out.write((char*)(&this->GetMeanVectorOutput()->Get()[0]), Dimension*sizeof(typename VectorType::ValueType));

	// Write the modes of variation
	for ( int i = 0; i < NumModes; i++ )
	{
		out.write((char*)(&this->GetModesOfVariationOutput()->Get()[i][0]), Dimension*sizeof(typename VectorType::ValueType));
	}

	// Write the weights
	out.write((char*)(&this->GetModesOfVariationWeightsOutput()->Get()[0]), NumModes*sizeof(typename VectorType::ValueType));

	out.close();
}


//----------------------------------------------------------------------------
template<class TVector, class TMatrix>
void
StatisticalModel<TVector,TMatrix>::Load(std::string s)
{
	std::ifstream in;
	in.open(s.c_str(),std::ios::in|std::ios::binary);

	// Read the size of each element, the dimension and number of modes
	int ElemSize;
	in >> ElemSize;
	int Dimension;
	in >> Dimension;
	int NumModes;
	in >> NumModes;

	std::cout << "SM Data:" << std::endl;
	std::cout << "	ElemSize=" << ElemSize << std::endl;
	std::cout << "	Dimension=" << Dimension << std::endl;
	std::cout << "	NumModes=" << NumModes << std::endl;

	in.seekg(((unsigned long)in.tellg())+1);

	// Read the mean vector
	VectorType Mean(Dimension);
	in.read((char*)&Mean[0], Dimension*ElemSize);

	// Read the modes of variation
	VectorContainerType MoVContainer;
	for ( int i = 0; i < NumModes; i++ )
	{
		VectorType MoV(Dimension);
		in.read((char*)&MoV[0], Dimension*ElemSize);
		MoVContainer.push_back(MoV);
	}

	// Read the weights
	VectorType Weights(NumModes);
	in.read((char*)&Weights[0], NumModes*ElemSize);
	for ( int i = 0; i < NumModes; i++ )
		std::cout << "Weight #" << i << ": " << Weights[i] << std::endl;

	in.close();

	// Set the data to the output
	this->GetMeanVectorOutput()->Set(Mean);
	this->GetModesOfVariationWeightsOutput()->Set(Weights);
	this->GetModesOfVariationOutput()->Set(MoVContainer);
	this->GetCorrelationCoeffMatrixOutput()->Set(MatrixType(100,Dimension));

	this->Modified();
}



