#pragma once
#define _USE_MATH_DEFINES
#include <iostream>
#include <cmath>

// libraries to generate images
#include "itkImage.h"
#include "itkImportImageFilter.h"
#include "itkResampleImageFilter.h"
#include "itkThresholdImageFilter.h"
#include "itkImageMaskSpatialObject.h"


// libraries to perform registration
#include "itkVersorRigid3DTransform.h"
#include "itkNearestNeighborInterpolateImageFunction.h"
#include "itkCorrelationImageToImageMetricv4.h"
#include "itkPowellOptimizerv4.h"
#include "itkImageRegistrationMethodv4.h"
#include "itkCenteredTransformInitializer.h"
#include "itkVersorRigid3DTransformOptimizer.h"

// libraries to write images
#include "itkImageFileWriter.h"
#include "itkMetaImageIOFactory.h"
#include "itkNrrdImageIOFactory.h"

enum IMAGESET {SETFIXED, SETMOVING};
const unsigned int Dimension = 3;
typedef itk::Image<unsigned char, Dimension>													CharImageType;
typedef itk::Image<float, Dimension>															FloatImageType;
typedef itk::ImportImageFilter<unsigned char, Dimension>										ImportFilterType;
typedef itk::ResampleImageFilter<CharImageType, CharImageType>									ResampleFilterType;
typedef itk::ThresholdImageFilter<CharImageType>												ThresholdFilterType;
typedef itk::ImageMaskSpatialObject< Dimension >												MaskType;
typedef itk::CastImageFilter< CharImageType, FloatImageType >									CastFilterType;

typedef itk::VersorRigid3DTransform<double>														TransformType;
typedef itk::CenteredTransformInitializer<TransformType, FloatImageType, FloatImageType>		TransformInitializerType;
typedef itk::NearestNeighborInterpolateImageFunction<CharImageType, double>						InterpolatorType;
typedef itk::CorrelationImageToImageMetricv4<FloatImageType, FloatImageType>					MetricType;
#define POWELL
#ifdef POWELL
typedef itk::PowellOptimizerv4<double>															OptimizerType;
#endif
//#define VERSOR_OPTIMIZER
#ifdef VERSOR_OPTIMIZER
typedef itk::VersorRigid3DTransformOptimizer													OptimizerType;
#endif
typedef itk::ImageRegistrationMethodv4<FloatImageType, FloatImageType, TransformType>			RegistrationType;

typedef itk::ImageFileWriter<CharImageType>														ImageFileWriter;

class ImageStitcherV2
{
public:
	ImageStitcherV2();
	~ImageStitcherV2();

	void GenerateImage(unsigned char *view1Image, int *view1ImageSize, double *view1ImageSpacing, enum IMAGESET);
	void SetFixedImage(unsigned char *view1Image, int *view1ImageSize, double *view1ImageSpacing);
	void SetMovingImage(unsigned char *view2Image, int *view2ImageSize, double *view2ImageSpacing);
	void SetDefaultImageSize();

	typedef itk::VersorRigid3DTransform<double>														TransformType;

	void WriteImage(std::string folderPath);														// NOTE: this overload for testing purposes
	void WriteImage(std::string folderPath, enum IMAGESET);
	void WriteFixedImage(std::string folderPath);
	void WriteMovingImage(std::string folderPath);


	void ApplyDownsampleFactor(int factor);
	void DownsampleFixedImage(int factor);
	void DownsampleMovingImage(int factor);
	void DownsampleImage(IMAGESET SET, int factor);

	void ApplyThresholdFilter(int threshold);
	void ThresholdFixedImage(int threshold);
	void ThresholdMovingImage(int threshold);
	void ThresholdFilter(IMAGESET SET, int threshold);

	void ApplyCastFilters();
	void InitializeTransform();
	void SetInitialAngle(double theta);
	void SetCentreOfRotation();

	void InitializeRegistration(int numberOfLevels, int shrinkFactorsPerLevel, int smoothingSigmasPerLevel);
	void SetPowellParameters(int stepLength, int maximumIterations, double angularScales, double tolerance);
#ifdef VERSOR_OPTIMIZER
	void ImageStitcherV2::SetVersorParameters(double angularScales);
#endif

	void RunImageRegistration();

private:
	CharImageType::Pointer																		fixedImage;
	CharImageType::Pointer																		movingImage;
	CharImageType::SizeType																		defaultImageSize;
	CastFilterType::Pointer																		fixedCastFilter;
	CastFilterType::Pointer																		movingCastFilter;

	bool downsampleApplied = false;
	CharImageType::Pointer																		fixedResampledImage;
	CharImageType::Pointer																		movingResampledImage;

	bool thresholdMaskApplied = false;
	MaskType::Pointer																			fixedImageMask;
	MaskType::Pointer																			movingImageMask;

	// Registration Components
	MetricType::Pointer																			metric;
	OptimizerType::Pointer																		optimizer;
	RegistrationType::Pointer																	registration;
	TransformType::Pointer																		transform;
	TransformInitializerType::Pointer															initializer;
	TransformType::InputPointType																centre;
	TransformType::MatrixType																	matrix;
	TransformType::TranslationType																translation;

	TransformType::ConstPointer																	finalTransform;

	std::ofstream errorFile;
};