/** @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.27 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-11-29 02:25:14 $
 *
 *  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
 *
 */
/*****************************************************************************/
/* 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                                           */
/*****************************************************************************/

psBool p_psMinLM_GuessABP(
    psImage  *Alpha,
    psVector *Beta,
    psVector *Params,
    const psImage  *alpha,
    const psVector *beta,
    const psVector *params,
    const psVector *paramMask,
    psMinimizeLMLimitFunc checkLimits,
    psF32 lambda)
{
    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_MASK, false);
    }

    // XXX the LU Decomposition code requires F64.  this is now incompatible
    // with the rest of this code. drop this segment when the changes are tested.
    # define USE_LU_DECOMP 0
    # if (USE_LU_DECOMP)
        psVector *LUv = NULL;
    psImage  *LUm = NULL;
    psImage  *A   = NULL;
    psF32    det;
    // LU decomposition version
    psTrace("psLib.math", 5, "using LUD version\n");

    // set new guess values (creates matrix A)
    A = psImageCopy(NULL, alpha, PS_TYPE_F64);
    for (int j = 0; j < params->n; j++) {
        if ((paramMask != NULL) && (paramMask->data.U8[j])) {
            continue;
        }
        A->data.F64[j][j] = alpha->data.F64[j][j] * (1.0 + lambda);
    }

    // solve A*beta = Beta (Alpha = 1/A)
    // these operations do not modify the input values (creates LUm, LUv)
    LUm   = psMatrixLUD(NULL, &LUv, A);
    if (LUm == NULL) {
        psError(PS_ERR_UNKNOWN, false, "psMatrixLUD() returned NULL\n");
    }
    Beta  = psMatrixLUSolve(Beta, LUm, beta, LUv);
    if (Beta == NULL) {
        psError(PS_ERR_UNKNOWN, false, "psMatrixLUSolve() returned NULL\n");
    }
    Alpha = psMatrixInvert(Alpha, A, &det);
    psFree(LUm);
    psFree(LUv);
    psFree(A);

    if (Alpha == NULL) {
        psError(PS_ERR_UNKNOWN, false, "psMatrixInvert() returned NULL\n");
        return(false);
    }
    # else
        // gauss-jordan version
        psTrace("psLib.math", 5, "using Gauss-J version");

    // set new guess values (creates matrix A)
    Beta = psVectorCopy(Beta, beta, PS_TYPE_F32);
    Alpha = psImageCopy(Alpha, alpha, PS_TYPE_F32);
    for (int j = 0; j < params->n; j++) {
        if ((paramMask != NULL) && (paramMask->data.U8[j])) {
            continue;
        }
        Alpha->data.F32[j][j] = alpha->data.F32[j][j] * (1.0 + lambda);
    }

    if (false == psMatrixGJSolveF32(Alpha, Beta)) {
        // psError(PS_ERR_UNKNOWN, false, "singular matrix in Guess ABP\n");
        psTrace ("psLib.math", 4, "singular matrix in Guess ABP\n");
        return(false);
    }
    # endif

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

        // compare new params to param limits
        if (checkLimits)
            checkLimits (PS_MINIMIZE_PARAM_MIN,  j, Params->data.F32, Beta->data.F32);
        if (checkLimits)
            checkLimits (PS_MINIMIZE_PARAM_MAX,  j, Params->data.F32, Beta->data.F32);
    }
    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  = psImageAlloc (params->n, params->n, PS_TYPE_F32);
    psImage  *Alpha  = psImageAlloc (params->n, params->n, PS_TYPE_F32);
    psVector *beta   = psVectorAlloc(params->n, PS_TYPE_F32);
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F32);
    psVector *dy     = NULL;
    psBool rc = 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 rcF32 = p_psMinLM_SetABX(alpha, beta, params, paramMask, x, y, dy, func);
    if (isnan(rcF32)) {
        psTrace ("psLib.math", 5, "p_psMinLM_SetABX() returned a NAN.\n");
        rc = false;
    }
    psTrace("psLib.math", 5, "p_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)");
    }

    psBool rcBool = p_psMinLM_GuessABP(Alpha, delta, Params, alpha, beta, params, paramMask, NULL, 0.0);
    if (rcBool == false) {
        psTrace ("psLib.math", 5, "p_psMinLM_GuessABP() returned FALSE.\n");
        rc = false;
    }
    psTrace("psLib.math", 5, "p_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(rc);
}

