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

#include <gsl/gsl_cblas.h>


NS_BEGIN(gdas)
using namespace mia; 
using namespace std; 

#define SUCCESS_XTOLA 1
#define SUCCESS_FTOLR 2

CGDSAMinimizer::CGDSAMinimizer(double min_step, double max_step, 
			       double xtol, double ftolr, 
			       unsigned int maxiter):
	m_min_step(min_step), 
	m_max_step(max_step), 
        m_ftolr(ftolr),
	m_xtol( xtol ), 
	m_maxiter( maxiter )
{
        add(property_gradient); 
        if (max_step <= min_step) {
                throw create_exception<invalid_argument>("CGDSAMinimizer max_step(", max_step,
                                       ") must be larger than min_step (", min_step,")"); 
        }
}

void CGDSAMinimizer::do_set_problem()
{
}

double CGDSAMinimizer::get_gradmax(CDoubleVector& x) 
{
        double gmax = 0.0; 
        for( auto ix = x.begin(); ix != x.end(); ++ix) {
                const double ax = fabs(*ix);
                if (ax > gmax) 
                        gmax = ax;
        }
        return gmax; 
}

int CGDSAMinimizer::do_run(CDoubleVector& x)
{
	TRACE_FUNCTION; 
	CDoubleVector g(x.size()); 
        CDoubleVector xwork(x.size()); 
        
	copy(x.begin(), x.end(), xwork.begin()); 
	
	double f_init = get_problem().fdf(xwork, g);
	double step = 0.5 * (m_max_step - m_min_step); 
	
	double f_old = f_init; 
	unsigned int iter = 0; 

        double gmax = get_gradmax(g); 
        int success = gmax < m_xtol ? SUCCESS_XTOLA : 0;
        int tries = 0; 
        
        while (iter++ < m_maxiter && !success) {
		cblas_daxpy(g.size(), -step / gmax, &g[0], 1, &xwork[0], 1);
                double f = get_problem().fdf(xwork, g);
                gmax = get_gradmax(g);
                if (gmax < m_xtol) 
                        success |= SUCCESS_XTOLA; 
                
                
		if (f < f_old) {
                        cvinfo() << "Successful step: [" << iter << "]: f=" << f 
                                 << ", gmax = " << gmax << ", step=" << step << "\n"; 

                        tries = 0; 
                        copy(xwork.begin(), xwork.end(), x.begin()); 
                        
                        if ( (f < 0.5 * f_old) && (step < m_max_step)) {
                                step *= 1.5; 
                                if (step > m_max_step) 
                                        step = m_max_step; 
                                cvdebug() << "Increase step size to " << step << " \n";
                        }

                        double dfrel = (f_old - f)/f_old; 
                        cvdebug() << "dfrel=" << dfrel << "\n"; 
                        if (dfrel < m_ftolr) {
                                success |= SUCCESS_FTOLR; 
                                break;
                        }
                        f_old = f; 
                } else {
                        cvinfo() << "Failed step: [" << iter << "]: f=" << f 
                                 << ", gmax = " << gmax << ", step=" << step << "\n"; 
                        
                        if (step > m_min_step) {
				
                                // restore last solution if current value larger 
				if (f > f_old) 
					copy(x.begin(), x.end(), xwork.begin()); 
				
                                step /= 2.0; 
                                if (step < m_min_step)  
                                        step = m_min_step; 
                                cvdebug() << "Decrease step size to " << step << " \n";
                        }else{
                                if (tries < 5) 
                                        ++tries;
                                else {
                                        cvmsg() << "At minimal stepsize and no improvent, stopping\n"; 
                                        break; 
                                }
                        }
                }
	}
	if (iter == m_maxiter) 
		cvmsg() << "Iteration stopped because maximum number of iterations was reached\n"; 
	
	if (success & SUCCESS_XTOLA) 
		cvmsg() << "Stop: dx below given limit\n"; 
		
	if (success & SUCCESS_FTOLR) 
		cvmsg() << "Stop: relative cost function value decrease below limit "<< m_ftolr << ".\n"; 

	return CMinimizer::success; 
}

CGDSAMinimizerPlugin::CGDSAMinimizerPlugin():
	CMinimizerPlugin("gdas"), 
	m_min_step(0.1), 
	m_max_step(2.0), 
        m_ftolr(0.0),
	m_xtol( 0.01 ), 
	m_maxiter( 200 )
{
	add_parameter("maxiter", new CUIBoundedParameter(m_maxiter, EParameterBounds::bf_min_closed, {1}, false, 
							 "Stopping criterion: the maximum number of iterations")); 
	
	add_parameter("min-step", new CDBoundedParameter(m_min_step, EParameterBounds::bf_min_open, {0.0}, false,
							 "Minimal absolute step size"));
	
	add_parameter("max-step", new CDBoundedParameter(m_max_step, EParameterBounds::bf_min_open, {0.0}, false,
							 "Maximal absolute step size"));
	
	add_parameter("xtola", new CDBoundedParameter(m_xtol, EParameterBounds::bf_min_closed, {0.0}, false, 
						      "Stop if the inf-norm of the change applied to x is below this value."));
	
	add_parameter("ftolr", new CDBoundedParameter(m_ftolr, EParameterBounds::bf_min_closed, {0.0},false, 
						      "Stop if the relative change of the criterion is below."));
}


CMinimizer *CGDSAMinimizerPlugin::do_create() const
{
	TRACE_FUNCTION; 
	return new CGDSAMinimizer(m_min_step, m_max_step, m_xtol, m_ftolr, m_maxiter);
}

const std::string CGDSAMinimizerPlugin::do_get_descr() const
{
	return "Gradient descent with automatic step size correction."; 
}

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


NS_END
