/* -*- mia-c++  -*-
 *
 * This file is part of MIA - a toolbox for medical image analysis 
 * Copyright (c) Leipzig, Madrid 1999-2017 Gert Wollny
 *
 * MIA is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with MIA; if not, see <http://www.gnu.org/licenses/>.
 *
 */

/*
  This program implements 3D gray scale image registration.
  Depending on the available plug-ins it can accomodate various models and cost-functions.
  So far input images can be given as PNG, TIF, BMP and OpenEXR (depending on the installed plug-ins)
  The input images must be of the same dimensions and gray scale (whatever bit-depth).
  The vector field will be saved as a EXR image with two frames of float values, one for the X
  and one for the Y values of the vectors.
  Other vector field output plug-ins might be supported depending on the installed plug-ins.
*/

#include <mia/core.hh>
#include <mia/3d.hh>
#include <sstream>
#include <iomanip>
#include <mia/3d/multireg.hh>
#include <boost/algorithm/minmax_element.hpp>

NS_MIA_BEGIN
using namespace std;


// an image copy operator that keeps the type of the internal images
struct FImageCopy: public TFilter<P3DImage> {
	template <typename T>
	P3DImage operator ()( const T3DImage<T>& image) const {
		return P3DImage(new T3DImage<T>(image));
	}
};


static size_t log2(size_t x)
{
	size_t result = 0;
	while (x >>= 1)
		++result;
	return result;
}

struct FMinMax: public TFilter <pair<float, float> > {
	template <typename T>
	pair<float, float> operator ()(const T3DImage<T>& image) const {
		pair<typename T3DImage<T>::const_iterator, typename T3DImage<T>::const_iterator>
			minmax = ::boost::minmax_element(image.begin(), image.end());
		return pair<float, float>(*minmax.first, *minmax.second);
	}
};

C3DMultiImageRegister::C3DMultiImageRegister(size_t start_size, size_t max_iter,
					     C3DRegModel& model, C3DRegTimeStep& time_step,
					     float outer_epsilon):
	m_start_size(start_size),
	m_max_iter(max_iter),
	m_model(model),
	m_time_step(time_step),
	m_outer_epsilon(outer_epsilon)
{
}

P3DTransformation C3DMultiImageRegister::operator () (C3DImageFatCostList& cost, P3DInterpolatorFactory ipf)
{
	TRACE("C3DMultiImageRegister::operator");


	size_t x_shift = log2(cost.get_size().x / m_start_size);
	size_t y_shift = log2(cost.get_size().y / m_start_size);
	size_t z_shift = log2(cost.get_size().z / m_start_size);


	P3DTransformation result(new C3DGridTransformation(C3DBounds(0,0,0)));

	while (x_shift || y_shift|| z_shift) {
		C3DBounds block_size(1 << x_shift, 1 << y_shift, 1 << z_shift);

		C3DImageFatCostList cost_scaled = cost.get_downscaled(block_size);

		result = result->upscale(cost_scaled.get_size());


		C3DGridTransformation& current_result = dynamic_cast<C3DGridTransformation&>(*result);
		if (m_time_step.has_regrid()) {
			reg_level_regrid(cost_scaled, current_result, *ipf);
		}else {
			reg_level(cost_scaled, current_result, *ipf);
		}

		if (x_shift)
			--x_shift;

		if (y_shift)
			--y_shift;

		if (z_shift)
			--z_shift;
	}

	result = result->upscale(cost.get_size());

	{
		C3DGridTransformation& current_result = dynamic_cast<C3DGridTransformation&>(*result);

		if (m_time_step.has_regrid()) {
			reg_level_regrid(cost, current_result, *ipf);
		}else
			reg_level(cost, current_result, *ipf);
	}

	return result;

}

