/*
 * This file is part of the GIRAFFE Pipeline
 * Copyright (C) 2002-2019 European Southern Observatory
 *
 * This program 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 2 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 this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#ifdef HAVE_CONFIG_H
#  include <config.h>
#endif

#include <math.h>

#include <cxtypes.h>
#include <cxmemory.h>

#include "gimath.h"
#include "gilevenberg.h"


/**
 * @defgroup gilevenberg Levenberg-Marquardt Fit Utilities
 *
 * TBD
 */

/**@{*/

inline static void
_giraffe_swap(cxdouble *a, cxdouble *b) {

    register cxdouble t = *a;

    *a = *b;
    *b = t;

    return;

}


inline static void
_giraffe_covsrt(cpl_matrix *covar, cxint ma, cxint ia[], cxint mfit)
{

    register cxint i, j, k;

    cxint nr = cpl_matrix_get_nrow(covar);

    cxdouble *_covar = cpl_matrix_get_data(covar);


    for (i = mfit; i < ma; i++) {
        for (j = 0; j <= i; j++) {
            _covar[i * nr + j] = _covar[j * nr + i] = 0.0;
        }
    }

    k = mfit - 1;

    for (j = (ma - 1); j >= 0; j--) {
        if (ia[j]) {
            for (i = 0; i < ma; i++) {
                _giraffe_swap(&_covar[i * nr + k], &_covar[i * nr + j]);
            }

            for (i = 0;i < ma; i++) {
                _giraffe_swap(&_covar[k * nr + i], &_covar[j * nr + i]);
            }

            k--;
        }
    }

}


/*
 * @brief
 *   LMRQ Chi Square Calculation
 *
 * @param x          -  X abcissa [ndata]
 * @param y          -  Y values [ndata]
 * @param sig        -  Y sigmas [ndata]
 * @param ndata      -  Number of values
 * @param a          -  Initial guesses for model parameters [ma]
 * @param r          -  Maximum deltat for modelparameters [ma]
 * @param ia         -  Flags for model parameters to be fitted [ma]
 * @param ma         -  Number of parameters to fit
 * @param alpha      -  Working space [ma,ma]
 * @param beta       -  Working space [ma,ma]
 * @param chisq      -  Chi Square value of fit
 * @param funcs      -  Non linear model to fit
 *
 * @return =0 if succesful, <0 if error an occured
 *
 * Used by @c giraffe_mrqmin() to evaluate the linearized fitting
 * matrix @a alpha and vector @a beta and calculate chi squared @a chisq.
 *
 * @see giraffe_mrqmin()
 */

inline static cxint
_giraffe_mrqcof(cpl_matrix  *x, cpl_matrix  *y, cpl_matrix  *sig,
                cxint ndata, cpl_matrix  *a, cxdouble r[], cxint ia[],
                cxint ma, cpl_matrix *alpha, cpl_matrix *beta,
                cxdouble *chisq, GiFitFunc funcs)
{

    register cxint i, j, k, l, m;
    register cxint mfit = 0;

    cxint nr_alpha = cpl_matrix_get_nrow(alpha);
    cxint nc_x = cpl_matrix_get_ncol(x);

    cxdouble ymod;
    cxdouble wt;
    cxdouble sig2i;
    cxdouble dy;
    cxdouble *dyda;
    cxdouble *pd_x     = cpl_matrix_get_data(x);
    cxdouble *pd_y     = cpl_matrix_get_data(y);
    cxdouble *pd_sig   = cpl_matrix_get_data(sig);
    cxdouble *pd_a     = cpl_matrix_get_data(a);
    cxdouble *pd_alpha = cpl_matrix_get_data(alpha);
    cxdouble *pd_beta  = cpl_matrix_get_data(beta);


    for (j = 0; j < ma; j++) {
        if (ia[j]) {
            mfit++;
        }
    }

    for (j = 0; j < mfit; j++) {
        for (k = 0; k <= j; k++) {
            pd_alpha[j * nr_alpha + k] = 0.0;
        }

        pd_beta[j] = 0.0;
    }

    *chisq = 0.0;

    dyda = cx_calloc(ma, sizeof(cxdouble));

    for (i = 0; i < ndata; i++) {

        (*funcs)(&ymod, &(pd_x[i * nc_x]), pd_a, ma, dyda, r);

        if (pd_sig[i] == 0.0) {
            continue;
        }

        sig2i = 1.0 / (pd_sig[i] * pd_sig[i]);
        dy = pd_y[i] - ymod;

        for (j = -1, l = 0; l < ma; l++) {

           if (ia[l]) {
                wt = dyda[l] * sig2i;
                for (j++, k = -1, m = 0; m <= l; m++) {
                    if (ia[m]) {
                        ++k;
                        pd_alpha[j * nr_alpha + k] += (wt * dyda[m]);
                    }
                }

                pd_beta[j] += (dy * wt);

            }
        }

        *chisq += (dy * dy * sig2i);

    }

    for (j = 1; j < mfit; j++) {
        for (k = 0; k < j; k++) {
            pd_alpha[k * nr_alpha + j] = pd_alpha[j * nr_alpha + k];
        }
    }


    cx_free(dyda);

    return 0;

}


