/* -*- 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/>.
 *
 */


#include <mia/internal/main.hh>
#include <mia/template/filter_chain.hh>
#include <mia/core/cmdlineparser.hh>
#include <mia/core/msgstream.hh>
#include <mia/core/filter.hh>
#include <mia/core/convergence_measure.hh>
#include <mia/core/paramarray.hh>
#include <mia/core/watch.hh>
#include <mia/2d/imageio.hh>
#include <mia/2d/filter.hh>
#include <mia/2d/transformio.hh>
#include <mia/2d/transformfactory.hh>
#include <mia/2d/vfregularizer.hh>
#include <mia/2d/cost.hh>
#include <mia/2d/deformer.hh>


NS_MIA_USE;
using namespace std;

const SProgramDescription g_description = {
        {pdi_group, "Registration, Comparison, and Transformation of 2D images"}, 
	{pdi_short, "Non-linear registration of 2D images by using a diffeomorphic SyN registration"}, 
	{pdi_description, "This program implements the registration of two gray scale 2D images. "
         "The transformation applied is a symmeric diffeomorpic fluid dynamic registration. "
         "To work most efficiently, this program makes a few assumptions about the imput data, i.e. "
         "the image must be of the same size, have the same voxel spacing, and any intensity "
         "range normalization or equalization should also be done before calling this program." 
        }, 
	{pdi_example_descr, "Register image test.v to image ref.v saving the transformation to transform.v "
         "and the inverse transform to inverse-transform.v. "
	 "Use three multiresolution levels, ssd as image cost function."}, 
	{pdi_example_code, "-i test.v -r ref.v -t transform.v -T inverse-transform.v  -l 3 ssd"}
};

void check_number_of_levels_consistency(size_t nvalues, size_t levels, const std::string& value_name)
{
        if (nvalues != 1 && nvalues != levels) {
                throw create_exception<invalid_argument>(levels, " ", 
                                                         value_name, "  values given, but ", nvalues, 
                                                         " multiresolution levels requested"); 
        }
}


typedef pair<P2DFVectorfield, P2DFVectorfield> TransformPair;

struct C2DSymScaledRegisterParams {
	uint32_t mg_levels; 
	P2DImageCost cost; 
        float current_step;
	P2DVectorfieldRegularizer regularizer;
	TPerLevelScalarParam<uint32_t> conv_count;
	TPerLevelScalarParam<double> stop_decline_rate;
	TPerLevelScalarParam<double> stop_cost;
        TPerLevelScalarParam<unsigned> iterations; 
	
	C2DSymScaledRegisterParams();

	void add_options(CCmdOptionList& options);
}; 

C2DSymScaledRegisterParams::C2DSymScaledRegisterParams():
	mg_levels(3), 
        current_step(0.25), 
	conv_count(10), 
	stop_decline_rate(0.1), 
	stop_cost(0.1), 
        iterations(100)
{
}

void C2DSymScaledRegisterParams::add_options(CCmdOptionList& options)
{
	options.set_group("Registration");
	options.add(make_opt( mg_levels, "mg-levels", 'l', "Number of multi-resolution levels to run the "
			      "registration on. Not that some registration parameters can be given as a "
			      "coma seperated list to indicate per level values. In these cases  if the "
			      "number of given values is smaller than the number of multi-resolution "
			      "levels (this parameter), the the last given value is used for all subsequest "
			      "multi-resolution levels."));

	options.add(make_opt(current_step,  EParameterBounds::bf_min_open | EParameterBounds::bf_max_closed,
			     {0.0f, 0.5f},  "step", 'S', "Initial step size for all levels")); 
	
	options.add(make_opt( cost, "ssd", "cost", 'c', "Image similarity function to be minimized"));
	options.add(make_opt( regularizer, "sor:kernel=fluid,maxiter=50", "regularizer", 'R', 
			      "Regularization for the force to transformation update")); 


	options.add(conv_count.
		    create_level_params_option("conv-test-interval",'T', EParameterBounds::bf_closed_interval, {4,40}, 
					       "Convergence test interations intervall: In order to measure "
					       "convergence the cost function value is averaged over this "
					       "amount of iterations, and the decline rate is evaluated based "
					       "on the linare regression of the cost function values in this "
					       "intervall. This parameter can be given as a coma-seperated "
					       "list with values corresponding to the multi-resolution levels, "
					       "see option --mg-levels for more information."
			    ));
	options.add(stop_decline_rate.
		    create_level_params_option("stop-decline-rate", 'D', EParameterBounds::bf_min_closed, {0},
					       "Stopping criterium for registration based on the cost decline rate. "
					       "If the rate below this value, the iteration is stopped. "
					       "This parameter can be given as a coma-seperated list with values "
					       "corresponding to the multi-resolution levels, see option "
					       "--mg-levels for more information."
			    ));
	
	options.add(stop_cost.
		    create_level_params_option("stop-cost", 'C',
					       "Stopping criterium for registration based on the cost value. "
					       "If the cost drops below this value, the iteration is stopped. "
					       "This parameter can be given as a coma-seperated list with values "
					       "corresponding to the multi-resolution levels, see option "
					       "--mg-levels for more information."
			    ));
	
	options.add(iterations.
		    create_level_params_option("iter", 'I', EParameterBounds::bf_min_closed, {4},
					       "Naximum number if iterations done on each multi-resolution level. "
					       "This parameter can be given as a coma-seperated list with values "
					       "corresponding to the multi-resolution levels, see option "
					       "--mg-levels for more information."));
}


