#ifndef STATISTICALMODEL_H__
#define STATISTICALMODEL_H__

#include <itkProcessObject.h>
#include <itkDataObject.h>
#include <itkSimpleDataObjectDecorator.h>

#include <itkArray.h>
#include <itkArray2D.h>

#include <iostream>
std::ostream& operator<<(std::ostream &out, const std::vector<itk::Array<float> > &container);
std::ostream& operator<<(std::ostream &out, const std::vector<itk::Array<double> > &container);



using namespace itk;

/**	@class	StatisticalModel
 *	@author	Jakob Wasza
 *	@brief	Statistical model
 *
 *	@details
 */
template<class TVector, class TMatrix>
class StatisticalModel : public itk::ProcessObject
{
public:
	/** @name	Typedefs */
	//@{

	/** @name	Standard typedefs */
	//@{
	typedef StatisticalModel					Self;
	typedef itk::ProcessObject					Superclass;
	typedef itk::SmartPointer<Self>				Pointer;
	typedef itk::SmartPointer<const Self>		ConstPointer;
	typedef itk::DataObject::Pointer			DataObjectPointer;
	//@}

	/**	@name	Typedefs for vectors and containers */
	//@{
	typedef TVector													VectorType;
	typedef itk::SimpleDataObjectDecorator<VectorType>				VectorDecoratorType;
	typedef typename VectorDecoratorType::Pointer					VectorDecoratorPointer;
	typedef typename VectorDecoratorType::ConstPointer				VectorDecoratorConstPointer;
	typedef std::vector<VectorType>									VectorContainerType;
	typedef itk::SimpleDataObjectDecorator<VectorContainerType>		VectorContainerDecoratorType;
	typedef typename VectorContainerDecoratorType::Pointer			VectorContainerDecoratorPointer;
	typedef typename VectorContainerDecoratorType::ConstPointer		VectorContainerDecoratorConstPointer;
	typedef TMatrix													MatrixType;
	typedef itk::SimpleDataObjectDecorator<MatrixType>				MatrixDecoratorType;
	typedef typename MatrixDecoratorType::Pointer					MatrixDecoratorPointer;
	typedef typename MatrixDecoratorType::ConstPointer				MatrixDecoratorConstPointer;
	//@}

	//@}

	/// Object creation
	itkNewMacro(Self);

	/// RTTI
	itkTypeMacro(StatisticalModel, itk::ProcessObject);

	/**	@name	Training data input */
	//@{
	itkSetDecoratedInputMacro(TrainingData, VectorContainerType, 0);
	//@}

	/**	Set the file name for the IO mode */
	//@{
	itkSetMacro(FileName, std::string);
	itkGetMacro(FileName, std::string);
	//@}

	/**	@name	Set the maximum number of modes used to generate the model */
	//@{
	itkSetMacro(MaxNumberOfModes, unsigned int);
	itkGetMacro(MaxNumberOfModes, unsigned int);
	itkGetConstMacro(MaxNumberOfModes, unsigned int);
	//@}

	/**	@name	Set the maximum variance (in percent) used to generate the model */
	//@{
	itkSetMacro(MaxVariance, double);
	itkGetMacro(MaxVariance, double);
	itkGetConstMacro(MaxVariance, double);
	//@}

	/**	@name	Set whether to use orthomax rotations after PCA */
	//@{
	itkSetMacro(UseOrthomax, bool);
	itkGetMacro(UseOrthomax, bool);
	itkGetConstMacro(UseOrthomax, bool);
	//@}

	/**	@name	Get the mean sample output */
	//@{
	VectorDecoratorPointer	GetMeanVectorOutput();
	VectorDecoratorConstPointer	GetMeanVectorOutput() const;
	//@}

	/**	@name	Get the modes of variation (MoV) output */
	//@{
	VectorContainerDecoratorPointer GetModesOfVariationOutput();
	VectorContainerDecoratorConstPointer GetModesOfVariationOutput() const;
	//@}

	/**	@name	Get MoV's weights output */
	//@{
	VectorDecoratorPointer	GetModesOfVariationWeightsOutput();
	VectorDecoratorConstPointer	GetModesOfVariationWeightsOutput() const;
	//@}

	/**	@name	Get the correlation coefficient matrix output */
	//@{
	MatrixDecoratorPointer GetCorrelationCoeffMatrixOutput();
	MatrixDecoratorConstPointer GetCorrelationCoeffMatrixOutput() const;
	//@}

	/**	@name IO */
	//@{
	void Save(std::string FileName);
	void Load(std::string FileName);
	//@}


protected:
	/// Constructor
	StatisticalModel();
	/// Destructor
	virtual ~StatisticalModel();

	/**	@name	Internal vector and matrix */
	//@{
	typedef itk::Array<double> InternalVectorType;
	typedef itk::Array2D<double> InternalMatrixType;
	//@}

	/// Creates the itk::ProcessObject's outputs for the modes of variation and its weights
	virtual DataObjectPointer MakeOutput(unsigned int idx);

	/// Data generation
	virtual void GenerateData();

	/// Compute the correlation coefficient matrix (called by GenerateData)
	virtual void GenerateCorrelationCoeffMatrix();

private:
	/// purposely not implemented
	StatisticalModel(Self&); 
	/// purposely not implemented
	void operator=(const Self&);


protected:
	/// The file name for the IO mode
	std::string m_FileName;
	/// Number of modes used to generate the model
	unsigned int m_MaxNumberOfModes;
	/// Maximum variance used to generate the model
	double m_MaxVariance;
	/// Use orthomax rotations
	bool m_UseOrthomax;

private:
	/// The output index for the mean vector
	unsigned int m_MeanVectorOutputIdx;
	/// The output index for the modes of variation
	unsigned int m_MoVOutputIdx;
	/// The output index for the modes of variation weights
	unsigned int m_MoVWeightsOutputIdx;
	/// The output index for the correlation coefficient matrix
	unsigned int m_CorrCoeffMatrixOutputIdx;
};

#endif // STATISTICALMODEL_H__
