/** @file  psMinimize.c
 *  \brief basic minimization functions
 *  @ingroup Math
 *
 *  This file will contain functions to minimize an arbitrary function at
 *  a data point, fit an arbitrary function to a set of data points, and
 *  fit a 1-D polynomial to a set of data points.
 *
 *  @author GLG, MHPCC
 *  @author EAM, IfA
 *
 *  @version $Revision: 1.36 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2009-01-27 06:39:38 $
 *
 *  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
 *
 */

#ifdef HAVE_CONFIG_H
# include "config.h"
#endif

/*****************************************************************************/
/* INCLUDE FILES                                                             */
/*****************************************************************************/
#include <stdio.h>
#include <float.h>
#include <math.h>

#include "psAbort.h"
#include "psAssert.h"
#include "psMinimizeLMM.h"
#include "psImage.h"
#include "psImageStructManip.h"
#include "psLogMsg.h"
/*****************************************************************************/
/* DEFINE STATEMENTS                                                         */
/*****************************************************************************/

/*****************************************************************************/
/* TYPE DEFINITIONS                                                          */
/*****************************************************************************/

/*****************************************************************************/
/* GLOBAL VARIABLES                                                          */
/*****************************************************************************/

/*****************************************************************************/
/* FILE STATIC VARIABLES                                                     */
/*****************************************************************************/

/*****************************************************************************/
/* FUNCTION IMPLEMENTATION - LOCAL                                           */
/*****************************************************************************/

// Alpha & Beta only represent unmasked values
bool psMinLM_GuessABP(
    psImage  *Alpha,
    psVector *Beta,
    psVector *Params,
    const psImage  *alpha,
    const psVector *beta,
    const psVector *params,
    const psVector *paramMask,
    psMinimizeLMLimitFunc checkLimits,
    psF32 lambda,
    psF32 *dLinear)
{
    PS_ASSERT_VECTOR_TYPE(Alpha,     PS_TYPE_F32,  false);
    PS_ASSERT_VECTOR_TYPE(Beta,      PS_TYPE_F32,  false);
    PS_ASSERT_VECTOR_TYPE(Params,    PS_TYPE_F32,  false);
    PS_ASSERT_VECTOR_TYPE(alpha,     PS_TYPE_F32,  false);
    PS_ASSERT_VECTOR_TYPE(beta,      PS_TYPE_F32,  false);
    PS_ASSERT_VECTOR_TYPE(params,    PS_TYPE_F32,  false);
    if (paramMask) {
        PS_ASSERT_VECTOR_TYPE(paramMask, PS_TYPE_VECTOR_MASK, false);
    }

    assert (alpha->numCols == beta->n);
    assert (alpha->numCols == alpha->numRows);

# define TESTGJ 0
# if (TESTGJ)
    lambda = 0.0;
#endif

    // set new guess values, applying (1+lambda) scaling to pivots
    Beta = psVectorCopy(Beta, beta, PS_TYPE_F32);
    Alpha = psImageCopy(Alpha, alpha, PS_TYPE_F32);
    for (int j = 0; j < Alpha->numCols; j++) {
        Alpha->data.F32[j][j] = alpha->data.F32[j][j] * (1.0 + lambda);
    }

    // error and clear above if kept?
    if (!psMatrixGJSolve(Alpha, Beta)) {
        psTrace ("psLib.math", 4, "singular matrix in Guess ABP\n");
        return(false);
    }

// XXX check that the GJ solver works:
# if (TESTGJ)
    psImage *out = psImageAlloc (alpha->numRows, alpha->numCols, PS_TYPE_F32);
    for (int oy = 0; oy < out->numRows; oy++) {
        for (int ox = 0; ox < out->numCols; ox++) {
            float value = 0;
            for (int i = 0; i < alpha->numCols; i++) {
                value += alpha->data.F32[i][ox]*Alpha->data.F32[oy][i];
            }
            out->data.F32[oy][ox] = value;
        }
    }

    psVector *vect = psVectorAlloc (beta->n, PS_TYPE_F32);
    for (int oy = 0; oy < vect->n; oy++) {
        float value = 0;
        for (int i = 0; i < alpha->numCols; i++) {
            value += alpha->data.F32[oy][i]*Beta->data.F32[i];
        }
        vect->data.F32[oy] = value;
    }

    psFree (out);
    psFree (vect);

# endif

    // check for non-finite entries in result
    for (int i = 0; i < Beta->n; i++) {
        if (!isfinite(Beta->data.F32[i])) {
            // psError(PS_ERR_BAD_PARAMETER_VALUE, 3, "Fit value diverges: vector[%d] is %.2f\n", i, Beta->data.F32[i]);
            return false;
        }
    }

    // measure linear model prediction
    // (we must do this before truncating Beta below)
    if (dLinear) {
        *dLinear = psMinLM_dLinear(Beta, beta, lambda);
    }

    // full-length Beta for checkLimits functions
    psVector *tmpBeta = psVectorAlloc(params->n, PS_TYPE_F32);
    psVectorInit (tmpBeta, 0.0);

    // set tmpBeta values which are not masked
    for (int j = 0, n = 0; j < params->n; j++) {
        if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[j])) continue;
        tmpBeta->data.F32[j] = Beta->data.F32[n];
        n++;
    }

    // apply Beta to get new Params values
    for (int j = 0; j < params->n; j++) {
        if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[j])) {
            Params->data.F32[j] = params->data.F32[j];
            continue;
        }
        // apply beta limits
        if (checkLimits) {
            checkLimits (PS_MINIMIZE_BETA_LIMIT, j, Params->data.F32, tmpBeta->data.F32);
        }

        Params->data.F32[j] = params->data.F32[j] - tmpBeta->data.F32[j];

        // compare new params to param limits
        if (checkLimits) {
            checkLimits (PS_MINIMIZE_PARAM_MIN,  j, Params->data.F32, tmpBeta->data.F32);
            checkLimits (PS_MINIMIZE_PARAM_MAX,  j, Params->data.F32, tmpBeta->data.F32);
        }
    }

    // apply tmpBeta after limits have been checked
    for (int j = 0, n = 0; j < params->n; j++) {
        if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[j])) continue;
        Beta->data.F32[n] = tmpBeta->data.F32[j];
        n++;
    }

    psFree (tmpBeta);
    return(true);
}

