// -*- coding: utf-8
// vim: set fileencoding=utf-8

// This file is part of Eigen, a lightweight C++ template library
// for linear algebra.
//
// Copyright (C) 2009 Thomas Capricelli <orzel@freehackers.org>
//
// Eigen is free software; you can redistribute it and/or
// modify it under the terms of the GNU Lesser General Public
// License as published by the Free Software Foundation; either
// version 3 of the License, or (at your option) any later version.
//
// Alternatively, 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 2 of
// the License, or (at your option) any later version.
//
// Eigen 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 Lesser General Public License or the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU Lesser General Public
// License and a copy of the GNU General Public License along with
// Eigen. If not, see <http://www.gnu.org/licenses/>.

#ifndef EIGEN_LEVENBERGMARQUARDT__H
#define EIGEN_LEVENBERGMARQUARDT__H


namespace LevenbergMarquardtSpace {
    enum Status {
        NotStarted = -2,
        Running = -1,
        ImproperInputParameters = 0,
        RelativeReductionTooSmall = 1,
        RelativeErrorTooSmall = 2,
        RelativeErrorAndReductionTooSmall = 3,
        CosinusTooSmall = 4,
        TooManyFunctionEvaluation = 5,
        FtolTooSmall = 6,
        XtolTooSmall = 7,
        GtolTooSmall = 8,
        UserAsked = 9
    };
}


/**
  * \ingroup NonLinearOptimization_Module
  * \brief Performs non linear optimization over a non-linear function,
  * using a variant of the Levenberg Marquardt algorithm.
  *
  * Check wikipedia for more information.
  * http://en.wikipedia.org/wiki/Levenberg%E2%80%93Marquardt_algorithm
  */
template<typename FunctorType, typename Scalar=double>
class LevenbergMarquardt
{
public:
    LevenbergMarquardt(FunctorType &_functor)
        : functor(_functor) { nfev = njev = iter = 0;  fnorm = gnorm = 0.; useExternalScaling=false; }

    struct Parameters {
        Parameters()
            : factor(Scalar(100.))
            , maxfev(400)
            , ftol(ei_sqrt(NumTraits<Scalar>::epsilon()))
            , xtol(ei_sqrt(NumTraits<Scalar>::epsilon()))
            , gtol(Scalar(0.))
            , epsfcn(Scalar(0.)) {}
        Scalar factor;
        int maxfev;   // maximum number of function evaluation
        Scalar ftol;
        Scalar xtol;
        Scalar gtol;
        Scalar epsfcn;
    };

    typedef Matrix< Scalar, Dynamic, 1 > FVectorType;
    typedef Matrix< Scalar, Dynamic, Dynamic > JacobianType;

    LevenbergMarquardtSpace::Status lmder1(
            FVectorType &x,
            const Scalar tol = ei_sqrt(NumTraits<Scalar>::epsilon())
            );

    LevenbergMarquardtSpace::Status minimize(FVectorType &x);
    LevenbergMarquardtSpace::Status minimizeInit(FVectorType &x);
    LevenbergMarquardtSpace::Status minimizeOneStep(FVectorType &x);

    static LevenbergMarquardtSpace::Status lmdif1(
            FunctorType &functor,
            FVectorType &x,
            int *nfev,
            const Scalar tol = ei_sqrt(NumTraits<Scalar>::epsilon())
            );

    LevenbergMarquardtSpace::Status lmstr1(
            FVectorType  &x,
            const Scalar tol = ei_sqrt(NumTraits<Scalar>::epsilon())
            );

    LevenbergMarquardtSpace::Status minimizeOptimumStorage(FVectorType  &x);
    LevenbergMarquardtSpace::Status minimizeOptimumStorageInit(FVectorType  &x);
    LevenbergMarquardtSpace::Status minimizeOptimumStorageOneStep(FVectorType  &x);

    void resetParameters(void) { parameters = Parameters(); }

