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

#define VSTREAM_DOMAIN "test-rigidregister"
#include <stdexcept>
#include <climits>

#include <mia/internal/autotest.hh>


#include <mia/3d/rigidregister.hh>
#include <mia/3d/transformfactory.hh>
#include <mia/core/spacial_kernel.hh>
#include <mia/core/minimizer.hh>
#include <mia/3d/filter.hh>
#include <mia/3d/imageio.hh>

NS_MIA_USE
using namespace std; 
namespace bfs=boost::filesystem;

class RigidRegisterFixture {
protected: 
	RigidRegisterFixture(); 
	void run(C3DTransformation& t, const std::string& minimizer_descr,  double accuracy); 
	const C3DBounds size;
}; 

void RigidRegisterFixture::run(C3DTransformation& t, const std::string& minimizer_descr, double accuracy)
{
	auto minimizer = CMinimizerPluginHandler::instance().produce(minimizer_descr); 
	P3DImageCost cost = C3DImageCostPluginHandler::instance().produce("ssd:norm=1");
	auto tr_creator = C3DTransformCreatorHandler::instance().produce(t.get_creator_string());

	C3DRigidRegister rr(cost, minimizer, tr_creator, 1);


	C3DFImage *psrc = new C3DFImage(size); 
	auto is = psrc->begin(); 
	const float hx = size.x/2.0;  
	const float hy = size.y/2.0;  
	const float hz = size.z/2.0;  
	for(size_t z = 0; z < size.z; ++z) {
		double fz = z - hz; 
		for(size_t y = 0; y < size.y; ++y) {
			double fy = 1.3 * y - size.y/2.0; 
			for(size_t x = 0; x < size.x; ++x,++is) {
				double fx = 1.5 * x - size.x/2.0; 
				*is = exp(-(fx*fx/(hx*hx) + fy*fy/(hy*hy) + fz*fz/(hz*hz))); 
			}
		}
	}
	P3DImage src(psrc);
	P3DImage ref = t(*src);

	const C3DFImage& rsrc = dynamic_cast<const C3DFImage&>(*src); 
	const C3DFImage& rref = dynamic_cast<const C3DFImage&>(*ref); 
	
	auto iss = rsrc.begin(); 
	auto irs = rref.begin(); 
	
	for(size_t z = 0; z < size.z; ++z) {
		for(size_t y = 0; y < size.y; ++y) {
			cvdebug() << "src:"; 
			for(size_t x = 0; x < size.x; ++x,++iss)
				cverb << setw(9) << *iss << " "; 
			cverb << "\n"; 
			cvdebug() << "ref:"; 
			for(size_t x = 0; x < size.x; ++x,++irs)
				cverb << setw(9) << *irs << " "; 
			cverb << "\n"; 
		}
		cverb << "\n"; 
	}
	
			

	P3DTransformation transform = rr.run(src, ref);
	auto params = transform->get_parameters();
	auto orig_params = t.get_parameters();

	BOOST_CHECK_EQUAL(params.size(), orig_params.size());

	for (size_t i = 0; i < params.size(); ++i) 
		if (fabs(orig_params[i]) > 0.01) 
			BOOST_CHECK_CLOSE(params[i], orig_params[i], accuracy);
		else
			BOOST_CHECK_CLOSE(1.0 + params[i], 1.0 + orig_params[i], 2*accuracy);

	if ( cverb.get_level() <= vstream::ml_info ) {
		save_image("src.hdr", src);
		save_image("ref.hdr", ref);
		P3DImage reg = (*transform)(*src); 
		stringstream out_name; 
		out_name << "reg-" << t.get_creator_string()
			 << "-" << minimizer->get_init_string() << ".hdr"; 
		save_image(out_name.str(), ref);
	}
}

RigidRegisterFixture::RigidRegisterFixture():
	size(20,30,30)
{
	
}

BOOST_FIXTURE_TEST_CASE( test_rigidreg_translate_nmsimplex, RigidRegisterFixture )
{

	auto tr_creator = C3DTransformCreatorHandler::instance().produce("translate");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] = 1.0;
	params[1] = 1.0;
	params[2] = 2.0;
	transformation->set_parameters(params); 

	run(*transformation, "gsl:opt=simplex,step=1.0", 1.0); 
}