// measure linear model prediction
psF32 p_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;
    for (int i = 0; i < beta->n; i++) {
        dLinear += lambda*PS_SQR(B[i]) + B[i]*b[i];
    }
    return(0.5*dLinear);
}

// alpha, beta, params are already allocated
psF32 p_psMinLM_SetABX(
    psImage  *alpha,
    psVector *beta,
    const psVector *params,
    const psVector *paramMask,
    const psArray  *x,
    const psVector *y,
    const psVector *dy,
    psMinimizeLMChi2Func func)
{
    // XXX: Check vector sizes.
    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, false);
    if (paramMask) {
        PS_ASSERT_VECTOR_TYPE(paramMask, PS_TYPE_MASK, false);
    }

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

    // zero alpha and beta for summing below
    for (psS32 j = 0; j < params->n; j++) {
        for (psS32 k = 0; k < params->n; k++) {
            alpha->data.F32[j][k] = 0;
        }
        beta->data.F32[j] = 0;
    }
    chisq = 0.0;

    // calculate chisq, alpha, beta
    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]))
            psAbort ("psMinLMM", "nan in weights");
        if (isnan(delta))
            psAbort ("psMinLMM", "nan in delta");
        if (isnan(chisq))
            psAbort ("psMinLMM", "nan in chisq");

        for (psS32 j = 0; j < params->n; j++) {
            if ((paramMask != NULL) && (paramMask->data.U8[j])) {
                continue;
            }
            weight = deriv->data.F32[j] * dy->data.F32[i];
            for (psS32 k = 0; k <= j; k++) {
                if ((paramMask != NULL) && (paramMask->data.U8[k])) {
                    continue;
                }
                alpha->data.F32[j][k] += weight * deriv->data.F32[k];
            }
            beta->data.F32[j] += weight * delta;
        }
    }

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

    // fill in pivots if we apply a mask
    if (paramMask != NULL) {
        for (psS32 j = 0; j < params->n; j++) {
            if (paramMask->data.U8[j]) {
                alpha->data.F32[j][j] = 1;
                beta->data.F32[j] = 1;
            }
        }
    }

    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?
  *****************************************************************************/
psBool 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_U8, 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   = psImageAlloc(params->n, params->n, PS_TYPE_F32);
    psImage *Alpha   = psImageAlloc(params->n, params->n, PS_TYPE_F32);
    psVector *beta   = psVectorAlloc(params->n, PS_TYPE_F32);
    psVector *Beta   = psVectorAlloc(params->n, PS_TYPE_F32);
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F32);
    psVector *dy     = NULL;
    psF32 Chisq = 0.0;
    psF32 lambda = 0.001;

    // 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 = p_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 (!p_psMinLM_GuessABP(Alpha, Beta, Params, alpha, beta, params, paramMask, checkLimits, lambda)) {
            min->iter ++;
            lambda *= 10.0;
            continue;
        }

        // measure linear model prediction
        psF32 dLinear = p_psMinLM_dLinear(Beta, beta, lambda);

        // 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)");
        }
        if (psTraceGetLevel("psLib.math") >= 5) {
            p_psVectorPrint(psTraceGetDestination(), Params, "params guess (1)");
        }

        // calculate Chisq for new guess, update Alpha & Beta
        Chisq = p_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, rho: %f\n", min->value,
                Chisq, min->lastDelta, rho);

        // 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)");
        }

        /* if (Chisq < min->value) {  */
        if (rho > 0.0) {
            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.1;
        } 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 (!p_psMinLM_GuessABP(covar, Beta, Params, alpha, beta, params, paramMask, NULL, 0.0)) {
            psTrace ("psLib.math", 5, "failure to calculate covariance matrix\n");
        }
    }

    // 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);
}

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()
{
    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);
}
