/** @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.20 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-08-26 04:34:28 $
 *
 *  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                                           */
/*****************************************************************************/

// XXX EAM : can we use static copies of LUv, LUm, A?
psBool p_psMinLM_GuessABP(
    psImage  *Alpha,
    psVector *Beta,
    psVector *Params,
    const psImage  *alpha,
    const psVector *beta,
    const psVector *params,
    const psVector *paramMask,
    const psVector *beta_lim,
    const psVector *params_min,
    const psVector *params_max,
    psF64 lambda)
{
    # 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_F64);
    Alpha = psImageCopy(Alpha, alpha, PS_TYPE_F64);
    for (int j = 0; j < params->n; j++) {
        if ((paramMask != NULL) && (paramMask->data.U8[j])) {
            continue;
        }
        Alpha->data.F64[j][j] = alpha->data.F64[j][j] * (1.0 + lambda);
    }

    if (false == psMatrixGJSolve(Alpha, Beta)) {
        // psError(PS_ERR_UNKNOWN, false, "singular matrix in Guess ABP\n");
        psTrace (__func__, 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])) {
            continue;
        }
        // Params->data.F32[j] = params->data.F32[j] - Beta->data.F64[j];
        // compare Beta to beta limits
        if (beta_lim != NULL) {
            if (fabs(Beta->data.F64[j]) > fabs(beta_lim->data.F32[j])) {
                Beta->data.F64[j] = (Beta->data.F64[j] > 0) ? fabs(beta_lim->data.F32[j]) : -fabs(beta_lim->data.F32[j]);
            }
        }
        Params->data.F32[j] = params->data.F32[j] - Beta->data.F64[j];
        // compare new params to param limits
        if (params_max != NULL) {
            Params->data.F32[j] = PS_MIN (Params->data.F32[j], params_max->data.F32[j]);
        }
        if (params_min != NULL) {
            Params->data.F32[j] = PS_MAX (Params->data.F32[j], params_min->data.F32[j]);
        }
    }


    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, "---- %s() begin ----\n", __func__);
    // allocate internal arrays (current vs Guess)
    psImage  *alpha  = psImageAlloc (params->n, params->n, PS_TYPE_F64);
    psImage  *Alpha  = psImageAlloc (params->n, params->n, PS_TYPE_F64);
    psVector *beta   = psVectorAlloc(params->n, PS_TYPE_F64);
    beta->n = beta->nalloc;
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F64);
    Params->n = Params->nalloc;
    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);
        dy->n = dy->nalloc;
        psVectorInit(dy, 1.0);
    }

    psF64 rcF64 = p_psMinLM_SetABX(alpha, beta, params, paramMask, x, y, dy, func);
    if (isnan(rcF64)) {
        psTrace ("psLib.math", 5, "p_psMinLM_SetABX() returned a NAN.\n");
        rc = false;
    }
    psTrace("psLib.math", 5, "p_psMinLM_SetABX() was succesful\n");

    psBool rcBool = p_psMinLM_GuessABP(Alpha, delta, Params, alpha, beta, params, paramMask, NULL, NULL, 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");

    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
psF64 p_psMinLM_dLinear(
    const psVector *Beta,
    const psVector *beta,
    psF64 lambda)
{

    /* get linear model prediction */
    psF64 dLinear = 0;
    psF64 *B = Beta->data.F64;
    psF64 *b = beta->data.F64;
    for (int i = 0; i < beta->n; i++) {
        dLinear += lambda*PS_SQR(B[i]) + B[i]*b[i];
    }
    return(0.5*dLinear);
}

// XXX EAM: this needs to respect the mask on params
// alpha, beta, params are already allocated
psF64 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);

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

    // 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.F64[j][k] = 0;
        }
        beta->data.F64[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.F64[j][k] += weight * deriv->data.F32[k];
            }
            beta->data.F64[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.F64[k][j] = alpha->data.F64[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.F64[j][j] = 1;
                beta->data.F64[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.
 
NOTES: EAM: this is my re-implementation of MinLM
 
XXX: EAM : constraints added 2006.02.15
 
XXX: Put the ASSERTS in.
 
XXX: This must work for both F32 and F64.  F32 is currently implemented.
     Note: since the LUD routines are only implemented in F64, then we
     will have to convert all F32 input vectors to F64 regardless.  So,
     the F64 port might be.
 
  *****************************************************************************/
psBool psMinimizeLMChi2(
    psMinimization *min,
    psImage *covar,
    psVector *params,
    psMinConstrain *constrain,
    const psArray *x,
    const psVector *y,
    const psVector *yWt,
    psMinimizeLMChi2Func func)
{
    psTrace("psLib.math", 3, "---- %s() begin ----\n", __func__);
    PS_ASSERT_PTR_NON_NULL(min, false);
    // XXX: If covar not NULL, do asserts...
    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;
    psVector *paramDelta = NULL;
    psVector *paramMin = NULL;
    psVector *paramMax = NULL;
    if (constrain != NULL) {
        // XXX EAM : fill out the asserts
        paramDelta = constrain->paramDelta;
        paramMin = constrain->paramMin;
        paramMax = constrain->paramMax;
        paramMask = constrain->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);

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

    // why is this needed here??? the initial guess on params is provided by the user
    Params = psVectorCopy(Params, params, PS_TYPE_F32);

    // 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);
        dy->n = dy->nalloc;
        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(__func__) >= 6) {
        p_psImagePrint(psTraceGetDestination(), alpha, "alpha guess (0)");
        p_psVectorPrint(psTraceGetDestination(), beta, "beta guess (0)");
        p_psVectorPrint(psTraceGetDestination(), params, "params guess (0)");
    }
    if (psTraceGetLevel (__func__) >= 6) {
        psTrace("psLib.math", 6, "The current Param vector: \n");
        for (psS32 i = 0 ; i < Params->n ; i++) {
            psTrace("psLib.math", 6, "Params[%d] is %f\n", i, Params->data.F32[i]);
        }
    }

    // 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,
                                paramDelta, paramMin, paramMax, lambda)) {
            min->iter ++;
            lambda *= 10.0;
            continue;
        }

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

        // dump some useful info if trace is defined
        if (psTraceGetLevel(__func__) >= 6) {
            p_psImagePrint(psTraceGetDestination(), Alpha, "alpha guess (1)");
            p_psVectorPrint(psTraceGetDestination(), Beta, "beta guess (1)");
            p_psVectorPrint(psTraceGetDestination(), Params, "params guess (1)");
        }
        if (psTraceGetLevel(__func__) >= 6) {
            if (psTraceGetLevel (__func__) >= 6) {
                psTrace("psLib.math", 6, "The current Param vector: \n");
                for (psS32 i = 0 ; i < Params->n ; i++) {
                    psTrace("psLib.math", 6, "Params[%d] is %f\n", i, Params->data.F32[i]);
                }
            }
        }

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

        // XXX EAM alternate convergence criterion:
        // compare the delta (min->value - Chisq) with the
        // expected delta from the linear model (dLinear)
        // accept new guess (if improvement), or increase lambda
        psF64 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(__func__) >= 6) {
            p_psImagePrint(psTraceGetDestination(), Alpha, "alpha guess (2)");
            p_psVectorPrint(psTraceGetDestination(), Beta, "beta guess (2)");
            p_psVectorPrint(psTraceGetDestination(), Params, "params 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_F64);
            beta   = psVectorCopy(beta, Beta, PS_TYPE_F64);
            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,
                                paramDelta, paramMin, paramMax, 0.0)) {
            psTrace (__func__, 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, "---- %s(false) end ----\n", __func__);
        return(false);
    }
    psTrace("psLib.math", 3, "---- %s(true) end ----\n", __func__);
    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)
{
    if (!is_psType(ptr)) {
        return false;
    }
    return( psMemGetDeallocator(ptr) == (psFreeFunc)minimizationFree );
}


static void constrainFree(psMinConstrain *tmp)
{
    // There are no dynamically allocated items
}

psMinConstrain* psMinConstrainAlloc()
{
    psMinConstrain *tmp = psAlloc(sizeof(psMinConstrain));
    tmp->paramMask = NULL;
    tmp->paramMax = NULL;
    tmp->paramMin = NULL;
    tmp->paramDelta = NULL;

    return(tmp);
}

bool psMemCheckConstrain(psPtr tmp)
{
    return(psMemGetDeallocator(tmp) == (psFreeFunc) constrainFree);
}
