/* -*- 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/gdsq.hh>
#include <gsl/gsl_cblas.h>


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


CGDSQMinimizer::CGDSQMinimizer(double start_step, double step_scale,
                               double xtol, double gtol, double ftolr,
                               unsigned int maxiter, double min_step):

       m_xtol( xtol ),
       m_gtol( gtol ),
       m_maxiter( maxiter ),
       m_step( start_step ),
       m_step_scale(step_scale),
       m_ftolr(ftolr),
       m_min_step(min_step)
{
       add(property_gradient);
}


void CGDSQMinimizer::do_set_problem()
{
}

int CGDSQMinimizer::test_tol(CDoubleVector& dx, double tol, int cause)const
{
       for (auto ix = dx.begin(); ix != dx.end(); ++ix) {
              if (fabs(*ix) >= tol)
                     return 0;
       }

       return cause;
}

int CGDSQMinimizer::do_run(CDoubleVector& x)
{
       TRACE_FUNCTION;
       CDoubleVector g(x.size());
       CDoubleVector gt(x.size());
       CDoubleVector xwork(x.size());
       copy(x.begin(), x.end(), xwork.begin());
       double f_init = get_problem().fdf(xwork, g);
       double step = m_step;
       double f_old = f_init;
       unsigned int iter = 0;
       int success = test_tol(g, m_xtol, SUCCESS_XTOLA);

       while (iter++ < m_maxiter && step > m_min_step && !success) {
              cblas_daxpy(g.size(), -step, &g[0], 1, &xwork[0], 1);
              double f = get_problem().fdf(xwork, gt);

              if (f < f_old) {
                     cblas_dscal(g.size(), -step, &g[0], 1);
                     double sum = cblas_ddot(g.size(), &g[0], 1, &gt[0], 1);
                     double aux = f_old - f + sum;

                     if (aux <= 0.0)
                            step *= m_step_scale;
                     else {
                            step = 0.5 * cblas_dnrm2(g.size(), &g[0], 1) / aux;
                     }

                     copy(xwork.begin(), xwork.end(), x.begin());
                     copy(gt.begin(), gt.end(), g.begin());
                     double frel = (f_old - f);
                     f_old = f;
                     success |= test_tol(gt, m_gtol, SUCCESS_GTOLA);

                     if (frel < m_ftolr)
                            success |= SUCCESS_FTOLR;
                     else
                            cvinfo() << "improovement = " << frel << "\n";
              } else {
                     // recover best solution
                     copy(x.begin(), x.end(), xwork.begin() );
                     step /= m_step_scale;
              }

              success |= test_tol(g, m_xtol / step, SUCCESS_XTOLA);
              cvinfo() << "[" << iter << "]: f=" << f << " step=" << step << "\n";
       }

       if (iter == m_maxiter)
              cvwarn() << "Iteration stopped because maximum number of iterations was reached\n";

       if (success & SUCCESS_XTOLA)
              cvmsg() << "Stop: XTOLA\n";

       if (success & SUCCESS_GTOLA)
              cvmsg() << "Stop: GTOLA\n";

       if (success & SUCCESS_FTOLR)
              cvmsg() << "Stop: FTOLR\n";

       cvmsg() << "Stop:" << success << " with " << iter << " of " << m_maxiter << "iterations\n";
       return CMinimizer::success;
}

CGDSQMinimizerPlugin::CGDSQMinimizerPlugin():
       CMinimizerPlugin("gdsq"),
       m_xtol(0.0),
       m_gtol(0.0),
       m_ftolr(0.0),
       m_maxiter(100),
       m_start_step(0.1),
       m_step_scale(2.0),
       m_min_step(1e-6)
{
       add_parameter("maxiter", new CUIBoundedParameter(m_maxiter, EParameterBounds::bf_min_closed, {1}, false,
                     "Stopping criterion: the maximum number of iterations"));
       add_parameter("step", new CDBoundedParameter(m_start_step, EParameterBounds::bf_min_open, {0.0f},
                     false, "Initial step size"));
       add_parameter("scale", new CDBoundedParameter(m_step_scale, EParameterBounds::bf_min_open, {1.0f},
                     false, "Fallback fixed step size scaling"));
       add_parameter("xtola", new CDBoundedParameter(m_xtol, EParameterBounds::bf_min_closed, {0.0}, false,
                     "Stop if the inf-norm of x-update is below this value."));
       add_parameter("gtola", new CDBoundedParameter(m_gtol, EParameterBounds::bf_min_closed, {0.0}, false,
                     "Stop if the inf-norm of the gradient 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 *CGDSQMinimizerPlugin::do_create() const
{
       TRACE_FUNCTION;
       return new CGDSQMinimizer(m_start_step, m_step_scale, m_xtol, m_gtol, m_ftolr, m_maxiter, m_min_step);
}

const std::string CGDSQMinimizerPlugin::do_get_descr() const
{
       return "Gradient descent with quadratic step estimation";
}

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




NS_END