BOOST_FIXTURE_TEST_CASE( test_rigid_simplex, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("rigid");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] =  1.0;
	params[1] =  3.0;
	params[2] =  1.0;
	params[3] =  0.3;
	params[4] =  0.3;
	params[5] = -0.2;
	
	transformation->set_parameters(params); 

	run(*transformation, "gsl:opt=simplex,step=1.0", 16.0); 
}

#ifdef HAVE_NLOPT

BOOST_FIXTURE_TEST_CASE( test_rigidreg_translate_gd, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("translate");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] = 1.0;
	params[1] = 1.0;
	params[2] = 2.0;
	transformation->set_parameters(params); 

	run(*transformation, "nlopt:opt=ld-var1,xtolr=0.01,ftolr=0.01", 0.8); 
}

BOOST_FIXTURE_TEST_CASE( test_rigid_ld_lbfgs, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("rigid");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] =  0.1;
	params[1] =  0.0;
	params[2] =  0.0;
	params[3] = -0.3;
	params[4] =  0.3;
	params[5] =  0.2;
	
	transformation->set_parameters(params); 

	run(*transformation, "nlopt:opt=ld-lbfgs,xtolr=0.001,ftolr=0.001", 4.0); 
}

#endif // HAVE_NLOPT

#if 0 
// the problem with tese tests is, that the images to be registered are 
// too ambigious, so that more then just one affine transformation 
// can result in a good registration - which means the parameters 
// will not be recovered. 
BOOST_FIXTURE_TEST_CASE( test_affine_simplex, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("affine");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] =  1.0;
	params[1] =  0.4;
	params[2] =  0.7;
	params[3] = -13.0;

	params[4] = -0.1;
	params[5] =  1.02;
	params[6] =  0.4;
	params[7] = -9.1;

	params[8] =  0.9;
	params[9] = -0.4;
	params[10] = 0.9;
	params[11] = -10.3;

	transformation->set_parameters(params); 

	run(*transformation, "gsl:opt=simplex,step=1.0", 1.0); 
}

BOOST_FIXTURE_TEST_CASE( test_affine_gd, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("affine");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] =  1.0;
	params[1] =  0.4;
	params[2] =  0.7;
	params[3] = -13.0;

	params[4] = -0.1;
	params[5] =  1.02;
	params[6] =  0.4;
	params[7] = -9.1;

	params[8] =  0.9;
	params[9] = -0.4;
	params[10] = 0.9;
	params[11] = -10.3;

	transformation->set_parameters(params); 

	run(*transformation, min_gd, 1.0); 
}

BOOST_FIXTURE_TEST_CASE( test_affine_bfgs, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("affine");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 

	params[0] =  1.0;
	params[1] =  0.4;
	params[2] =  0.7;
	params[3] = -13.0;

	params[4] = -0.1;
	params[5] =  1.02;
	params[6] =  0.4;
	params[7] = -9.1;

	params[8] =  0.9;
	params[9] = -0.4;
	params[10] = 0.9;
	params[11] = -10.3;

	transformation->set_parameters(params); 

	run(*transformation, min_bfgs, 1.0); 
}

BOOST_FIXTURE_TEST_CASE( test_affine_cg_fr, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("affine");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] =  1.0;
	params[1] =  0.4;
	params[2] =  0.7;
	params[3] = -13.0;

	params[4] = -0.1;
	params[5] =  1.02;
	params[6] =  0.4;
	params[7] = -9.1;

	params[8] =  0.9;
	params[9] = -0.4;
	params[10] = 0.9;
	params[11] = -10.3;

	transformation->set_parameters(params); 

	run(*transformation, min_cg_fr, 1.0); 
}

BOOST_FIXTURE_TEST_CASE( test_affine_cg_pr, RigidRegisterFixture )
{
	auto tr_creator = C3DTransformCreatorHandler::instance().produce("affine");
	auto transformation = tr_creator->create(size); 
	auto params = transformation->get_parameters(); 
	params[0] =  1.0;
	params[1] =  0.4;
	params[2] =  0.7;
	params[3] = -13.0;

	params[4] = -0.1;
	params[5] =  1.02;
	params[6] =  0.4;
	params[7] = -9.1;

	params[8] =  0.9;
	params[9] = -0.4;
	params[10] = 0.9;
	params[11] = -10.3;

	transformation->set_parameters(params); 

	run(*transformation, min_cg_pr, 1.0); 
}

#endif