/*
 * @brief
 *   Levenberg-Marquardt non-linear fit routine
 *
 * @param x          -  X abcissa [ndata]
 * @param y          -  Y values [ndata]
 * @param sig        -  Y sigmas [ndata]
 * @param ndata      -  Number of values
 * @param a          -  Initial guesses for model parameters [ma]
 * @param r          -  Maximum delta for model parameters [ma]
 * @param ia         -  Flags fot model parameters to be fitted [ma]
 * @param ma         -  Number of parameters to fit
 * @param covar      -  Covariance matrix [ma,ma]
 * @param alpha      -  Working space [ma,ma]
 * @param chisq      -  Chi Square of fit
 * @param funcs      -  Non linear model to fit
 * @param alamda     -  Control parameter of fit
 *
 * @return  0 if succesful, < 0 if an error occured
 *
 * Levenberg-Marquardt non linear fit method, based upon attempting to
 * reduce the value @em CHISQ of a fit between a set of data points
 * @a x[1..ndata], @a y[1..ndata] with individual standard deviations
 * @a sig[1..ndata], and a nonlinear function @a funcs dependent on
 * @a ma coefficients @a a[1..ma].
 * @par Fit Control Parameters:
 * The input array @a a[1..ma] contains initial guesses for the parameters
 * to be fitted.
 * The input array @a ia[1..ma] indicates by nonzero entries those components
 * of @a a[1..ma] that should be fitted for, and by zero entries those
 * components that should be held fixed at their input values.
 *
 * The program returns current best-fit values for the parameters @a a[1..ma],
 * and @em CHISQ=chisq. The arrays @a covar[1..ma][1..ma] and
 * @a alpha[1..ma][1..ma] are used as working space during most iterations.
 *
 * Supply a routine @a funcs(x,a,yfit,dyda,ma) that evaluates the fitting
 * function yfit, and its derivatives @em dyda[1..ma] with respect to the
 * fitting parameters @a a at @a x. On the first call provide an initial
 * guess for the parameters @a a, and set @a alamda<0  for initialization
 * (which then sets @a alamda=.001). If a step succeeds @a chisq becomes
 * smaller and @a alamda decreases by a factor of 10. If a step fails
 * @a alamda grows by a factor of 10.
 *
 * You @em must call this routine repeatedly until convergence is achieved.
 * Then, make one final call with @a alamda=0, so that @a covar[1..ma][1..ma]
 * returns the covariance matrix, and @a alpha[1..ma][1..ma] the
 * curvature matrix.
 *
 * Parameters held fixed will return zero covariances.
 *
 * @see _giraffe_mrqcof()
 *
 */