    Parameters parameters;
    FVectorType  fvec, qtf, diag;
    JacobianType fjac;
    PermutationMatrix<Dynamic,Dynamic> permutation;
    int nfev;
    int njev;
    int iter;
    Scalar fnorm, gnorm;
    bool useExternalScaling; 

    Scalar lm_param(void) { return par; }
private:
    FunctorType &functor;
    int n;
    int m;
    FVectorType wa1, wa2, wa3, wa4;

    Scalar par, sum;
    Scalar temp, temp1, temp2;
    Scalar delta;
    Scalar ratio;
    Scalar pnorm, xnorm, fnorm1, actred, dirder, prered;
};

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::lmder1(
        FVectorType  &x,
        const Scalar tol
        )
{
    n = x.size();
    m = functor.values();

    /* check the input parameters for errors. */
    if (n <= 0 || m < n || tol < 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    resetParameters();
    parameters.ftol = tol;
    parameters.xtol = tol;
    parameters.maxfev = 100*(n+1);

    return minimize(x);
}


template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::minimize(FVectorType  &x)
{
    LevenbergMarquardtSpace::Status status = minimizeInit(x);
    if (status==LevenbergMarquardtSpace::ImproperInputParameters)
        return status;
    do {
        status = minimizeOneStep(x);
    } while (status==LevenbergMarquardtSpace::Running);
    return status;
}

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::minimizeInit(FVectorType  &x)
{
    n = x.size();
    m = functor.values();

    wa1.resize(n); wa2.resize(n); wa3.resize(n);
    wa4.resize(m);
    fvec.resize(m);
    fjac.resize(m, n);
    if (!useExternalScaling)
        diag.resize(n);
    assert( (!useExternalScaling || diag.size()==n) || "When useExternalScaling is set, the caller must provide a valid 'diag'");
    qtf.resize(n);

    /* Function Body */
    nfev = 0;
    njev = 0;

    /*     check the input parameters for errors. */
    if (n <= 0 || m < n || parameters.ftol < 0. || parameters.xtol < 0. || parameters.gtol < 0. || parameters.maxfev <= 0 || parameters.factor <= 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    if (useExternalScaling)
        for (int j = 0; j < n; ++j)
            if (diag[j] <= 0.)
                return LevenbergMarquardtSpace::ImproperInputParameters;

    /*     evaluate the function at the starting point */
    /*     and calculate its norm. */
    nfev = 1;
    if ( functor(x, fvec) < 0)
        return LevenbergMarquardtSpace::UserAsked;
    fnorm = fvec.stableNorm();

    /*     initialize levenberg-marquardt parameter and iteration counter. */
    par = 0.;
    iter = 1;

    return LevenbergMarquardtSpace::NotStarted;
}

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::minimizeOneStep(FVectorType  &x)
{
    assert(x.size()==n); // check the caller is not cheating us

    /* calculate the jacobian matrix. */
    int df_ret = functor.df(x, fjac);
    if (df_ret<0)
        return LevenbergMarquardtSpace::UserAsked;
    if (df_ret>0)
        // numerical diff, we evaluated the function df_ret times
        nfev += df_ret;
    else njev++;

    /* compute the qr factorization of the jacobian. */
    wa2 = fjac.colwise().blueNorm();
    ColPivHouseholderQR<JacobianType> qrfac(fjac);
    fjac = qrfac.matrixQR();
    permutation = qrfac.colsPermutation();

    /* on the first iteration and if external scaling is not used, scale according */
    /* to the norms of the columns of the initial jacobian. */
    if (iter == 1) {
        if (!useExternalScaling)
            for (int j = 0; j < n; ++j)
                diag[j] = (wa2[j]==0.)? 1. : wa2[j];

        /* on the first iteration, calculate the norm of the scaled x */
        /* and initialize the step bound delta. */
        xnorm = diag.cwiseProduct(x).stableNorm();
        delta = parameters.factor * xnorm;
        if (delta == 0.)
            delta = parameters.factor;
    }

    /* form (q transpose)*fvec and store the first n components in */
    /* qtf. */
    wa4 = fvec;
    wa4.applyOnTheLeft(qrfac.householderQ().adjoint());
    qtf = wa4.head(n);

    /* compute the norm of the scaled gradient. */
    gnorm = 0.;
    if (fnorm != 0.)
        for (int j = 0; j < n; ++j)
            if (wa2[permutation.indices()[j]] != 0.)
                gnorm = std::max(gnorm, ei_abs( fjac.col(j).head(j+1).dot(qtf.head(j+1)/fnorm) / wa2[permutation.indices()[j]]));

    /* test for convergence of the gradient norm. */
    if (gnorm <= parameters.gtol)
        return LevenbergMarquardtSpace::CosinusTooSmall;

    /* rescale if necessary. */
    if (!useExternalScaling)
        diag = diag.cwiseMax(wa2);

    do {

        /* determine the levenberg-marquardt parameter. */
        ei_lmpar2<Scalar>(qrfac, diag, qtf, delta, par, wa1);

        /* store the direction p and x + p. calculate the norm of p. */
        wa1 = -wa1;
        wa2 = x + wa1;
        pnorm = diag.cwiseProduct(wa1).stableNorm();

        /* on the first iteration, adjust the initial step bound. */
        if (iter == 1)
            delta = std::min(delta,pnorm);

        /* evaluate the function at x + p and calculate its norm. */
        if ( functor(wa2, wa4) < 0)
            return LevenbergMarquardtSpace::UserAsked;
        ++nfev;
        fnorm1 = wa4.stableNorm();

        /* compute the scaled actual reduction. */
        actred = -1.;
        if (Scalar(.1) * fnorm1 < fnorm)
            actred = 1. - ei_abs2(fnorm1 / fnorm);

        /* compute the scaled predicted reduction and */
        /* the scaled directional derivative. */
        wa3 = fjac.template triangularView<Upper>() * (qrfac.colsPermutation().inverse() *wa1);
        temp1 = ei_abs2(wa3.stableNorm() / fnorm);
        temp2 = ei_abs2(ei_sqrt(par) * pnorm / fnorm);
        prered = temp1 + temp2 / Scalar(.5);
        dirder = -(temp1 + temp2);

        /* compute the ratio of the actual to the predicted */
        /* reduction. */
        ratio = 0.;
        if (prered != 0.)
            ratio = actred / prered;

        /* update the step bound. */
        if (ratio <= Scalar(.25)) {
            if (actred >= 0.)
                temp = Scalar(.5);
            if (actred < 0.)
                temp = Scalar(.5) * dirder / (dirder + Scalar(.5) * actred);
            if (Scalar(.1) * fnorm1 >= fnorm || temp < Scalar(.1))
                temp = Scalar(.1);
            /* Computing MIN */
            delta = temp * std::min(delta, pnorm / Scalar(.1));
            par /= temp;
        } else if (!(par != 0. && ratio < Scalar(.75))) {
            delta = pnorm / Scalar(.5);
            par = Scalar(.5) * par;
        }

        /* test for successful iteration. */
        if (ratio >= Scalar(1e-4)) {
            /* successful iteration. update x, fvec, and their norms. */
            x = wa2;
            wa2 = diag.cwiseProduct(x);
            fvec = wa4;
            xnorm = wa2.stableNorm();
            fnorm = fnorm1;
            ++iter;
        }

        /* tests for convergence. */
        if (ei_abs(actred) <= parameters.ftol && prered <= parameters.ftol && Scalar(.5) * ratio <= 1. && delta <= parameters.xtol * xnorm)
            return LevenbergMarquardtSpace::RelativeErrorAndReductionTooSmall;
        if (ei_abs(actred) <= parameters.ftol && prered <= parameters.ftol && Scalar(.5) * ratio <= 1.)
            return LevenbergMarquardtSpace::RelativeReductionTooSmall;
        if (delta <= parameters.xtol * xnorm)
            return LevenbergMarquardtSpace::RelativeErrorTooSmall;

        /* tests for termination and stringent tolerances. */
        if (nfev >= parameters.maxfev)
            return LevenbergMarquardtSpace::TooManyFunctionEvaluation;
        if (ei_abs(actred) <= NumTraits<Scalar>::epsilon() && prered <= NumTraits<Scalar>::epsilon() && Scalar(.5) * ratio <= 1.)
            return LevenbergMarquardtSpace::FtolTooSmall;
        if (delta <= NumTraits<Scalar>::epsilon() * xnorm)
            return LevenbergMarquardtSpace::XtolTooSmall;
        if (gnorm <= NumTraits<Scalar>::epsilon())
            return LevenbergMarquardtSpace::GtolTooSmall;

    } while (ratio < Scalar(1e-4));

    return LevenbergMarquardtSpace::Running;
}

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::lmstr1(
        FVectorType  &x,
        const Scalar tol
        )
{
    n = x.size();
    m = functor.values();

    /* check the input parameters for errors. */
    if (n <= 0 || m < n || tol < 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    resetParameters();
    parameters.ftol = tol;
    parameters.xtol = tol;
    parameters.maxfev = 100*(n+1);

    return minimizeOptimumStorage(x);
}

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::minimizeOptimumStorageInit(FVectorType  &x)
{
    n = x.size();
    m = functor.values();

    wa1.resize(n); wa2.resize(n); wa3.resize(n);
    wa4.resize(m);
    fvec.resize(m);
    // Only R is stored in fjac. Q is only used to compute 'qtf', which is
    // Q.transpose()*rhs. qtf will be updated using givens rotation,
    // instead of storing them in Q.
    // The purpose it to only use a nxn matrix, instead of mxn here, so
    // that we can handle cases where m>>n :
    fjac.resize(n, n);
    if (!useExternalScaling)
        diag.resize(n);
    assert( (!useExternalScaling || diag.size()==n) || "When useExternalScaling is set, the caller must provide a valid 'diag'");
    qtf.resize(n);

    /* Function Body */
    nfev = 0;
    njev = 0;

    /*     check the input parameters for errors. */
    if (n <= 0 || m < n || parameters.ftol < 0. || parameters.xtol < 0. || parameters.gtol < 0. || parameters.maxfev <= 0 || parameters.factor <= 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    if (useExternalScaling)
        for (int j = 0; j < n; ++j)
            if (diag[j] <= 0.)
                return LevenbergMarquardtSpace::ImproperInputParameters;

    /*     evaluate the function at the starting point */
    /*     and calculate its norm. */
    nfev = 1;
    if ( functor(x, fvec) < 0)
        return LevenbergMarquardtSpace::UserAsked;
    fnorm = fvec.stableNorm();

    /*     initialize levenberg-marquardt parameter and iteration counter. */
    par = 0.;
    iter = 1;

    return LevenbergMarquardtSpace::NotStarted;
}


template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::minimizeOptimumStorageOneStep(FVectorType  &x)
{
    assert(x.size()==n); // check the caller is not cheating us

    int i, j;
    bool sing;

    /* compute the qr factorization of the jacobian matrix */
    /* calculated one row at a time, while simultaneously */
    /* forming (q transpose)*fvec and storing the first */
    /* n components in qtf. */
    qtf.fill(0.);
    fjac.fill(0.);
    int rownb = 2;
    for (i = 0; i < m; ++i) {
        if (functor.df(x, wa3, rownb) < 0) return LevenbergMarquardtSpace::UserAsked;
        ei_rwupdt<Scalar>(fjac, wa3, qtf, fvec[i]);
        ++rownb;
    }
    ++njev;

    /* if the jacobian is rank deficient, call qrfac to */
    /* reorder its columns and update the components of qtf. */
    sing = false;
    for (j = 0; j < n; ++j) {
        if (fjac(j,j) == 0.)
            sing = true;
        wa2[j] = fjac.col(j).head(j).stableNorm();
    }
    permutation.setIdentity(n);
    if (sing) {
        wa2 = fjac.colwise().blueNorm();
        // TODO We have no unit test covering this code path, do not modify
        // until it is carefully tested
        ColPivHouseholderQR<JacobianType> qrfac(fjac);
        fjac = qrfac.matrixQR();
        wa1 = fjac.diagonal();
        fjac.diagonal() = qrfac.hCoeffs();
        permutation = qrfac.colsPermutation();
        // TODO : avoid this:
        for(int ii=0; ii< fjac.cols(); ii++) fjac.col(ii).segment(ii+1, fjac.rows()-ii-1) *= fjac(ii,ii); // rescale vectors

        for (j = 0; j < n; ++j) {
            if (fjac(j,j) != 0.) {
                sum = 0.;
                for (i = j; i < n; ++i)
                    sum += fjac(i,j) * qtf[i];
                temp = -sum / fjac(j,j);
                for (i = j; i < n; ++i)
                    qtf[i] += fjac(i,j) * temp;
            }
            fjac(j,j) = wa1[j];
        }
    }

    /* on the first iteration and if external scaling is not used, scale according */
    /* to the norms of the columns of the initial jacobian. */
    if (iter == 1) {
        if (!useExternalScaling)
            for (j = 0; j < n; ++j)
                diag[j] = (wa2[j]==0.)? 1. : wa2[j];

        /* on the first iteration, calculate the norm of the scaled x */
        /* and initialize the step bound delta. */
        xnorm = diag.cwiseProduct(x).stableNorm();
        delta = parameters.factor * xnorm;
        if (delta == 0.)
            delta = parameters.factor;
    }

    /* compute the norm of the scaled gradient. */
    gnorm = 0.;
    if (fnorm != 0.)
        for (j = 0; j < n; ++j)
            if (wa2[permutation.indices()[j]] != 0.)
                gnorm = std::max(gnorm, ei_abs( fjac.col(j).head(j+1).dot(qtf.head(j+1)/fnorm) / wa2[permutation.indices()[j]]));

    /* test for convergence of the gradient norm. */
    if (gnorm <= parameters.gtol)
        return LevenbergMarquardtSpace::CosinusTooSmall;

    /* rescale if necessary. */
    if (!useExternalScaling)
        diag = diag.cwiseMax(wa2);

    do {

        /* determine the levenberg-marquardt parameter. */
        ei_lmpar<Scalar>(fjac, permutation.indices(), diag, qtf, delta, par, wa1);

        /* store the direction p and x + p. calculate the norm of p. */
        wa1 = -wa1;
        wa2 = x + wa1;
        pnorm = diag.cwiseProduct(wa1).stableNorm();

        /* on the first iteration, adjust the initial step bound. */
        if (iter == 1)
            delta = std::min(delta,pnorm);

        /* evaluate the function at x + p and calculate its norm. */
        if ( functor(wa2, wa4) < 0)
            return LevenbergMarquardtSpace::UserAsked;
        ++nfev;
        fnorm1 = wa4.stableNorm();

        /* compute the scaled actual reduction. */
        actred = -1.;
        if (Scalar(.1) * fnorm1 < fnorm)
            actred = 1. - ei_abs2(fnorm1 / fnorm);

        /* compute the scaled predicted reduction and */
        /* the scaled directional derivative. */
        wa3 = fjac.corner(TopLeft,n,n).template triangularView<Upper>() * (permutation.inverse() * wa1);
        temp1 = ei_abs2(wa3.stableNorm() / fnorm);
        temp2 = ei_abs2(ei_sqrt(par) * pnorm / fnorm);
        prered = temp1 + temp2 / Scalar(.5);
        dirder = -(temp1 + temp2);

        /* compute the ratio of the actual to the predicted */
        /* reduction. */
        ratio = 0.;
        if (prered != 0.)
            ratio = actred / prered;

        /* update the step bound. */
        if (ratio <= Scalar(.25)) {
            if (actred >= 0.)
                temp = Scalar(.5);
            if (actred < 0.)
                temp = Scalar(.5) * dirder / (dirder + Scalar(.5) * actred);
            if (Scalar(.1) * fnorm1 >= fnorm || temp < Scalar(.1))
                temp = Scalar(.1);
            /* Computing MIN */
            delta = temp * std::min(delta, pnorm / Scalar(.1));
            par /= temp;
        } else if (!(par != 0. && ratio < Scalar(.75))) {
            delta = pnorm / Scalar(.5);
            par = Scalar(.5) * par;
        }

        /* test for successful iteration. */
        if (ratio >= Scalar(1e-4)) {
            /* successful iteration. update x, fvec, and their norms. */
            x = wa2;
            wa2 = diag.cwiseProduct(x);
            fvec = wa4;
            xnorm = wa2.stableNorm();
            fnorm = fnorm1;
            ++iter;
        }

        /* tests for convergence. */
        if (ei_abs(actred) <= parameters.ftol && prered <= parameters.ftol && Scalar(.5) * ratio <= 1. && delta <= parameters.xtol * xnorm)
            return LevenbergMarquardtSpace::RelativeErrorAndReductionTooSmall;
        if (ei_abs(actred) <= parameters.ftol && prered <= parameters.ftol && Scalar(.5) * ratio <= 1.)
            return LevenbergMarquardtSpace::RelativeReductionTooSmall;
        if (delta <= parameters.xtol * xnorm)
            return LevenbergMarquardtSpace::RelativeErrorTooSmall;

        /* tests for termination and stringent tolerances. */
        if (nfev >= parameters.maxfev)
            return LevenbergMarquardtSpace::TooManyFunctionEvaluation;
        if (ei_abs(actred) <= NumTraits<Scalar>::epsilon() && prered <= NumTraits<Scalar>::epsilon() && Scalar(.5) * ratio <= 1.)
            return LevenbergMarquardtSpace::FtolTooSmall;
        if (delta <= NumTraits<Scalar>::epsilon() * xnorm)
            return LevenbergMarquardtSpace::XtolTooSmall;
        if (gnorm <= NumTraits<Scalar>::epsilon())
            return LevenbergMarquardtSpace::GtolTooSmall;

    } while (ratio < Scalar(1e-4));

    return LevenbergMarquardtSpace::Running;
}

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::minimizeOptimumStorage(FVectorType  &x)
{
    LevenbergMarquardtSpace::Status status = minimizeOptimumStorageInit(x);
    if (status==LevenbergMarquardtSpace::ImproperInputParameters)
        return status;
    do {
        status = minimizeOptimumStorageOneStep(x);
    } while (status==LevenbergMarquardtSpace::Running);
    return status;
}

template<typename FunctorType, typename Scalar>
LevenbergMarquardtSpace::Status
LevenbergMarquardt<FunctorType,Scalar>::lmdif1(
        FunctorType &functor,
        FVectorType  &x,
        int *nfev,
        const Scalar tol
        )
{
    int n = x.size();
    int m = functor.values();

    /* check the input parameters for errors. */
    if (n <= 0 || m < n || tol < 0.)
        return LevenbergMarquardtSpace::ImproperInputParameters;

    NumericalDiff<FunctorType> numDiff(functor);
    // embedded LevenbergMarquardt
    LevenbergMarquardt<NumericalDiff<FunctorType> > lm(numDiff);
    lm.parameters.ftol = tol;
    lm.parameters.xtol = tol;
    lm.parameters.maxfev = 200*(n+1);

    LevenbergMarquardtSpace::Status info = LevenbergMarquardtSpace::Status(lm.minimize(x));
    if (nfev)
        * nfev = lm.nfev;
    return info;
}

//vim: ai ts=4 sts=4 et sw=4
#endif // EIGEN_LEVENBERGMARQUARDT__H