class C2DSymFluidRegistration {
public: 
        C2DSymFluidRegistration(const C2DSymScaledRegisterParams& params);
        
        TransformPair run(const C2DFImage& src, const C2DFImage& ref) const; 
private:
	const C2DSymScaledRegisterParams& m_params;
}; 

P2DTransformation wrap_vectorfield_in_transformation(const C2DFVectorfield& field, 
                                                     const C2DTransformCreator& vftcreator) 
{
        
	// the final transformation is twice the distance 
	auto transform = vftcreator.create(field.get_size()); 
	CDoubleVector buffer(transform->degrees_of_freedom(), false);

	double max_transform = 0.0;
	C2DFVector max_vec; 
	
	auto ib = buffer.begin(); 
	for (auto ivf = field.begin(); ivf != field.end(); ++ivf) {
		cvdebug() << *ivf << "\n"; 
		*ib++ = 2.0f * ivf->x; 
		*ib++ = 2.0f * ivf->y; 

		double vm = ivf->norm();
		if (vm > max_transform) {
			max_transform = vm;
			max_vec = *ivf; 
		}
	}

	cvmsg() << "Max-transform = " << max_vec << " norm=" << max_transform << "\n"; 
	transform->set_parameters(buffer); 
        return transform; 
}

C2DFImage get_asfloat_pixel(P2DImage image) 
{
        if (image->get_pixel_type() != it_float)
                image = run_filters(image, "convert:map=copy,repn=float");
        return dynamic_cast<const C2DFImage&>(*image); 
}

int do_main( int argc, char *argv[] )
{
        string src_filename;
	string ref_filename;
        
	string out_transform_filename;
	string out_inv_transform_filename;


        const auto& imageio = C2DImageIOPluginHandler::instance();
        const auto& transio = C2DTransformationIOPluginHandler::instance();

        CCmdOptionList options(g_description);
        options.set_group("IO"); 
	options.add(make_opt( src_filename, "in-image", 'i', "test image", 
			      CCmdOptionFlags::required_input, &imageio));
	options.add(make_opt( ref_filename, "ref-image", 'r', "reference image", 
			      CCmdOptionFlags::required_input, &imageio));
	options.add(make_opt( out_transform_filename, "transform", 'o', "output transformation", 
			      CCmdOptionFlags::required_output, &transio));
	options.add(make_opt( out_inv_transform_filename, "inverse-transform", 'O', 
                              "inverse output transformation", 
			      CCmdOptionFlags::required_output, &transio));

	C2DSymScaledRegisterParams params; 
	params.add_options(options); 
        
                              
        if (options.parse(argc, argv) != CCmdOptionList::hr_no)
		return EXIT_SUCCESS; 

        
        P2DImage src = load_image2d(src_filename); 
        P2DImage ref = load_image2d(ref_filename); 
        
        C2DSymFluidRegistration registration(params); 
        
        auto result = registration.run(get_asfloat_pixel(src), get_asfloat_pixel(ref));
        
        auto vftranscreator  = produce_2dtransform_factory("vf:imgkernel=[bspline:d=1],imgboundary=zero");
        
        P2DTransformation t_src_ref = wrap_vectorfield_in_transformation(*result.first, *vftranscreator); 
        
        if (!C2DTransformationIOPluginHandler::instance().save(out_transform_filename, *t_src_ref)) 
		throw create_exception<runtime_error>( "Unable to save transformation to '", 
                                                       out_transform_filename, "'"); 
        
        P2DTransformation t_ref_src = wrap_vectorfield_in_transformation(*result.second, *vftranscreator);
        
        if (!C2DTransformationIOPluginHandler::instance().save(out_inv_transform_filename, *t_ref_src)) 
		throw create_exception<runtime_error>( "Unable to save transformation to '", 
                                                       out_inv_transform_filename, "'"); 
        return EXIT_SUCCESS; 
}