static cxint
_giraffe_mrqmin(cpl_matrix *x, cpl_matrix *y, cpl_matrix *sig, cxint ndata,
                cpl_matrix *a, cxdouble r[], cxint ia[], cxint ma,
                cpl_matrix *covar, cpl_matrix *alpha, cxdouble *chisq,
                GiFitFunc funcs, cxdouble *alamda)
{

    register cxint gj, j, k, l, m;

    static cxint nr_covar, nr_alpha, nr_moneda, mfit;

    static cxdouble *pd_a, *pd_covar, *pd_alpha;
    static cxdouble *atry, *beta, *da, *oneda, ochisq;

    static cpl_matrix *matry, *mbeta, *mda, *moneda;


    pd_a     = cpl_matrix_get_data(a);
    pd_covar = cpl_matrix_get_data(covar);
    pd_alpha = cpl_matrix_get_data(alpha);
    nr_covar = cpl_matrix_get_nrow(covar);
    nr_alpha = cpl_matrix_get_nrow(alpha);

    if (*alamda < 0.0) {

        matry = cpl_matrix_new(ma, 1);
        atry  = cpl_matrix_get_data(matry);

        mbeta = cpl_matrix_new(ma, 1);
        beta  = cpl_matrix_get_data(mbeta);

        mda = cpl_matrix_new(ma, 1);
        da  = cpl_matrix_get_data(mda);

        for (mfit = 0, j = 0; j < ma; j++) {
            if (ia[j]) {
                mfit++;
            }
        }

        moneda = cpl_matrix_new(1, mfit);
        oneda  = cpl_matrix_get_data(moneda);

        *alamda = 0.001;

        gj = _giraffe_mrqcof(x, y, sig, ndata, a, r, ia, ma, alpha, mbeta,
                             chisq, funcs);

        if (gj != 0) {
            cpl_matrix_delete(moneda);
            moneda = NULL;
            oneda = NULL;

            cpl_matrix_delete(mda);
            mda = NULL;
            da = NULL;

            cpl_matrix_delete(mbeta);
            mbeta = NULL;
            beta = NULL;

            cpl_matrix_delete(matry);
            matry = NULL;
            atry = NULL;

            return gj;
        }

        ochisq = (*chisq);

        for (j = 0; j < ma; j++) {
            atry[j] = pd_a[j];
        }

    }

    nr_moneda = cpl_matrix_get_nrow(moneda);

    for (j = -1, l = 0; l < ma; l++) {
        if (ia[l]) {
            for (j++, k = -1, m = 0; m < ma; m++) {
                if (ia[m]) {
                    k++;
                    pd_covar[j * nr_covar + k] = pd_alpha[j * nr_alpha + k];
                }
            }

            pd_covar[j * nr_covar + j] = pd_alpha[j * nr_alpha + j] *
                (1.0 + (*alamda));

            oneda[j * nr_moneda + 0] = beta[j];
        }
    }

    gj = giraffe_gauss_jordan(covar, mfit, moneda, 1);

    if (gj != 0) {
        cpl_matrix_delete(moneda);
        moneda = NULL;
        oneda = NULL;

        cpl_matrix_delete(mda);
        mda = NULL;
        da = NULL;

        cpl_matrix_delete(mbeta);
        mbeta = NULL;
        beta = NULL;

        cpl_matrix_delete(matry);
        matry = NULL;
        atry = NULL;

        return gj;
    }

    for (j = 0; j < mfit; j++) {
        da[j] = oneda[j * nr_moneda + 0];
    }

    if (*alamda == 0.0) {
        _giraffe_covsrt(covar, ma, ia, mfit);

        cpl_matrix_delete(moneda);
        moneda = NULL;
        oneda = NULL;

        cpl_matrix_delete(mda);
        mda = NULL;
        da = NULL;

        cpl_matrix_delete(mbeta);
        mbeta = NULL;
        beta = NULL;

        cpl_matrix_delete(matry);
        matry = NULL;
        atry = NULL;

        return 0;
    }

    for (j = -1, l = 0; l < ma; l++) {
        if (ia[l]) {
            atry[l] = pd_a[l] + da[++j];
        }
    }

    gj = _giraffe_mrqcof(x, y, sig, ndata, matry, r, ia, ma, covar, mda,
                         chisq, funcs);

    if (gj != 0) {
        cpl_matrix_delete(moneda);
        moneda = NULL;
        oneda = NULL;

        cpl_matrix_delete(mda);
        mda = NULL;
        da = NULL;

        cpl_matrix_delete(mbeta);
        mbeta = NULL;
        beta = NULL;

        cpl_matrix_delete(matry);
        matry = NULL;
        atry = NULL;

        return gj;
    }

    if (*chisq < ochisq) {

        *alamda *= 0.1;
        ochisq = *chisq;

        for (j = -1, l = 0; l < ma; l++) {
            if (ia[l]) {
                for (j++, k = -1, m = 0; m < ma; m++) {
                    if (ia[m]) {
                        k++;
                        pd_alpha[j * nr_alpha + k] =
                            pd_covar[j * nr_covar + k];
                    }
                }

                beta[j] = da[j];
                pd_a[l] = atry[l];
            }
        }

    }
    else {

        *alamda *= 10.0;
        *chisq = ochisq;

    }

    return 0;

}


