/** @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.5 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-02-07 23:14:21 $
 *
 *  Copyright 2004-2005 Maui High Performance Computing Center, University of Hawaii
 *
 */
/*****************************************************************************/
/* INCLUDE FILES                                                             */
/*****************************************************************************/
#include <stdio.h>
#include <float.h>
#include <math.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 1
    # if (USE_LU_DECOMP)
        psVector *LUv = NULL;
    psImage  *LUm = NULL;
    psImage  *A   = NULL;
    psF32    det;

    // LU decomposition version
    psTrace(__func__, 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(__func__, 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 == psGaussJordan(Alpha, Beta)) {
        psError(PS_ERR_UNKNOWN, false, "psMatrixInvert() returned NULL\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(__func__, 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);
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F64);
    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);
    }

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

    psBool rcBool = p_psMinLM_GuessABP(Alpha, delta, Params, alpha, beta, params, paramMask, NULL, NULL, NULL, 0.0);
    if (rcBool == false) {
        psError(PS_ERR_UNKNOWN, false, "p_psMinLM_GuessABP() retruned FALSE.\n");
        rc = false;
    }
    psTrace(__func__, 5, "p_psMinLM_GuessABP() was succesful\n", __func__);

    psFree(alpha);
    psFree(Alpha);
    psFree(beta);
    psFree(Params);
    if (yWt == NULL) {
        psFree(dy);
    }
    psTrace(__func__, 3, "---- %s() end ----\n", __func__);
    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);

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

        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: Must implement code to handle the constrain->min, ->max, ->delta members.
 
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(__func__, 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;
    if (constrain != NULL) {
        paramMask = constrain->paramMask;
        if (paramMask != NULL) {
            PS_ASSERT_VECTOR_TYPE(paramMask, PS_TYPE_U8, false);
            PS_ASSERT_VECTORS_SIZE_EQUAL(params, paramMask, false);
        }
        if (constrain->paramMin != NULL) {
            psLogMsg(__func__, PS_LOG_WARN, "WARNING: the constrain->paramMin vector is currently ignored.\n");
        }
        if (constrain->paramMax != NULL) {
            psLogMsg(__func__, PS_LOG_WARN, "WARNING: the constrain->paramMax vector is currently ignored.\n");
        }
        if (constrain->paramDelta != NULL) {
            psLogMsg(__func__, PS_LOG_WARN, "WARNING: the constrain->paramDelta vector is currently ignored.\n");
        }
    }
    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);
    psVector *Beta   = psVectorAlloc(params->n, PS_TYPE_F64);
    psVector *Params = psVectorAlloc(params->n, PS_TYPE_F32);
    psVector *dy     = NULL;
    psF64 Chisq = 0.0;
    psF64 lambda = 0.001;
    // XXX: Code this properly.  Don't use mustFree00.
    psBool mustFree00 = false;
    psVector *beta_lim = NULL;
    psVector *param_min = NULL;
    psVector *param_max = NULL;

    // if we are provided a covar image, we expect to find these three vectors in first three rows
    if (covar != NULL) {
        mustFree00 = true;
        beta_lim  = psVectorAlloc(params->n, PS_TYPE_F32);
        param_min = psVectorAlloc(params->n, PS_TYPE_F32);
        param_max = psVectorAlloc(params->n, PS_TYPE_F32);
        for (int i = 0; i < params->n; i++) {
            beta_lim->data.F32[i] = covar->data.F64[0][i];
            param_min->data.F32[i] = covar->data.F64[1][i];
            param_max->data.F32[i] = covar->data.F64[2][i];
        }
        psImageRecycle(covar, params->n, params->n, PS_TYPE_F64);
    }

    // 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);
        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(__func__, 6, "The current Param vector: \n");
        for (psS32 i = 0 ; i < Params->n ; i++) {
            psTrace(__func__, 6, "Params[%d] is %f\n", 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(__func__, 5, "Iteration number %d.  (max iterations is %d).\n", min->iter, min->maxIter);
        psTrace(__func__, 5, "Last delta is %f.  Min->tol is %f.\n", min->lastDelta, min->tol);

        // set a new guess for Alpha, Beta, Params
        p_psMinLM_GuessABP(Alpha, Beta, Params, alpha, beta, params, paramMask,
                           beta_lim, param_min, param_max, lambda);

        // 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(__func__, 6, "The current Param vector: \n");
                for (psS32 i = 0 ; i < Params->n ; i++) {
                    psTrace(__func__, 6, "Params[%d] is %f\n", Params->data.F32[i]);
                }
            }
        }

        // calculate Chisq for new guess, update Alpha & Beta
        Chisq = p_psMinLM_SetABX(Alpha, Beta, Params, paramMask, x, y, dy, func);

        // 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(__func__, 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(__func__, 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) {
        p_psMinLM_GuessABP(covar, Beta, Params, alpha, beta, params, paramMask,
                           beta_lim, param_min, param_max, 0.0);
    }

    // free the internal temporary data
    psFree(alpha);
    psFree(Alpha);
    psFree(beta);
    psFree(Beta);
    psFree(Params);
    if (yWt == NULL) {
        psFree(dy);
    }
    if (mustFree00 == true) {
        psFree(beta_lim);
        psFree(param_min);
        psFree(param_max);
    }
    if (min->iter == min->maxIter) {
        psTrace(__func__, 3, "---- %s(false) end ----\n", __func__);
        return(false);
    }
    psTrace(__func__, 3, "---- %s(true) end ----\n", __func__);
    return(true);
}