MIA_MAIN(do_main); 


C2DSymFluidRegistration::C2DSymFluidRegistration(const C2DSymScaledRegisterParams& params):
	m_params(params)
{
}

P2DFVectorfield upscale( const C2DFVectorfield& vf, C2DBounds size)
{
	P2DFVectorfield Result(new C2DFVectorfield(size));
	float x_mult = float(size.x) / (float)vf.get_size().x;
	float y_mult = float(size.y) / (float)vf.get_size().y;
	float ix_mult = 1.0f / x_mult;
	float iy_mult = 1.0f / y_mult;

	double max_defo = 0.0; 
	auto callback = [&](const C1DParallelRange& range){
		//	CThreadMsgStream thread_stream;
		
		for (auto y = range.begin(); y != range.end();++y) { 
			auto i = Result->begin_at(0,y);
			for (unsigned int x = 0; x < size.x; x++,++i){
				C2DFVector help(ix_mult * x, iy_mult * y);
				C2DFVector val = vf.get_interpol_val_at(help);
				*i = C2DFVector(val.x * x_mult,val.y * y_mult);
				
				double defo = i->norm();
				if (max_defo < defo)
					max_defo = defo; 
			}
		}
	}; 
	pfor(C1DParallelRange(0, size.y, 1), callback);
	// race condition !! 
	cvmsg() << "max-defo at size (" << size << ")=" << max_defo << "\n"; 
	
	return Result;
}



class C2DSymScaledRegister {
public: 
        C2DSymScaledRegister(unsigned level, const C2DSymScaledRegisterParams& params); 
        void run (const C2DFImage& src, const C2DFImage& ref, TransformPair& transforms) const; 
private: 
        void deform(const C2DFImage& src, const C2DFVectorfield& t, C2DFImage& result) const;

	unsigned m_level;
	const C2DSymScaledRegisterParams& m_params;
	C2DInterpolatorFactory m_ipfac; 
        
}; 

TransformPair
C2DSymFluidRegistration::run(const C2DFImage& src, const C2DFImage& ref) const 
{
	const CWatch& watch = CWatch::instance();

	double start_time_global = watch.get_seconds(); 

	TransformPair transforms;
        
        for (unsigned l = 0; l < m_params.mg_levels; ++l) {
		double start_time_level = watch.get_seconds(); 
                unsigned scale_block = 1 << (m_params.mg_levels - l - 1);
                
                stringstream downscale_descr;
                downscale_descr << "downscale:bx=" << scale_block
                                << ",by=" << scale_block;
                
                auto downscaler =
                        C2DFilterPluginHandler::instance().produce(downscale_descr.str().c_str());
                
                auto scaled_src_p = downscaler->filter(src); 
                auto scaled_ref_p = downscaler->filter(ref); 

                const C2DFImage& scaled_src = dynamic_cast<const C2DFImage&>(*scaled_src_p); 
                const C2DFImage& scaled_ref = dynamic_cast<const C2DFImage&>(*scaled_ref_p); 

                if (!transforms.first) {
                        transforms.first.reset(new C2DFVectorfield(scaled_src.get_size()));
                        transforms.second.reset(new C2DFVectorfield(scaled_src.get_size()));
                }else{
                        transforms.first = upscale(*transforms.first, scaled_src.get_size());
                        transforms.second = upscale(*transforms.second, scaled_src.get_size());
                }
                
                C2DSymScaledRegister level_worker(l, m_params); 
                
                level_worker.run(scaled_src, scaled_ref, transforms);
		cvmsg() << "Level " << l << " time=" << watch.get_seconds() - start_time_level << "\n"; 
                
        }
	cvmsg() << "Global run time " << watch.get_seconds() - start_time_global << "\n"; 
        return transforms; 
}