bool psMinimizeGaussNewtonDelta(
    psVector *delta,
    const psVector *params,
    const psVector *paramMask,
    const psArray  *x,
    const psVector *y,
    const psVector *yWt,
    psMinimizeLMChi2Func func)
{
    psTrace("psLib.math", 3, "---- begin ----\n");

    // allocate internal arrays (current vs Guess)
    psImage *Alpha = NULL;
    psVector *Beta = NULL;

    psVectorInit (delta, 0.0);

    // Alpha & Beta only contain elements to represent the unmasked parameters
    // if none are available, return false
    if (!psMinLM_AllocAB (&Alpha, &Beta, params, paramMask)) {
        return false;
    }

    psImage *alpha   = psImageAlloc(Alpha->numCols, Alpha->numRows, PS_TYPE_F32);
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F32);

    psVector *dy     = NULL;
    bool retValue = true;

    // the user provides the error or NULL.  we need to convert
    // to appropriate weights
    if (yWt != NULL) {
        dy = (psVector *) yWt;
    } else {
        dy = psVectorAlloc(y->n, PS_TYPE_F32);
        psVectorInit(dy, 1.0);
    }

    // XXX should we give up if chisq is nan?
    psF32 chisq = psMinLM_SetABX(alpha, Beta, params, paramMask, x, y, dy, func);
    if (isnan(chisq)) {
        psTrace ("psLib.math", 5, "psMinLM_SetABX() returned a NAN chisq.\n");
        psVectorInit (delta, NAN);
        retValue = false;
    }

    psTrace("psLib.math", 5, "psMinLM_SetABX() was succesful\n");
    // dump some useful info if trace is defined
    if (psTraceGetLevel("psLib.math") >= 6) {
        p_psImagePrint(psTraceGetDestination(), alpha, "alpha guess (0)");
        p_psVectorPrint(psTraceGetDestination(), Beta, "beta guess (0)");
        p_psVectorPrint(psTraceGetDestination(), params, "params guess (0)");
    }

    bool status = psMinLM_GuessABP(Alpha, delta, Params, alpha, Beta, params, paramMask, NULL, 0.0, NULL);
    if (!status) {
        psTrace ("psLib.math", 5, "psMinLM_GuessABP() returned FALSE.\n");
        psVectorInit (delta, NAN);
        retValue = false;
    }
    psTrace("psLib.math", 5, "psMinLM_GuessABP() was succesful\n");
    if (psTraceGetLevel("psLib.math") >= 6) {
        p_psImagePrint(psTraceGetDestination(), Alpha, "alpha guess (1)");
        p_psVectorPrint(psTraceGetDestination(), delta, "delta guess (1)");
        p_psVectorPrint(psTraceGetDestination(), Params, "params guess (1)");
    }

    psFree(alpha);
    psFree(Alpha);
    psFree(Beta);
    psFree(Params);
    if (yWt == NULL) {
        psFree(dy);
    }
    psTrace("psLib.math", 3, "---- end ----\n");
    return(retValue);
}