void C3DMultiImageRegister::reg_level_regrid(C3DImageFatCostList& cost,
					     C3DGridTransformation& result,
					     const C3DInterpolatorFactory& /*ipf*/)
{
	TRACE("C3DMultiImageRegister::reg_level");
	double cost_value;
	size_t iter = 0;
	size_t inertia = 5;
	bool better = false;

	double initial_cost;

	C3DGridTransformation local_shift(cost.get_size());
	C3DGridTransformation best_local_shift(cost.get_size());

	// local_shift.clear();

	cost.transform(result);
	double new_cost_value = initial_cost = cost.value();

	C3DFVectorfield v(cost.get_size());
	C3DFVectorfield force(cost.get_size());


	const char lend = cverb.get_level() == vstream::ml_debug ? '\n' : '\n';
	cvmsg() << cost.get_size() << "@" << iter << "|cost:" << new_cost_value << lend;


	float best_cost_value = new_cost_value;

	float force_scale = m_model.get_force_scale();
	do {
		++iter;
		cost_value = new_cost_value;

		force.clear();
		cost.evaluate_force(force);

		for (C3DFVectorfield::iterator i = force.begin(); i != force.end(); ++i)
			*i *= force_scale;

		// solve for the force to get a velocity or deformation field
		//v.clear();
		m_model.solve(force, v);

		// obtain the maximum shift. v might be updated according to the time step model
		float maxshift = m_time_step.calculate_pertuberation(v, local_shift);
		if (maxshift <= 0) {
			cvmsg() << "zero time step, perfect registration?\n";
			break;
		}

		// get the time step based on the maximum shift
		float delta = m_time_step.get_delta(maxshift);
		cvdebug() << "maxshift = " << maxshift << " Timestep = " << delta << "\n";

		// now check, whether the new deformation would result in a
		// non-topologic distortion of the image. if so, regrid
		// some time step models might always return false
		if (m_time_step.regrid_requested(local_shift, v, delta)) {
			result += local_shift;
			fill(local_shift.field_begin(), local_shift.field_end(), C3DFVector(0.0f, 0.0f, 0.0f));
			better = false;
		}

		m_time_step.apply(v, local_shift, delta);

		C3DGridTransformation current_trans = result + local_shift;

		cost.transform(current_trans);

		new_cost_value = cost.value();

		// if the cost is reduced, we allow for some search steps
		// with increasing cost again, and we save the current best deformation
		if (new_cost_value < best_cost_value) {
			better = true;
			best_local_shift = local_shift;
			best_cost_value = new_cost_value;
			inertia = 5;
		}

		// if the cost reduction was big, increase the possible time step
		if (new_cost_value < cost_value * 0.9)
			m_time_step.increase();

		// if the cost increased, reduce the time step if possible,
		// or count down the allowed search in sub-optimal areas
		if (new_cost_value > best_cost_value)
			if (!m_time_step.decrease())
				--inertia;

		cvmsg() << iter << "@" << cost.get_size()
			<< ": best ratio=" << best_cost_value / initial_cost
			<< ", absolute cost=" << new_cost_value << "                     "
			<< lend;

	} while ( ((inertia > 0) || (new_cost_value < 0.999 * best_cost_value)) &&
		  new_cost_value / initial_cost > m_outer_epsilon && iter < m_max_iter);

	cvmsg() << '\n';

	// if re-gridding is possible, we need to accumulate teh deformation one last time

	if (better) {
		cvdebug() << "reg_level: sum the final result\n";

		result += best_local_shift;
	}

}

void C3DMultiImageRegister::reg_level(C3DImageFatCostList& cost,
                                      C3DGridTransformation& result,
                                      const C3DInterpolatorFactory& /*ipf*/)
{
	TRACE("C3DMultiImageRegister::reg_level");
	double cost_value;

	size_t iter = 0;
	size_t inertia = 5;
	bool better = false;

	double initial_cost;

	C3DGridTransformation local_shift(cost.get_size());
	C3DGridTransformation best_local_shift(cost.get_size());

	cost.transform(result);

	double new_cost_value = initial_cost = cost.value();

	C3DGridTransformation v(cost.get_size());
	C3DGridTransformation force(cost.get_size());


	const char lend = cverb.get_level() == vstream::ml_debug ? '\n' : '\r';
	cvmsg() << cost.get_size() << "@" << iter << "|cost:" << new_cost_value << lend;

	float best_cost_value = new_cost_value;

	float force_scale = m_model.get_force_scale();
	cvmsg() << "Force scale = " << force_scale << "\n";

	do {
		++iter;
		cost_value = new_cost_value;

		cost.evaluate_force(force);
		for (C3DFVectorfield::iterator i = force.field_begin(); i != force.field_end(); ++i)
			*i *= force_scale;


		// solve for the force to get a velocity or deformation field
		fill(v.field_begin(), v.field_end(), C3DFVector(0.0f, 0.0f, 0.0f));
		m_model.solve(force, v);

		// obtain the maximum shift. v might be updated according to the time step model
		float maxshift = m_time_step.calculate_pertuberation(v, local_shift);
		if (maxshift <= 0) {
			cvmsg() << "zero time step, perfect registration?\n";
			break;
		}

		// get the time step based on the maximum shift
		float delta = m_time_step.get_delta(maxshift);
		cvdebug() << "maxshift = " << maxshift << " Timestep = " << delta << "\n";

		m_time_step.apply(v, local_shift, delta);

		force = result +  local_shift;

		cost.transform(force);

		new_cost_value = cost.value();

		// if the cost is reduced, we allow for some search steps
		// with increasing cost again, and we save the current best deformation
		if (new_cost_value < best_cost_value) {
			better = true;
			best_local_shift = local_shift;
			best_cost_value = new_cost_value;
			inertia = 5;
		}

		// if the cost reduction was big, increase the possible time step
		if (new_cost_value < cost_value * 0.9)
			m_time_step.increase();

		// if the cost increased, reduce the time step if possible,
		// or count down the allowed search in sub-optimal areas
		if (new_cost_value > best_cost_value)
			if (!m_time_step.decrease())
				--inertia;

		cvmsg() << iter << "@" << cost.get_size()
			<< ": best ratio=" << best_cost_value / initial_cost
			<< ", absolute cost=" << new_cost_value << "                     "
			<< lend;

	} while ( ((inertia > 0) || (new_cost_value < 0.999 * best_cost_value)) &&
		  new_cost_value / initial_cost > m_outer_epsilon && iter < m_max_iter);

	cvmsg() << '\n';

	// if re-gridding is possible, we need to accumulate teh deformation one last time

	if (better) {
		cvdebug() << "reg_level: sum the final result\n";

		result += best_local_shift;
	}
}

NS_MIA_END