C2DSymScaledRegister::C2DSymScaledRegister(unsigned level, const C2DSymScaledRegisterParams& params):
	m_level(level),
	m_params(params),
	m_ipfac("bspline:d=1", "zero")
{

}

void C2DSymScaledRegister::run (const C2DFImage& src, const C2DFImage& ref, TransformPair& transforms) const
{
	TransformPair transforms_best = transforms; 
	cvmsg() << "Enter registration at level "<< m_level << " with size " << src.get_size() << "\n";
	cvmsg() << "  Running at most " << m_params.iterations[m_level] << " iterations\n"; 

	auto current_step = m_params.current_step; 
	CConvergenceMeasure conv_measure(m_params.conv_count[m_level]);

	m_params.regularizer->set_size(src.get_size()); 
        
        // deform the images with the input transformations
        C2DFImage src_tmp(src.get_size()); 
        deform(src, *transforms.first, src_tmp); 
        
        C2DFImage ref_tmp(ref.get_size()); 
        deform(ref, *transforms.second, ref_tmp); 

        // now this runs ping-pong 
        unsigned iter = 0; 

	C2DFVectorfield grad(src.get_size()); 
        C2DFVectorfield v(src.get_size()); 

	double decline_rate = numeric_limits<double>::max();
	double avg_cost = numeric_limits<double>::max();

	double old_cost = numeric_limits<double>::max();
	
        while (!conv_measure.is_full_size() || (
	       decline_rate > m_params.stop_decline_rate[m_level] &&
	       avg_cost > m_params.stop_cost[m_level] &&
	       iter < m_params.iterations[m_level] &&
	       current_step > 0.0001)) {
		
		double start_time_level = CWatch::instance().get_seconds(); 
                ++iter; 

		m_params.cost->set_reference(ref_tmp);
		double cost_fw = m_params.cost->evaluate_force(src_tmp, grad); 
		float max_v_fw = m_params.regularizer->run(v, grad, *transforms.first); 
		transforms.first->update_by_velocity(v, current_step / max_v_fw);
		transforms.second->update_as_inverse_of(*transforms.first, 1e-5, 20);
		
                deform(src, *transforms.first, src_tmp); 
                deform(ref, *transforms.second, ref_tmp);
		
		m_params.cost->set_reference(src_tmp);
		double cost_bw = m_params.cost->evaluate_force(ref_tmp, grad);
		double cost_val = cost_bw + cost_fw; 

                // update step length 
		if ( cost_val <= old_cost ) {
			
			transforms_best = transforms;

			old_cost = cost_val;
			conv_measure.push(cost_val);
			
			decline_rate = - conv_measure.rate();
			avg_cost = conv_measure.value();
			
			cvmsg() << "[" << setw(4) << iter << "]:"
				<< "cost = "<< cost_fw + cost_bw
				<< ", step = " << current_step
				<< ", cost_avg(n="<< conv_measure.fill()<< ")=" << avg_cost
				<< ", rate=" <<  decline_rate
				<< ", ct=" << CWatch::instance().get_seconds() - start_time_level << "s"
				<< "\n";
			
		}
		
		if ( cost_val < 0.9 * old_cost ) {
			if (current_step < 0.25) {
				current_step *= 1.25;
				if (current_step > 0.25)
					current_step = 0.25;
			}
		}else if ( cost_val > old_cost ) {
			current_step *= 0.5;
			cvmsg() << "Increasing cost = " << cost_val
				<< ", Discard step and retry with step size" << current_step <<"\n";

			// reset transform and re-evaluate gradient 
			transforms = transforms_best;
			deform(src, *transforms.first, src_tmp); 
			deform(ref, *transforms.second, ref_tmp);
			m_params.cost->set_reference(src_tmp);
			m_params.cost->evaluate_force(ref_tmp, grad);
		}
				
		float max_v_bw = m_params.regularizer->run(v, grad, *transforms.second); 
		
		transforms.second->update_by_velocity(v, current_step / max_v_bw);  
		transforms.first->update_as_inverse_of(*transforms.second, 1e-5, 20);
                
                deform(src, *transforms.first, src_tmp); 
                deform(ref, *transforms.second, ref_tmp);
        }
	transforms = transforms_best; 
}

void C2DSymScaledRegister::deform(const C2DFImage& src, const C2DFVectorfield& t, C2DFImage& result) const
{
        FDeformer2D src_deformer(t, m_ipfac); 
        src_deformer(src,result); 
}