// measure linear model prediction
psF32 psMinLM_dLinear(
    const psVector *Beta,
    const psVector *beta,
    psF32 lambda)
{

    /* get linear model prediction */
    psF32 dLinear = 0;
    psF32 *B = Beta->data.F32;
    psF32 *b = beta->data.F32;

    float dh = 0.0, sh = 0.0, Sh = 0.0;

    // beta only counts unmasked parameters
    for (int i = 0; i < beta->n; i++) {
        dh = lambda*B[i] + b[i];
        sh = 0.5*B[i]*dh;
        Sh += sh;
        dLinear += lambda*PS_SQR(B[i]) + B[i]*b[i];
    }
    return(0.5*dLinear);
}

// alpha, beta, params are already allocated
psF32 psMinLM_SetABX(
    psImage  *alpha,
    psVector *beta,
    const psVector *params,
    const psVector *paramMask,
    const psArray  *x,
    const psVector *y,
    const psVector *dy,
    psMinimizeLMChi2Func func)
{
    PS_ASSERT_IMAGE_NON_NULL(alpha, NAN);
    PS_ASSERT_VECTOR_NON_NULL(beta, NAN);
    PS_ASSERT_VECTOR_NON_NULL(params, NAN);
    PS_ASSERT_PTR_NON_NULL(x, NAN);
    PS_ASSERT_VECTOR_NON_NULL(y, NAN);
    PS_ASSERT_VECTOR_NON_NULL(dy, NAN);

    PS_ASSERT_VECTOR_TYPE(params, PS_TYPE_F32, NAN);
    if (paramMask) {
        PS_ASSERT_VECTOR_TYPE(paramMask, PS_TYPE_VECTOR_MASK, NAN);
    }

    psF32 chisq;
    psF32 delta;
    psF32 weight;
    psF32 ymodel;
    psVector *deriv = psVectorAlloc(params->n, PS_TYPE_F32);

    // zero alpha, beta, and chisq for summing below
    psImageInit (alpha, 0.0);
    psVectorInit (beta, 0.0);
    chisq = 0.0;

    // calculate chisq, alpha, beta. alpha & beta only represent unmasked parameters; skip
    // masked ones
    for (psS32 i = 0; i < y->n; i++) {
        ymodel = func(deriv, params, (psVector *) x->data[i]);

        delta = ymodel - y->data.F32[i];
        chisq += PS_SQR(delta) * dy->data.F32[i];

        if (isnan(dy->data.F32[i])) return NAN;
        if (isnan(delta)) return NAN;
        if (isnan(chisq)) return NAN;

        // we track alpha,beta and params,deriv separately
        for (int j = 0, J = 0; j < params->n; j++) {
            if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[j])) continue;

            weight = deriv->data.F32[j] * dy->data.F32[i];

            for (int k = 0, K = 0; k <= j; k++) {
                if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[k])) continue;
                alpha->data.F32[J][K] += weight * deriv->data.F32[k];
                K++;
            }
            beta->data.F32[J] += weight * delta;
            J++;
        }
    }

    // calculate lower-left half of alpha
    for (int j = 1; j < alpha->numCols; j++) {
        for (int k = 0; k < j; k++) {
            alpha->data.F32[k][j] = alpha->data.F32[j][k];
        }
    }

    psFree(deriv);
    return(chisq);
}


