/* -*- 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 <memory>
#include <mia/core/minimizer/gsl.hh>

NS_BEGIN(gsl)

using namespace std;
using namespace mia; 


const TDictMap<CGSLMinimizerPlugin::EGSLOptimizer>::Table minimizer_table[] = {
	{"simplex", CGSLMinimizerPlugin::opt_simplex2, "Simplex algorithm of Nelder and Mead"},
	{"cg-fr",   CGSLMinimizerPlugin::opt_cg_fr, "Flecher-Reeves conjugate gradient algorithm"},
	{"cg-pr",   CGSLMinimizerPlugin::opt_cg_pr, "Polak-Ribiere conjugate gradient algorithm"},
	{"bfgs",    CGSLMinimizerPlugin::opt_bfgs, "Broyden-Fletcher-Goldfarb-Shann"},
	{"bfgs2",   CGSLMinimizerPlugin::opt_bfgs2, "Broyden-Fletcher-Goldfarb-Shann (most efficient version)"},
	{"gd",      CGSLMinimizerPlugin::opt_gd, "Gradient descent."},
	{NULL, CGSLMinimizerPlugin::opt_unknown, ""}
};


CGSLFDFMinimizer::CGSLFDFMinimizer(const gsl_multimin_fdfminimizer_type *ot, 
				   double gorth_tolerance, 
				   double stop_eps, 
				   unsigned int maxiter, 
				   float start_step):
	m_ot(ot), 
	m_s(NULL), 
	m_gorth_tolerance(gorth_tolerance), 
	m_stop_eps(stop_eps), 
	m_maxiter(maxiter), 
	m_start_step(start_step)
{
	add(property_gradient); 

	m_func.n = 0; 
	m_func.f = nullptr; 
	m_func.df = nullptr; 
	m_func.fdf = nullptr; 
	m_func.params = nullptr;
	
}

CGSLFDFMinimizer::~CGSLFDFMinimizer()
{
	if (m_s) 
		gsl_multimin_fdfminimizer_free (m_s);	
}

double CGSLFDFMinimizer::f(const gsl_vector * x, void * obj)
{
	assert(x); 
	assert(obj); 
	
	CMinimizer::Problem *o = (Problem *)obj; 
	return o->f(x->size, x->data); 
}

void CGSLFDFMinimizer::df(const gsl_vector * x, void * obj, gsl_vector * g)
{
	assert(x); 
	assert(obj); 
	assert(g); 
	assert(x->size == g->size); 
	
	CMinimizer::Problem *o = (Problem *)obj; 
	o->df(x->size, x->data, g->data); 
}

void CGSLFDFMinimizer::fdf(const gsl_vector * x, void * obj, double * f, gsl_vector * g)
{
	assert(x); 
	assert(obj); 
	assert(g); 
	assert(x->size == g->size); 
	
	CMinimizer::Problem *o = (Problem *)obj; 
	*f = o->fdf(x->size, x->data, g->data); 
}


void CGSLFDFMinimizer::do_set_problem()
{
	m_func.n = size(); 
	m_func.f = CGSLFDFMinimizer::f; 
	m_func.df = CGSLFDFMinimizer::df; 
	m_func.fdf = CGSLFDFMinimizer::fdf; 
	m_func.params = (void *)get_problem_pointer();

	cvinfo() << "Set problem size to " << size() << " parameters\n"; 
	if (m_s) 
		gsl_multimin_fdfminimizer_free (m_s);
	m_s = gsl_multimin_fdfminimizer_alloc (m_ot, size()); 	
	if (!m_s) 
		throw std::runtime_error("CFDFMinimizer:Not enough memory to allocate the minimizer"); 
}

struct DeallocGSLVector {
	void operator ()(gsl_vector*v) {
		gsl_vector_free(v); 
	}
}; 


int CGSLFDFMinimizer::do_run(CDoubleVector& x)
{
	int iter = 0; 
	int status = GSL_CONTINUE; 
	shared_ptr<gsl_vector> init_x(gsl_vector_alloc(x.size()),DeallocGSLVector()); 
	copy(x.begin(), x.end(), init_x->data);  
	cvinfo() << "Start GSL optimizer with " << x.size() << " parameters\n"; 
	gsl_multimin_fdfminimizer_set (m_s, &m_func, init_x.get(), m_start_step, m_gorth_tolerance);
	do {
		++iter; 
		status = gsl_multimin_fdfminimizer_iterate (m_s);
		if (status) 
			break; 
		
		status = gsl_multimin_test_gradient (m_s->gradient, m_stop_eps);
	} while (status == GSL_CONTINUE && iter < m_maxiter); 
	
	// copy best solution 
	gsl_vector * help = gsl_multimin_fdfminimizer_x (m_s); 
	std::copy(help->data, help->data + size(), x.begin()); 
	
	cvmsg() << "\n"; 
	if (status) {
		if (iter == m_maxiter)  
			cvinfo() << "Maximum number of iterations reached\n"; 
		else 
			cvmsg() << "Iteration stopped because '" << gsl_strerror(status) << "'\n";
	}else{
		cvinfo() << "Convergence reached\n"; 
	}
	
	return status == GSL_SUCCESS ? CMinimizer::success : CMinimizer::failure;
}

CGSLFMinimizer::CGSLFMinimizer(const gsl_multimin_fminimizer_type *ot,
			       double stop_eps, 
			       unsigned int maxiter, double start_step):
	m_ot(ot), 
	m_s(NULL),
	m_stop_eps(stop_eps), 
	m_maxiter(maxiter), 
	m_start_step(start_step), 
	m_step_init(NULL)
{
	m_func.n = 0; 
	m_func.f = nullptr; 
	m_func.params = nullptr;
}

CGSLFMinimizer::~CGSLFMinimizer()
{
	gsl_multimin_fminimizer_free (m_s);
	gsl_vector_free(m_step_init); 	
}

double CGSLFMinimizer::f(const gsl_vector * x, void *obj)
{
	assert(x); 
	assert(obj); 
	
	CMinimizer::Problem *o = reinterpret_cast<Problem *>(obj); 
	return o->f(x->size, x->data); 
}


void CGSLFMinimizer::do_set_problem()
{
	m_func.n = size(); 
	m_func.f = CGSLFDFMinimizer::f; 
	m_func.params = (void *)get_problem_pointer();
	
	if (m_s) 
		gsl_multimin_fminimizer_free (m_s);
	m_s = gsl_multimin_fminimizer_alloc (m_ot, size()); 	
	if (!m_s) 
		throw std::runtime_error("CFMinimizer:Not enough memory to allocate the minimizer"); 
	
	if (m_step_init) {
		if (m_step_init->size != size()) {
			gsl_vector_free(m_step_init); 
			m_step_init= gsl_vector_alloc(size());
		}
	}else 
		m_step_init= gsl_vector_alloc(size());
	
}

int CGSLFMinimizer::do_run(CDoubleVector& x)
{
	unsigned iter = 0; 
	int status;  
	shared_ptr<gsl_vector> init_x(gsl_vector_alloc(x.size()), DeallocGSLVector()); 
	copy(x.begin(), x.end(), init_x->data);  
	
	
	gsl_vector_set_all (m_step_init, m_start_step);
	gsl_multimin_fminimizer_set (m_s, &m_func, init_x.get(), m_step_init);
	do {
		++iter; 
		status = gsl_multimin_fminimizer_iterate (m_s);
		if (status) 
			break; 
		const double size = gsl_multimin_fminimizer_size (m_s);
		status = gsl_multimin_test_size (size, m_stop_eps);
		
	} while (status == GSL_CONTINUE && iter < m_maxiter); 
	
	// copy best solution 
	gsl_vector * help = gsl_multimin_fminimizer_x (m_s); 
	std::copy(help->data, help->data + size(), x.begin()); 
	if (status) {
		if (iter == m_maxiter) 
			cvmsg() << "\nMaximum number of iterations reached\n"; 
		else 
			cvwarn() << "\nIteration stopped because '" << gsl_strerror(status) << "\n";
	}else{
		cvmsg() << "\nConvergence reached\n"; 
	}
	
	return status == GSL_SUCCESS ? CMinimizer::success : CMinimizer::failure;
}



typedef CDictParameter<CGSLMinimizerPlugin::EGSLOptimizer> CGSLMinimizerParam; 
typedef TDictMap<CGSLMinimizerPlugin::EGSLOptimizer> CGSLMinimizerDict;  

CGSLMinimizerPlugin::CGSLMinimizerPlugin():
	CMinimizerPlugin("gsl"), 
	m_ot(opt_gd), 
	m_gorth_tolerance(0.1),  
	m_stop_eps(0.01),  
	m_maxiter(100), 
	m_start_step(0.001)
	
{
	add_parameter("opt", new CGSLMinimizerParam(m_ot, CGSLMinimizerDict(minimizer_table), 
						    "Specific optimizer to be used."));
	
	add_parameter("tol", new CDBoundedParameter(m_gorth_tolerance, EParameterBounds::bf_min_open, {0.0f}, false, 
						    "some tolerance parameter"));
	
	add_parameter("eps", new CDBoundedParameter(m_stop_eps, EParameterBounds::bf_min_open, {0.0f}, 
						    false, 
						    "gradient based optimizers: stop when |grad| < eps, "
						    "simplex: stop when simplex size < eps."));
	
	add_parameter("iter", new CUIBoundedParameter(m_maxiter, EParameterBounds::bf_min_closed, {1},
						      false, "maximum number of iterations"));
	
	add_parameter("step", new CDBoundedParameter(m_start_step, EParameterBounds::bf_min_open, {0.0f},  false, 
						     "initial step size"));
}						

const std::string CGSLMinimizerPlugin::do_get_descr() const
{
	return "optimizer plugin based on the multimin optimizers of "
		"the GNU Scientific Library (GSL) https://www.gnu.org/software/gsl/"; 
}

CMinimizer *CGSLMinimizerPlugin::do_create() const
{
	const gsl_multimin_fdfminimizer_type *ot; 
	
	switch (m_ot) {
	case opt_simplex2 : return new CGSLFMinimizer(gsl_multimin_fminimizer_nmsimplex2, 
							     m_stop_eps, m_maxiter, m_start_step);
	case opt_cg_fr: ot = gsl_multimin_fdfminimizer_conjugate_fr; break;   
	case opt_cg_pr: ot = gsl_multimin_fdfminimizer_conjugate_pr; break;   
	case opt_bfgs: ot = gsl_multimin_fdfminimizer_vector_bfgs; break;   
	case opt_bfgs2: ot = gsl_multimin_fdfminimizer_vector_bfgs2; break;   
	case opt_gd: ot = gsl_multimin_fdfminimizer_steepest_descent; break;  
	default:
		throw invalid_argument("Unknown GSL optimizer given"); 
	}
	return new CGSLFDFMinimizer(ot, m_gorth_tolerance, 
				    m_stop_eps, m_maxiter, 
				    m_start_step);
}

extern "C" EXPORT CPluginBase *get_plugin_interface()
{
	return new CGSLMinimizerPlugin();
}


NS_END
