/*
 * This file is part of the ESO Common Pipeline Library
 * Copyright (C) 2001-2008,2014 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
 */

/**
 * @defgroup cpl_matrix Matrices
 *
 * This module provides functions to create, destroy and use a @em cpl_matrix.
 * The elements of a @em cpl_matrix with M rows and N columns are counted 
 * from 0,0 to M-1,N-1. The matrix element 0,0 is the one at the upper left
 * corner of a matrix. The CPL matrix functions work properly only in the 
 * case the matrices elements do not contain garbage (such as @c NaN or 
 * infinity).
 *
 * @par Synopsis:
 * @code
 *   #include <flames_lsfit.h>
 * @endcode
 */

#include <cpl.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include <flames_lsfit.h>
#include <flames_gauss_jordan.h>
#include <flames_covariance_reorder.h>
#include <flames_newmatrix.h>
#include <uves_dfs.h>
/**@{*/

cpl_matrix * cpl_matrix_product_normal_create(const cpl_matrix * self);
cpl_error_code cpl_matrix_product_transpose(cpl_matrix * self,
                                            const cpl_matrix * ma,
                                            const cpl_matrix * mb);


/*-----------------------------------------------------------------------------
                        Private function prototypes
 -----------------------------------------------------------------------------*/

/*-----------------------------------------------------------------------------
                              Function codes
 -----------------------------------------------------------------------------*/

/* ---------------------------------------------------------------------------*/
/**
 * @brief generic 1d vandermonde matrix
 *
 * @param sample  sampling positions
 * @param degree  degree of polynomial
 * @param func    function evaluating polynomials from [0, degree] at
 *                sampling point
 * @param offset  offset of func (1 for indexing starting with 1)
 * @return matrix containing the vandermonde matrix
 */
/* ---------------------------------------------------------------------------*/
cpl_matrix * vander1d(
                const cpl_vector * sample,
                cpl_size           degree,
                void (*func)(double, double *, int),
                size_t offset)
{
    size_t i;
    const size_t nr = cpl_vector_get_size(sample);
    const size_t nc = degree + 1;
    cpl_matrix * V = cpl_matrix_new(nr, nc);
    double * v = cpl_matrix_get_data(V);
    const double * d = cpl_vector_get_data_const(sample);
    for (i = 0; i < nr; i++) {
        if (offset != 0) {
            double tmp[nc + offset];
            func(d[i], tmp, nc);
            memcpy(&v[i * nc], &tmp[offset], nc * sizeof(*v));
        }
        else {
            func(d[i], &v[i * nc], nc);
        }
    }

    return V;
}

cpl_matrix * vander2d(
                const cpl_vector * sample_x,
                const cpl_vector * sample_y,
                cpl_size           degree,
                void (*func)(double, double, double *, int),
                size_t offset)
{
    size_t i;
    const size_t nr = cpl_vector_get_size(sample_x);
    const size_t nc = degree + 1;
    cpl_matrix * V = cpl_matrix_new(nr, nc);
    double * v = cpl_matrix_get_data(V);
    const double * dx = cpl_vector_get_data_const(sample_x);
    const double * dy = cpl_vector_get_data_const(sample_y);
    assert(cpl_vector_get_size(sample_y) == nr);
    for (i = 0; i < nr; i++) {
        if (offset != 0) {
            double tmp[nc + offset];
            func(dx[i], dy[i], tmp, nc);
            memcpy(&v[i * nc], &tmp[offset], nc * sizeof(*v));
        }
        else {
            func(dx[i], dy[i], &v[i * nc], nc);
        }
    }

    return V;
}

static void flames_polynomial(double x, double * p, int ncoefs)
{
    int i;
    p[0] = 1.;
    for (i = 1; i < ncoefs; i++) {
        p[i] = pow(x, i);
    }
}

cpl_matrix * polyvander1d(
                const cpl_vector * sample,
                cpl_size           degree)
{
    return vander1d(sample, degree, &flames_polynomial, 0);
}

