#include "ImageStitcherV2.h"


ImageStitcherV2::ImageStitcherV2()
{
	itk::MetaImageIOFactory::RegisterOneFactory();
	itk::NrrdImageIOFactory::RegisterOneFactory();
	fixedImage = CharImageType::New();
	movingImage = CharImageType::New();
	errorFile.open("errorFile.txt");
}

ImageStitcherV2::~ImageStitcherV2()
{
}

void ImageStitcherV2::SetFixedImage(unsigned char *view1Image, int *view1ImageSize, double *view1ImageSpacing)
{
	GenerateImage(view1Image, view1ImageSize, view1ImageSpacing, SETFIXED);
}

void ImageStitcherV2::SetMovingImage(unsigned char *view2Image, int *view2ImageSize, double *view2ImageSpacing)
{
	GenerateImage(view2Image, view2ImageSize, view2ImageSpacing, SETMOVING);
}

void ImageStitcherV2::GenerateImage(unsigned char *imageBuffer, int *imageSize, double *imageSpacing, IMAGESET SET)
{

	ImportFilterType::SizeType size;
	size[0] = imageSize[0];
	size[1] = imageSize[1];
	size[2] = imageSize[2];

	ImportFilterType::SpacingType spacing;
	spacing[0] = imageSpacing[0];
	spacing[1] = imageSpacing[1];
	spacing[2] = imageSpacing[2];
	ImportFilterType::IndexType start;
	ImportFilterType::IndexType::IndexValueType value;
	start.Fill(0);

	CharImageType::PointType originPoint;
	originPoint.Fill(0);
	ImportFilterType::RegionType region;
	region.SetSize(size);
	region.SetIndex(start);

	ImportFilterType::Pointer importFilter = ImportFilterType::New();
	importFilter->SetRegion(region);


	importFilter->SetOrigin(originPoint);
	importFilter->SetSpacing(spacing);

	unsigned int numberOfPixels = imageSize[0] * imageSize[1] * imageSize[2];
	const bool importImageFilterWillOwnTheBuffer = false; // setting this to true results in the original pointer being destroyed and the program crashing

	//the import filter is assigned the input volume buffer
	importFilter->SetImportPointer(imageBuffer, numberOfPixels, importImageFilterWillOwnTheBuffer);

	try
	{
		if (SET == SETFIXED)
		{
			fixedImage = importFilter->GetOutput();
			fixedImage->Update();
		}
		else
		{
			movingImage = importFilter->GetOutput();
			movingImage->Update();
		}
	}
	catch (itk::ExceptionObject &err){
		errorFile << err << std::endl;
	}
}

void ImageStitcherV2::SetDefaultImageSize()
{
	CharImageType::SizeType fixedImageSize = fixedImage->GetLargestPossibleRegion().GetSize();
	CharImageType::SizeType movingImageSize = movingImage->GetLargestPossibleRegion().GetSize();
	defaultImageSize[0] = fixedImageSize[0] > movingImageSize[0] ? fixedImageSize[0] : movingImageSize[0];
	defaultImageSize[1] = fixedImageSize[1] > movingImageSize[1] ? fixedImageSize[1] : movingImageSize[1];
	defaultImageSize[2] = fixedImageSize[2] > movingImageSize[2] ? fixedImageSize[2] : movingImageSize[2];
}

void ImageStitcherV2::WriteFixedImage(std::string folderPath)
{
	WriteImage(folderPath, SETFIXED);
}

void ImageStitcherV2::WriteMovingImage(std::string folderPath)
{
	WriteImage(folderPath, SETMOVING);
}

void ImageStitcherV2::WriteImage(std::string folderPath, enum IMAGESET SET)
{
	ImageFileWriter::Pointer writer = ImageFileWriter::New();
	std::string filename;
	switch (SET)
	{
	case SETFIXED:
		filename = "fixedImage.mha";
		writer->SetInput(fixedImage);
		break;
	case SETMOVING:
		filename = "movingImage.mha";
		writer->SetInput(movingImage);
		break;
	}
	writer->SetFileName(folderPath + filename);
	try
	{
		writer->Update();
	}
	catch (itk::ExceptionObject e){
		errorFile << e << std::endl;
	}
}