/**
 * @brief
 *   Levenberg-Marquardt non-linear fit driver
 *
 * @param x      X abcissa [ndata]
 * @param y      Y values [ndata]
 * @param sigma  Y sigmas [ndata]
 * @param ndata  Number of values
 * @param a      Initial guesses for model parameters [ma]
 * @param delta  Maximum delta for model parameters [ma]
 * @param ia     Flags for model parameters to be fitted [ma]
 * @param ma     Number of model parameters to fit
 * @param alpha  Fitted parameters [ma,ma]
 * @param chisq  Chi square of fit
 * @param funcs  Non linear model to fit
 * @param setup  Non linear fit parameters
 *
 * @return
 *   Number of iterations performed or negative error code
 *
 * @see GiFitParams
 * @see GiFitFunc
 */

cxint
giraffe_nlfit(cpl_matrix *x, cpl_matrix *y, cpl_matrix *sigma,
              cxint ndata, cpl_matrix *a, cpl_matrix *delta, cxint *ia,
              cxint ma, cpl_matrix *alpha, cxdouble *chisq, GiFitFunc funcs,
              const GiFitParams *setup)
{

    cxint itst;
    cxint n;
    cxint res;

    cxdouble alamda = -1.;
    cxdouble *r = NULL;

    cpl_matrix *beta = cpl_matrix_new(ma, ma);


    if (delta) {
        r = cpl_matrix_get_data(delta);
    }

    res = _giraffe_mrqmin(x, y, sigma, ndata, a, r, ia, ma, alpha, beta,
                          chisq, funcs, &alamda);

    if (res != 0) {
        cpl_matrix_delete(beta);
        beta = NULL;

        return res;
    }

    itst=0;

    for (n = 1; n <= setup->iterations; n++) {

        cxdouble ochisq = *chisq;

        res = _giraffe_mrqmin(x, y, sigma, ndata, a, r, ia, ma, alpha, beta,
                              chisq, funcs, &alamda);

        if (res != 0) {
            cpl_matrix_delete(beta);
            beta = NULL;

            return res;
        }

        if (*chisq > ochisq) {
            itst = 0;
        }
        else if (fabs(ochisq - *chisq) < setup->dchisq) {
            itst++;
        }

        if (itst > setup->tests) {
            break;
        }

    }


    /*
     * Get covariance matrix
     */

    alamda=0.0;

    res = _giraffe_mrqmin(x, y, sigma, ndata, a, r, ia, ma, alpha, beta,
                          chisq, funcs, &alamda);

    if (res != 0) {
        cpl_matrix_delete(beta);
        beta = NULL;

        return res;
    }

    cpl_matrix_delete(beta);
    beta = NULL;

    return n;

}
/**@}*/