void lsqfit(
                const cpl_matrix * design,
                const cpl_vector * values,
                const cpl_vector * errors,
                cpl_matrix **coef)
{
    /* weight response and design */
    cpl_vector * vrhs = cpl_vector_duplicate(errors);
    cpl_vector_power(vrhs, -1);
    cpl_matrix * wdesign = cpl_matrix_duplicate(design);
    long i, j;
    for (i = 0; i < cpl_vector_get_size(errors); i++) {
        double w = cpl_vector_get(vrhs,i);
        for (j = 0; j < cpl_matrix_get_ncol(wdesign); j++) {
            cpl_matrix_set(wdesign, i, j,
                            cpl_matrix_get(wdesign, i, j) * w);
        }
    }

    cpl_vector_multiply(vrhs, values);
    cpl_matrix * rhs = cpl_matrix_wrap(cpl_vector_get_size(vrhs), 1,
                    cpl_vector_get_data(vrhs));

    /* solve Ax = b */
    /* cpl_matrix_solve_normal(design, rhs) + covariance */
    {
        cpl_matrix * At  = cpl_matrix_transpose_create(wdesign);
        cpl_matrix * AtA = cpl_matrix_product_normal_create(At);

        /* RRt = AtA */
        cpl_matrix_decomp_chol(AtA);
        /* solve for pseudo inverse: (RRt)P=At*/
        cpl_matrix_solve_chol(AtA, At);
        /* compute solution to system Ax=b -> x=Pb */
        *coef = cpl_matrix_product_create(At, rhs);
        /* compute covariance matrix cov(b) = PPt */
        //cpl_matrix * cov = cpl_matrix_new(cpl_matrix_get_ncol(At),
        //                        cpl_matrix_get_ncol(At));
        //cpl_matrix_product_transpose(cov, At, At);

        cpl_matrix_delete(At);
        cpl_matrix_delete(AtA);
    }

    cpl_matrix_unwrap(rhs);
    cpl_vector_delete(vrhs);
    cpl_matrix_delete(wdesign);

    return;
}

void lsqfit_nr(
                double x[], double y[], double sig[], int32_t ndat, double a[],
                int ma,
                void (*funcs)(double, double [], int))
{
    int i;
    cpl_vector * vX = cpl_vector_wrap(ndat, &x[1]);
    cpl_vector * vY = cpl_vector_wrap(ndat, &y[1]);
    cpl_vector * vS;
    cpl_matrix * design = vander1d(vX, ma -1, funcs, 1);
    cpl_matrix * mcoef;

    if (sig) {
        vS = cpl_vector_wrap(ndat, &sig[1]);
    }
    else {
        vS = cpl_vector_new(ndat);
        for (i= 0; i < ndat; i++)
            cpl_vector_set(vS, i, 1.);
    }
    lsqfit(design, vY, vS, &mcoef);

    for (i = 1; i <= ma; i++) {
        a[i] = cpl_matrix_get(mcoef, i -1, 0);
    }
    cpl_vector_unwrap(vX);
    cpl_vector_unwrap(vY);
    if (sig) {
        cpl_vector_unwrap(vS);
    }
    else {
        cpl_vector_delete(vS);
    }
    cpl_matrix_delete(design);
    cpl_matrix_delete(mcoef);
}


void lsqfit2d_nr(
                double x[], double y[], double z[], double sig[], int ndat, double a[],
                int ma,
                void (*funcs)(double, double, double [], int))
{
    int i;
    cpl_vector * vX = cpl_vector_wrap(ndat, &x[1]);
    cpl_vector * vY = cpl_vector_wrap(ndat, &y[1]);
    cpl_vector * vZ = cpl_vector_wrap(ndat, &z[1]);
    cpl_vector * vS;
    cpl_matrix * design = vander2d(vX, vY, ma - 1, funcs, 1);
    cpl_matrix * mcoef;

    if (sig) {
        vS = cpl_vector_wrap(ndat, &sig[1]);
    }
    else {
        vS = cpl_vector_new(ndat);
        for (i= 0; i < ndat; i++)
            cpl_vector_set(vS, i, 1.);
    }
    lsqfit(design, vZ, vS, &mcoef);

    for (i = 1; i <= ma; i++) {
        a[i] = cpl_matrix_get(mcoef, i - 1, 0);
    }
    cpl_vector_unwrap(vX);
    cpl_vector_unwrap(vY);
    cpl_vector_unwrap(vZ);
    if (sig) {
        cpl_vector_unwrap(vS);
    }
    else {
        cpl_vector_delete(vS);
    }
    cpl_matrix_delete(design);
    cpl_matrix_delete(mcoef);
}