void ImageStitcherV2::WriteImage(std::string folderPath)
{
	ImageFileWriter::Pointer writer = ImageFileWriter::New();
	std::string filename = "testImage.mha";
	
//	SET HERE TO OUTPUT A TEST IMAGE

	writer->SetFileName(folderPath + filename);
	try
	{
		writer->Update();
	}
	catch (itk::ExceptionObject e){
		errorFile << e << std::endl;
	}
}

void ImageStitcherV2::ApplyDownsampleFactor(int factor)
{
	downsampleApplied = true;
	DownsampleFixedImage(factor);
	DownsampleMovingImage(factor);
}

void ImageStitcherV2::DownsampleFixedImage(int factor)
{
	DownsampleImage(SETFIXED, factor);
}

void ImageStitcherV2::DownsampleMovingImage(int factor)
{
	DownsampleImage(SETMOVING, factor);
}

void ImageStitcherV2::DownsampleImage(IMAGESET SET, int factor)
{
	ResampleFilterType::Pointer filter = ResampleFilterType::New();
	TransformType::Pointer transform = TransformType::New();
	InterpolatorType::Pointer interpolator = InterpolatorType::New();

	filter->SetInterpolator(interpolator);
	filter->SetDefaultPixelValue(0);

	CharImageType::SpacingType newSpacing;
	CharImageType::SizeType newSize;
	CharImageType::Pointer image;
	CharImageType::Pointer resampledImage = CharImageType::New();

	switch(SET)
	{
	case SETFIXED:
		image = fixedImage;
		break;
	case SETMOVING:
		image = movingImage;
		break;
	}

	newSpacing = image->GetSpacing();
	newSpacing *= factor;

	newSize = image->GetLargestPossibleRegion().GetSize();
	for (int i = 0; i < 3; i++)
	{
		newSize[i] /= factor;
	}

	filter->SetSize(newSize);
	filter->SetOutputSpacing(newSpacing);
	filter->SetInput(image);
	try
	{
		switch (SET)
		{
		case SETFIXED:
			fixedResampledImage = filter->GetOutput();
			fixedResampledImage->Update();
			break;
		
		case SETMOVING:
			movingResampledImage = filter->GetOutput();
			movingResampledImage->Update();
			break;
		}
		
	}
	catch (itk::ExceptionObject e)
	{
		errorFile << e << std::endl;
	}
}

void ImageStitcherV2::ApplyThresholdFilter(int threshold)
{
	thresholdMaskApplied = true;
	ThresholdFixedImage(threshold);
	ThresholdMovingImage(threshold);
}

void ImageStitcherV2::ThresholdFixedImage(int threshold)
{
	ThresholdFilter(SETFIXED, threshold);
}

void ImageStitcherV2::ThresholdMovingImage(int threshold)
{
	ThresholdFilter(SETMOVING, threshold);
}

void ImageStitcherV2::ThresholdFilter(IMAGESET SET, int threshold)
{
	ThresholdFilterType::Pointer thresholdFilter = ThresholdFilterType::New();
	thresholdFilter->ThresholdBelow(threshold);
	try{
		switch (SET)
		{
		case SETFIXED:
			downsampleApplied ? thresholdFilter->SetInput(fixedResampledImage) : thresholdFilter->SetInput(fixedImage);
			fixedImageMask = MaskType::New();
			fixedImageMask->SetImage(thresholdFilter->GetOutput());
			fixedImageMask->Update();
			break;
		case SETMOVING:
			downsampleApplied ? thresholdFilter->SetInput(movingResampledImage) : thresholdFilter->SetInput(movingImage);
			movingImageMask = MaskType::New();
			movingImageMask->SetImage(thresholdFilter->GetOutput());
			movingImageMask->Update();
			break;
		}
	}
	catch (itk::ExceptionObject e)
	{
		errorFile << e << std::endl;
	}
}

