#pragma once
#include <itkCommand.h>
#include <itkRegularStepGradientDescentOptimizer.h>
#include <itkRegularStepGradientDescentOptimizerv4.h>

class CommandIterationUpdate : public itk::Command
{
public:
	using Self = CommandIterationUpdate;
	using Superclass = itk::Command;
	using Pointer = itk::SmartPointer<Self>;
	itkNewMacro(Self);

protected:
	CommandIterationUpdate() = default;

public:
	using OptimizerType = itk::RegularStepGradientDescentOptimizerv4<double>;
	using OptimizerPointer = const OptimizerType*;

	void
		Execute(itk::Object* caller, const itk::EventObject& event) override
	{
		Execute((const itk::Object*)caller, event);
	}

	void
		Execute(const itk::Object* object, const itk::EventObject& event) override
	{
		auto optimizer = static_cast<OptimizerPointer>(object);
		if (!(itk::IterationEvent().CheckEvent(&event)))
		{
			return;
		}
		std::cout << optimizer->GetCurrentIteration() << "   ";
		std::cout << optimizer->GetValue() << "   ";
		std::cout << optimizer->GetCurrentPosition() << "   ";
		std::cout << m_CumulativeIterationIndex++ << std::endl;
	}

private:
	unsigned int m_CumulativeIterationIndex{ 0 };
};

template <typename TRegistration>
class RegistrationInterfaceCommand : public itk::Command
{
public:
	using Self = RegistrationInterfaceCommand;
	using Superclass = itk::Command;
	using Pointer = itk::SmartPointer<Self>;
	itkNewMacro(Self);

protected:
	RegistrationInterfaceCommand() = default;

public:
	using RegistrationType = TRegistration;
	using RegistrationPointer = RegistrationType*;
	using OptimizerType = itk::RegularStepGradientDescentOptimizerv4<double>;
	using OptimizerPointer = OptimizerType*;

	void
		Execute(itk::Object* object, const itk::EventObject& event) override
	{
		// First we verify that the event invoked is of the right type,
		// \code{itk::MultiResolutionIterationEvent()}.
		// If not, we return without any further action.
		if (!(itk::MultiResolutionIterationEvent().CheckEvent(&event)))
		{
			return;
		}

		// We then convert the input object pointer to a RegistrationPointer.
		// Note that no error checking is done here to verify the
		// \code{dynamic\_cast} was successful since we know the actual object
		// is a registration method. Then we ask for the optimizer object
		// from the registration method.
		auto registration = static_cast<RegistrationPointer>(object);
		auto optimizer =
			static_cast<OptimizerPointer>(registration->GetModifiableOptimizer());

		unsigned int currentLevel = registration->GetCurrentLevel();
		typename RegistrationType::ShrinkFactorsPerDimensionContainerType
			shrinkFactors =
			registration->GetShrinkFactorsPerDimension(currentLevel);
		typename RegistrationType::SmoothingSigmasArrayType smoothingSigmas =
			registration->GetSmoothingSigmasPerLevel();

		std::cout << "-------------------------------------" << std::endl;
		std::cout << " Current level = " << currentLevel << std::endl;
		std::cout << "    shrink factor = " << shrinkFactors << std::endl;
		std::cout << "    smoothing sigma = ";
		std::cout << smoothingSigmas[currentLevel] << std::endl;
		std::cout << std::endl;

		// If this is the first resolution level we set the learning rate
		// (representing the first step size) and the minimum step length
		// (representing the convergence criterion) to large values.  At each
		// subsequent resolution level, we will reduce the minimum step length by
		// a factor of 5 in order to allow the optimizer to focus on progressively
		// smaller regions. The learning rate is set up to the current step
		// length. In this way, when the optimizer is reinitialized at the
		// beginning of the registration process for the next level, the step
		// length will simply start with the last value used for the previous
		// level. This will guarantee the continuity of the path taken by the
		// optimizer through the parameter space.
		double currStepLen0 = /*16.00*/  /*8.0*/ 16.0;
		double miniStepLen0 = /*2.5*/ /*0.04*/ 0.08;
		if (registration->GetCurrentLevel() == 0)
		{
			optimizer->SetLearningRate(currStepLen0);
			optimizer->SetMinimumStepLength(miniStepLen0);

			std::cout << " Current Step Length = " << currStepLen0
				<< ",  Minimum Step Length = " << miniStepLen0 << std::endl;
		}
		else
		{
			double currStepLen = optimizer->GetCurrentStepLength();
			double miniStepLen = optimizer->GetMinimumStepLength();
			optimizer->SetLearningRate(currStepLen);
			optimizer->SetMinimumStepLength(miniStepLen * 0.2);

			std::cout << " Current Step Length = " << currStepLen
				<< ",  Minimum Step Length = " << miniStepLen << std::endl;
		}
	}
	// Another version of the \code{Execute()} method accepting a \code{const}
	// input object is also required since this method is defined as pure
	// virtual in the base class.  This version simply returns without taking any action.
	void
		Execute(const itk::Object*, const itk::EventObject&) override
	{
		return;
	}
};

//  The following section of code implements a Command observer
//  that will monitor the configurations of the registration process
//  at every change of stage and resolution level.
template <typename TRegistration>
class RegistrationInterfaceCommand1 : public itk::Command
{
public:
	using Self = RegistrationInterfaceCommand1;
	using Superclass = itk::Command;
	using Pointer = itk::SmartPointer<Self>;
	itkNewMacro(Self);

protected:
	RegistrationInterfaceCommand1() = default;

public:
	using RegistrationType = TRegistration;