cpl_vector * eval_poly(
                cpl_matrix * design,
                cpl_matrix * coef)
{
    cpl_matrix * fvalues = cpl_matrix_product_create(design, coef);
    cpl_vector * res = cpl_vector_wrap(cpl_matrix_get_nrow(fvalues),
                    cpl_matrix_get_data(fvalues));
    cpl_matrix_unwrap(fvalues);
    return res;
}


/* fill upper-left matrix elements as mirror of lower-right elements
 * below matrix diagonal
 */
static void
flames_matrix_mirror_lu(double** mat,const int nraws)
{
    int j,k;
    for (j=2;j<=nraws;j++)
        for (k=1;k<j;k++)
            mat[k][j]=mat[j][k];

    return;
}


/* compute chi2 */
static double
flames_compute_chi2(double* data, double* value, double* errs, const int ndat,
                    double* par, const int npar,double* afunc,
                    void (*model)(double, double [], int)) {
    double chi2=0,sum=0;
    int i,j;
    for (i=1;i<=ndat;i++) {

        (*model)(data[i],afunc,npar);
        for (sum=0.0,j=1;j<=npar;j++) sum += par[j]*afunc[j];
        double scatter = (value[i]-sum)/errs[i];
        chi2 += scatter*scatter;
    }

    return chi2;

}

/* build design matrix: pointer to be allocated */
static void
flames_matrix_design(double* data,double* value,double* ps,const int ndat,
                     double* par,int* par_sw, const int npar,const int nfit,
                     double** covar,double** mat, double* afunc,
                     void (*model)(double, double [], int))
{
    int j,k,l,m;

    for (int i=1;i<=ndat;i++) {
        (*model)(data[i],afunc,npar);
        double value_model=value[i];

        if (nfit < npar) {
            for (j=1;j<=npar;j++)
                if (!par_sw[j]) value_model -= par[j]*afunc[j];
        }
        double sig2i=1.0/(ps[i]*ps[i]);
        for (j=0,l=1;l<=npar;l++) {
            if (par_sw[l]) {
                double weight=afunc[l]*sig2i;
                for (j++,k=0,m=1;m<=l;m++)
                    if (par_sw[m]) covar[j][++k] += weight*afunc[m];
                mat[j][1] += value_model*weight;
            }
        }
    }

    return;
}

void flames_lfit(cpl_vector* vx, cpl_vector* vy, cpl_vector* vs, int32_t ndat,
                 double a[],
                 int ia[], int ma, double **covar, double *chisq,
                 void (*funcs)(double, double [], int))
{
    int j=0;
    int l=0;
    int mfit=0;
    double **beta=0;
    double *afunc=0;

    double *px = cpl_vector_get_data(vx);
    double *py = cpl_vector_get_data(vy);
    double *ps = cpl_vector_get_data(vs);

    beta=dmatrix(1,ma,1,1);
    afunc=dvector(1,ma);
    /* Not needed as calloc also init to 0.
     for (i=1; i<=ma; i++) {
       beta[i][1] = 0;
       afunc[i] = 0;
     }
     */
    for (j=1;j<=ma;j++)
        if (ia[j]) mfit++;
    if (mfit == 0) nrerror("lfit: no parameters to be fitted");
    /* Not needed as calloc also init to 0.
     for (j=1;j<=mfit;j++) {
         for (k=1;k<=mfit;k++) covar[j][k]=0.0;
         beta[j][1]=0.0;
     }
     */
    flames_matrix_design(px,py,ps,ndat,a,ia,ma,mfit,covar,beta,afunc,funcs);
    flames_matrix_mirror_lu(covar,mfit);
    flames_gauss_jordan(covar,mfit,beta,1);

    for (j=0,l=1;l<=ma;l++)
        if (ia[l]) a[l]=beta[++j][1];
    *chisq=flames_compute_chi2(px,py,ps,ndat,a,ma,afunc,funcs);
    flames_covariance_reorder(covar,ma,ia,mfit);
    free_dvector(afunc,1,ma);
    free_dmatrix(beta,1,ma,1,1);
    return;
}