/******************************************************************************
psMinimizeLMChi2():  This routine will take an procedure which calculates an
arbitrary function and it's derivative and minimize the chi-squared match
between that function at the specified coords and the specified value at those
coords.

This requires F32 input data; all internal calls use F32.
XXX Make an F64 version?
  *****************************************************************************/
bool psMinimizeLMChi2(
    psMinimization *min,
    psImage *covar,
    psVector *params,
    psMinConstraint *constraint,
    const psArray *x,
    const psVector *y,
    const psVector *yWt,
    psMinimizeLMChi2Func func)
{
    psTrace("psLib.math", 3, "---- begin ----\n");
    PS_ASSERT_PTR_NON_NULL(min, false);
    PS_ASSERT_VECTOR_NON_NULL(params, false);
    PS_ASSERT_VECTOR_NON_EMPTY(params, false);
    PS_ASSERT_VECTOR_TYPE(params, PS_TYPE_F32, false);
    psVector *paramMask = NULL;
    if (constraint != NULL) {
        paramMask = constraint->paramMask;
        if (paramMask != NULL) {
            PS_ASSERT_VECTOR_TYPE(paramMask, PS_TYPE_VECTOR_MASK, false);
            PS_ASSERT_VECTORS_SIZE_EQUAL(params, paramMask, false);
        }
    }
    PS_ASSERT_PTR_NON_NULL(x, false);
    for (psS32 i = 0 ; i < x->n ; i++) {
        psVector *coord = (psVector *) (x->data[i]);
        PS_ASSERT_VECTOR_NON_NULL(coord, false);
        PS_ASSERT_VECTOR_TYPE(coord, PS_TYPE_F32, false);
    }
    PS_ASSERT_VECTOR_NON_NULL(y, false);
    PS_ASSERT_VECTOR_NON_EMPTY(y, false);
    PS_ASSERT_VECTOR_TYPE(y, PS_TYPE_F32, false);
    PS_ASSERT_VECTORS_SIZE_EQUAL(x, y, false);
    if (yWt != NULL) {
        PS_ASSERT_VECTOR_TYPE(yWt, PS_TYPE_F32, false);
        PS_ASSERT_VECTORS_SIZE_EQUAL(y, yWt, false);
    }
    PS_ASSERT_PTR_NON_NULL(func, false);

    psMinimizeLMLimitFunc checkLimits = NULL;
    if (constraint) {
        checkLimits = constraint->checkLimits;
    }

    // this function has test and current values for several things
    // the current best value is in lower case
    // the next guess value is in upper case

    // allocate internal arrays (current vs Guess)
    psImage *Alpha = NULL;
    psVector *Beta = NULL;

    // Alpha & Beta only contain elements to represent the unmasked parameters
    if (!psMinLM_AllocAB (&Alpha, &Beta, params, paramMask)) {
        psAbort ("programming error: no unmasked parameters to be fit\n");
    }

    psImage *alpha   = psImageAlloc(Alpha->numCols, Alpha->numRows, PS_TYPE_F32);
    psVector *beta   = psVectorAlloc(Beta->n, PS_TYPE_F32);
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F32);

    psVector *dy     = NULL;
    psF32 Chisq = 0.0;
    psF32 lambda = 0.001;
    psF32 dLinear = 0.0;

    // the user provides the error or NULL.  we need to convert
    // to appropriate weights
    if (yWt != NULL) {
        dy = (psVector *) yWt;
    } else {
        dy = psVectorAlloc(y->n, PS_TYPE_F32);
        psVectorInit(dy, 1.0);
    }

    // calculate initial alpha and beta, set chisq (min->value)
    min->value = psMinLM_SetABX(alpha, beta, params, paramMask, x, y, dy, func);
    if (isnan(min->value)) {
        min->iter = min->maxIter;
        return(false);
    }
    // dump some useful info if trace is defined
    if (psTraceGetLevel("psLib.math") >= 6) {
        p_psImagePrint(psTraceGetDestination(), alpha, "alpha guess (0)");
        p_psVectorPrint(psTraceGetDestination(), beta, "beta guess (0)");
    }
    if (psTraceGetLevel("psLib.math") >= 5) {
        p_psVectorPrint(psTraceGetDestination(), params, "params guess (0)");
    }

    // iterate until the tolerance is reached, or give up
    while ((min->iter < min->maxIter) && ((min->lastDelta > min->tol) || !isfinite(min->lastDelta))) {
        psTrace("psLib.math", 5, "Iteration number %d.  (max iterations is %d).\n", min->iter, min->maxIter);
        psTrace("psLib.math", 5, "Last delta is %f.  Min->tol is %f.\n", min->lastDelta, min->tol);

        // set a new guess for Alpha, Beta, Params
        if (!psMinLM_GuessABP(Alpha, Beta, Params, alpha, beta, params, paramMask, checkLimits, lambda, &dLinear)) {
            min->iter ++;
            lambda *= 10.0;
            continue;
        }

        // dump some useful info if trace is defined
        if (psTraceGetLevel("psLib.math") >= 6) {
            p_psImagePrint(psTraceGetDestination(), Alpha, "Alpha guess (1)");
            p_psVectorPrint(psTraceGetDestination(), Beta, "Beta guess (1)");
            p_psVectorPrint(psTraceGetDestination(), beta, "beta current (1)");
        }
        if (psTraceGetLevel("psLib.math") >= 5) {
            p_psVectorPrint(psTraceGetDestination(), Params, "params guess (1)");
        }

        // calculate Chisq for new guess, update Alpha & Beta
        Chisq = psMinLM_SetABX(Alpha, Beta, Params, paramMask, x, y, dy, func);
        if (isnan(Chisq)) {
            min->iter ++;
            lambda *= 10.0;
            continue;
        }

        // convergence criterion:
        // compare the delta (min->value - Chisq) with the
        // expected delta from the linear model (dLinear)
        // accept new guess if it is an improvement (rho > 0), or else increase lambda
        psF32 rho = (min->value - Chisq) / dLinear;

        psTrace("psLib.math", 5, "last chisq: %f, new chisq %f, delta: %f, dLinear: %f, rho: %f, lambda: %f\n", min->value,
                Chisq, min->lastDelta, dLinear, rho, lambda);

        psTrace("psLib.math.dLinear", 5, "last chisq: %f, new chisq %f, delta: %f, dLinear: %f, rho: %f, lambda: %f\n", min->value,
                Chisq, min->lastDelta, dLinear, rho, lambda);

        // dump some useful info if trace is defined
        if (psTraceGetLevel("psLib.math") >= 6) {
            p_psImagePrint(psTraceGetDestination(), Alpha, "alpha guess (2)");
            p_psVectorPrint(psTraceGetDestination(), Beta, "beta guess (2)");
        }

        /* rho is positive if the new chisq is smaller; allow for some insignificant change (slight negative rho) */
        if (rho >= -1e-6) {
            min->lastDelta = (min->value - Chisq) / (dy->n - params->n);
            min->value = Chisq;
            alpha  = psImageCopy(alpha, Alpha, PS_TYPE_F32);
            beta   = psVectorCopy(beta, Beta, PS_TYPE_F32);
            params = psVectorCopy(params, Params, PS_TYPE_F32);
            lambda *= 0.25;
        } else {
            lambda *= 10.0;
        }
        min->iter++;
    }
    psTrace("psLib.math", 5, "chisq: %f, last delta: %f, Niter: %d\n", min->value, min->lastDelta, min->iter);

    // construct & return the covariance matrix (if requested)
    if (covar != NULL) {
        if (!psMinLM_GuessABP(Alpha, Beta, Params, alpha, beta, params, paramMask, NULL, 0.0, NULL)) {
            psTrace ("psLib.math", 5, "failure to calculate covariance matrix\n");
        }
        // set covar values which are not masked
        psImageInit (covar, 0.0);
        for (int j = 0, J = 0; j < params->n; j++) {
            if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[j])) {
                covar->data.F32[j][j] = 1.0;
                continue;
            }
            for (int k = 0, K = 0; k < params->n; k++) {
                if (paramMask && (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[k])) continue;
                covar->data.F32[j][k] = Alpha->data.F32[J][K];
                K++;
            }
            J++;
        }
    }

    // free the internal temporary data
    psFree(alpha);
    psFree(Alpha);
    psFree(beta);
    psFree(Beta);
    psFree(Params);
    if (yWt == NULL) {
        psFree(dy);
    }
    if (min->iter == min->maxIter) {
        psTrace("psLib.math", 3, "---- end (false) ----\n");
        return(false);
    }
    psTrace("psLib.math", 3, "---- end (true) ----\n");
    return(true);
}