	// The Execute function simply calls another version of the \code{Execute()}
	// method accepting a \code{const} input object
	void
		Execute(itk::Object* object, const itk::EventObject& event) override
	{
		Execute((const itk::Object*)object, event);
	}

	void
		Execute(const itk::Object* object, const itk::EventObject& event) override
	{
		if (!(itk::MultiResolutionIterationEvent().CheckEvent(&event)))
		{
			return;
		}

		std::cout << "\nObserving from class " << object->GetNameOfClass();
		if (!object->GetObjectName().empty())
		{
			std::cout << " \"" << object->GetObjectName() << "\"" << std::endl;
		}

		const auto* registration = static_cast<const RegistrationType*>(object);

		unsigned int currentLevel = registration->GetCurrentLevel();
		typename RegistrationType::ShrinkFactorsPerDimensionContainerType
			shrinkFactors =
			registration->GetShrinkFactorsPerDimension(currentLevel);
		typename RegistrationType::SmoothingSigmasArrayType smoothingSigmas =
			registration->GetSmoothingSigmasPerLevel();

		std::cout << "-------------------------------------" << std::endl;
		std::cout << " Current multi-resolution level = " << currentLevel
			<< std::endl;
		std::cout << "    shrink factor = " << shrinkFactors << std::endl;
		std::cout << "    smoothing sigma = " << smoothingSigmas[currentLevel]
			<< std::endl;
		std::cout << std::endl;
	}
};

//  The following section of code implements an observer
//  that will monitor the evolution of the registration process.
class CommandIterationUpdate1 : public itk::Command
{
public:
	using Self = CommandIterationUpdate1;
	using Superclass = itk::Command;
	using Pointer = itk::SmartPointer<Self>;
	itkNewMacro(Self);

protected:
	CommandIterationUpdate1() = default;

public:
	using OptimizerType = itk::GradientDescentOptimizerv4Template<double>;
	using OptimizerPointer = const OptimizerType*;

	void
		Execute(itk::Object* caller, const itk::EventObject& event) override
	{
		Execute((const itk::Object*)caller, event);
	}

	void
		Execute(const itk::Object* object, const itk::EventObject& event) override
	{
		auto optimizer = static_cast<OptimizerPointer>(object);
		if (!(itk::IterationEvent().CheckEvent(&event)))
		{
			return;
		}
		std::cout << optimizer->GetCurrentIteration() << "   ";
		std::cout << optimizer->GetValue() << "   ";
		std::cout << optimizer->GetCurrentPosition() << "  "
			<< m_CumulativeIterationIndex++ << std::endl;
	}
private:
	unsigned int m_CumulativeIterationIndex{ 0 };
};


class CommandIterationUpdate2 : public itk::Command
{
public:
	using Self = CommandIterationUpdate2;
	using Superclass = itk::Command;
	using Pointer = itk::SmartPointer<Self>;
	itkNewMacro(Self);

protected:
	CommandIterationUpdate2() = default;

public:
	using OptimizerType = itk::RegularStepGradientDescentOptimizer;
	using OptimizerPointer = const OptimizerType*;

	void
		Execute(itk::Object* caller, const itk::EventObject& event) override
	{
		Execute((const itk::Object*)caller, event);
	}

	void
		Execute(const itk::Object* object, const itk::EventObject& event) override
	{
		auto optimizer = static_cast<OptimizerPointer>(object);
		if (!(itk::IterationEvent().CheckEvent(&event)))
		{
			return;
		}
		std::cout << optimizer->GetCurrentIteration() << "   ";
		std::cout << optimizer->GetValue() << "   ";
		std::cout << optimizer->GetCurrentPosition() << std::endl;
	}
};

template <typename TRegistration>
class RegistrationInterfaceCommand2 : public itk::Command
{
public:
	using Self = RegistrationInterfaceCommand2;
	using Superclass = itk::Command;
	using Pointer = itk::SmartPointer<Self>;
	itkNewMacro(Self);

protected:
	RegistrationInterfaceCommand2() = default;

public:
	using RegistrationType = TRegistration;
	using RegistrationPointer = RegistrationType*;
	using OptimizerType = itk::RegularStepGradientDescentOptimizer;
	using OptimizerPointer = OptimizerType*;

	void Execute(itk::Object* object, const itk::EventObject& event) override
	{
		if (!(itk::IterationEvent().CheckEvent(&event)))
		{
			return;
		}
		auto registration = static_cast<RegistrationPointer>(object);
		if (registration == nullptr)
		{
			return;
		}
		auto optimizer =
			static_cast<OptimizerPointer>(registration->GetModifiableOptimizer());

		std::cout << "-------------------------------------" << std::endl;
		std::cout << "MultiResolution Level : " << registration->GetCurrentLevel()
			<< std::endl;
		std::cout << std::endl;

		if (registration->GetCurrentLevel() == 0)
		{
			optimizer->SetMaximumStepLength(16.00);
			optimizer->SetMinimumStepLength(0.01);
		}
		else
		{
			optimizer->SetMaximumStepLength(optimizer->GetMaximumStepLength() / 4.0);
			optimizer->SetMinimumStepLength(optimizer->GetMinimumStepLength() / 10.0);
		}
	}

	void Execute(const itk::Object*, const itk::EventObject&) override
	{
		return;
	}
};