void ImageStitcherV2::ApplyCastFilters()
{
	fixedCastFilter = CastFilterType::New();
	fixedCastFilter->SetInput(fixedImage);
	movingCastFilter = CastFilterType::New();
	movingCastFilter->SetInput(movingImage);
}

void ImageStitcherV2::SetInitialAngle(double theta)
{
	matrix.Fill(0);
	matrix(0, 0) = cos(theta);
	matrix(2, 0) = -sin(theta);
	matrix(0, 2) = sin(theta);
	matrix(2, 2) = cos(theta);
	matrix(1, 1) = 1;

	transform->SetMatrix(matrix);
}

void ImageStitcherV2::InitializeTransform()
{
	transform = TransformType::New();
}

void ImageStitcherV2::SetCentreOfRotation()
{

	FloatImageType::SizeType size;
	FloatImageType::SpacingType spacing;

	centre = movingCastFilter->GetOutput()->GetOrigin();
	centre[0] += size[0] / 2 * spacing[0];
	//centre[1] += size[1] / 2 * spacing[1];
	centre[2] += size[2] / 2 * spacing[2];

	transform->SetCenter(centre);
}

void ImageStitcherV2::InitializeRegistration(int numberOfLevels, int shrinkFactors, int smoothingSigmas)
{
	metric = MetricType::New();
	optimizer = OptimizerType::New();
	registration = RegistrationType::New();
	initializer = TransformInitializerType::New();
	
	metric->SetFixedImage(fixedCastFilter->GetOutput());
	metric->SetMovingImage(movingCastFilter->GetOutput());
	if (thresholdMaskApplied)
	{
		metric->SetFixedImageMask(fixedImageMask);
		metric->SetMovingImageMask(movingImageMask);
	}

	registration->SetMetric(metric);
	registration->SetOptimizer(optimizer);
	initializer->SetTransform(transform);

	initializer->SetFixedImage(fixedCastFilter->GetOutput());
	initializer->SetMovingImage(movingCastFilter->GetOutput());
	registration->SetFixedImage(fixedCastFilter->GetOutput());
	registration->SetMovingImage(movingCastFilter->GetOutput());

	registration->InPlaceOn();
	initializer->MomentsOn();
	initializer->InitializeTransform();

	RegistrationType::ShrinkFactorsArrayType shrinkFactorsPerLevel;
	shrinkFactorsPerLevel.SetSize(shrinkFactors);
	shrinkFactorsPerLevel[0] = 1;
	RegistrationType::SmoothingSigmasArrayType smoothingSigmasPerLevel;
	smoothingSigmasPerLevel.SetSize(smoothingSigmas);
	smoothingSigmasPerLevel[0] = 0;
	registration->SetNumberOfLevels(numberOfLevels);
	registration->SetSmoothingSigmasPerLevel(smoothingSigmasPerLevel);
	registration->SetShrinkFactorsPerLevel(shrinkFactorsPerLevel);

	registration->SetInitialTransform(transform);

}
#ifdef POWELL
void ImageStitcherV2::SetPowellParameters(int stepLength, int maximumIterations, double angularScales, double tolerance)
{


	OptimizerType::ScalesType optimizerScales(transform->GetNumberOfParameters());
	const double translationScale = 1.0;

	// scales fall apart when increased to 1/0.3
	optimizerScales[0] = 1 / angularScales;
	optimizerScales[1] = 1 / angularScales;
	optimizerScales[2] = 1 / angularScales;
	optimizerScales[3] = translationScale;
	optimizerScales[4] = translationScale;
	optimizerScales[5] = translationScale;
	optimizer->SetScales(optimizerScales);

	optimizer->SetStepTolerance(tolerance);
	optimizer->SetValueTolerance(tolerance);
	optimizer->SetStepLength(stepLength);
	optimizer->SetMaximumIteration(maximumIterations);
}
#endif

void ImageStitcherV2::RunImageRegistration()
{
	try
	{
		registration->Update();
	}
	catch (itk::ExceptionObject e)
	{
		errorFile << e << std::endl;
	}
}