// This is a temporary gauss-jordan solver based on gene's
// version based on the Numerical Recipes version
bool psGaussJordan(
    psImage *a,
    psVector *b)
{
    int *indxc,*indxr,*ipiv;
    int Nx, icol, irow;
    int i, j, k, l, ll;
    float big, dum, pivinv;
    psF64 *vector;
    psF64 **matrix;

    Nx = a->numCols;
    matrix = a->data.F64;
    vector = b->data.F64;

    indxc = psAlloc(Nx*sizeof(int));
    indxr = psAlloc(Nx*sizeof(int));
    ipiv  = psAlloc(Nx*sizeof(int));
    for (j = 0; j < Nx; j++) {
        ipiv[j] = 0;
    }

    irow = icol = 0;
    big = fabs(matrix[0][0]);

    for (i = 0; i < Nx; i++) {
        big = 0.0;
        for (j = 0; j < Nx; j++) {
            if (!isfinite(matrix[i][j])) {
                psError(PS_ERR_UNKNOWN, false, "Input matrix contains NaNs: matrix[%d][%d] is %.2f\n", i, j, matrix[i][j]);
                goto fescape;
            }
            if (ipiv[j] != 1) {
                for (k = 0; k < Nx; k++) {
                    if (ipiv[k] == 0) {
                        if (fabs (matrix[j][k]) >= big) {
                            big  = fabs(matrix[j][k]);
                            irow = j;
                            icol = k;
                        }
                    } else {
                        if (ipiv[k] > 1) {
                            psError(PS_ERR_UNKNOWN, false, "Singular Matrix (1).\n");
                            goto fescape;
                        }
                    }
                }
            }
        }
        ipiv[icol]++;
        if (irow != icol) {
            for (l = 0; l < Nx; l++) {
                PS_SWAP(matrix[irow][l], matrix[icol][l]);
            }
            PS_SWAP(vector[irow], vector[icol]);
        }
        indxr[i] = irow;
        indxc[i] = icol;
        if (matrix[icol][icol] == 0.0) {
            psError(PS_ERR_UNKNOWN, false, "Singular Matrix (2).\n");
            goto fescape;
        }
        pivinv = 1.0 / matrix[icol][icol];
        matrix[icol][icol] = 1.0;
        for (l = 0; l < Nx; l++) {
            matrix[icol][l] *= pivinv;
        }
        vector[icol] *= pivinv;

        for (ll = 0; ll < Nx; ll++) {
            if (ll != icol) {
                dum = matrix[ll][icol];
                matrix[ll][icol] = 0.0;
                for (l = 0; l < Nx; l++) {
                    matrix[ll][l] -= matrix[icol][l]*dum;
                }
                vector[ll] -= vector[icol]*dum;
            }
        }
    }

    for (l = Nx - 1; l >= 0; l--) {
        if (indxr[l] != indxc[l]) {
            for (k = 0; k < Nx; k++) {
                PS_SWAP(matrix[k][indxr[l]], matrix[k][indxc[l]]);
            }
        }
    }
    psFree(ipiv);
    psFree(indxr);
    psFree(indxc);
    return(true);

fescape:
    psFree(ipiv);
    psFree(indxr);
    psFree(indxc);
    return(false);
}

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