bool psMinLM_AllocAB (psImage **Alpha, psVector **Beta, const psVector *params, const psVector *paramMask) {

    assert (Alpha);
    assert (Beta);
    assert (params);

    int nParams = params->n;

    // count unmasked parameters
    if (paramMask) {
        nParams = 0;
        for (int i = 0; i < paramMask->n; i++) {
            if (paramMask->data.PS_TYPE_VECTOR_MASK_DATA[i]) continue;
            nParams ++;
        }
    }

    if (nParams == 0) {
        return false;
    }

    *Alpha = psImageAlloc(nParams, nParams, PS_TYPE_F32);
    *Beta  = psVectorAlloc(nParams, PS_TYPE_F32);
    return true;
}

static void minimizationFree(psMinimization *min)
{
    // There are no dynamically allocated items
}

psMinimization *psMinimizationAlloc(int maxIter,
                                    float tol)
{
    PS_ASSERT_INT_NONNEGATIVE(maxIter, NULL);

    psMinimization *min = psAlloc(sizeof(psMinimization));
    psMemSetDeallocator(min, (psFreeFunc)minimizationFree);
    P_PSMINIMIZATION_SET_MAXITER(min,maxIter);
    P_PSMINIMIZATION_SET_TOL(min,tol);
    min->value = 0.0;
    min->iter = 0;
    min->lastDelta = NAN;

    return(min);
}

bool psMemCheckMinimization(psPtr ptr)
{
    PS_ASSERT_PTR(ptr, false);
    return( psMemGetDeallocator(ptr) == (psFreeFunc)minimizationFree );
}


static void constraintFree(psMinConstraint *tmp)
{
    if (tmp == NULL)
        return;

    psFree (tmp->paramMask);
}

psMinConstraint* psMinConstraintAlloc(void)
{
    psMinConstraint *tmp = psAlloc(sizeof(psMinConstraint));
    psMemSetDeallocator(tmp, (psFreeFunc)constraintFree);
    tmp->paramMask = NULL;
    tmp->checkLimits = NULL;

    return(tmp);
}

bool psMemCheckConstraint(psPtr tmp)
{
    return(psMemGetDeallocator(tmp) == (psFreeFunc) constraintFree);
}
