/****

This file contains the LM functions from my cycle 5 modifications and notes.

- I added the psMinimizeGaussNewtonDelta function which is used to judge the 
  error on a parameter which is not being fitted 

- I added the p_psMinLM_dLinear function which is used to perform the gain
  factor test for convergence

- I changed Params in psMinimizeLMChi2 to type F32 from F64 since it must match 
  the type of the input vector params 

- I adjusted the tracing levels to make the more detailed information a higher 
  level of tracing.

- I added the paramMask arguments to the GuessAPB and SetABX functions

- I added the gain ratio test for convergence (better defined than the simple 
  delta chisq threshold)

- I added the construction of the covariance matrix at the end of psMinimizeLMChi2 
  (if requested with covar != NULL)

- I added a test of convergence and set the output bool to false if convergence 
  is not reached.

****/

// measure the distance to the minimum assuming a linear model
bool psMinimizeGaussNewtonDelta (psVector *delta,
				 const psVector *params,
				 const psVector *paramMask,
				 const psArray  *x,
				 const psVector *y,
				 const psVector *yErr,
				 psMinimizeLMChi2Func 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     = psVectorAlloc (y->n, PS_TYPE_F32);

    // the user provides the error or NULL.  we need to convert
    // to appropriate weights
    if (yErr != NULL) {
        for (int i = 0; i < dy->n; i++) {
            dy->data.F32[i] = 1.0 / PS_SQR (yErr->data.F32[i]);
        }
    } else {
        for (int i = 0; i < dy->n; i++) {
            dy->data.F32[i] = 1.0;
        }
    }

    p_psMinLM_SetABX (alpha, beta, params, paramMask, x, y, dy, func);
    p_psMinLM_GuessABP (Alpha, delta, Params, alpha, beta, params, paramMask, 0.0);

    psFree (alpha);
    psFree (Alpha);
    psFree (beta);
    psFree (Params);
    psFree (dy);
    return (true);
}

// 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 is my re-implementation of MinLM
psBool psMinimizeLMChi2(psMinimization *min,
                        psImage *covar,
                        psVector *params,
                        const psVector *paramMask,
                        const psArray *x,
                        const psVector *y,
                        const psVector *yErr,
                        psMinimizeLMChi2Func func)
{
    PS_PTR_CHECK_NULL(min, NULL);
    PS_VECTOR_CHECK_NULL(params, NULL);
    PS_VECTOR_CHECK_EMPTY(params, NULL);
    PS_PTR_CHECK_NULL(x, NULL);
    PS_VECTOR_CHECK_NULL(y, NULL);
    PS_VECTOR_CHECK_EMPTY(y, NULL);
    PS_VECTOR_CHECK_SIZE_EQUAL(x, y, NULL);
    PS_PTR_CHECK_NULL(func, NULL);

    // 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 EAM: why is this needed here? the value is not used, and the memory 
    // is allocated above.  However, if I drop it, I get weird answers or crashes.
    Params = psVectorCopy (Params, params, PS_TYPE_F32);

    // the user provides the error or NULL.  we need to convert
    // to appropriate weights
    dy = psVectorAlloc (y->n, PS_TYPE_F32);
    if (yErr != NULL) {
        for (int i = 0; i < dy->n; i++) {
            dy->data.F32[i] = 1.0 / PS_SQR (yErr->data.F32[i]);
        }
    } else {
        for (int i = 0; i < dy->n; i++) {
            dy->data.F32[i] = 1.0;
        }
    }

    // calculate initial alpha and beta, set chisq (min->value)
    min->value = p_psMinLM_SetABX (alpha, beta, params, paramMask, x, y, dy, func);
    # ifndef PS_NO_TRACE
    // dump some useful info if trace is defined
    if (psTraceGetLevel (".psLib.dataManip.psMinimizeLMChi2") >= 5) {
        p_psImagePrint  (psTraceGetDestination(), alpha, "alpha guess");
        p_psVectorPrint (psTraceGetDestination(), beta, "beta guess");
        p_psVectorPrint (psTraceGetDestination(), params, "params guess");
    }
    if (psTraceGetLevel (".psLib.dataManip.psMinimizeLMChi2") == 4) {
	p_psVectorPrintRow (psTraceGetDestination(), Params, "params guess");
    }
    # endif /* PS_NO_TRACE */

    // iterate until the tolerance is reached, or give up
    while ((min->lastDelta > min->tol) && (min->iter < min->maxIter)) {

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

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

        # ifndef PS_NO_TRACE
        // dump some useful info if trace is defined
        if (psTraceGetLevel (".psLib.dataManip.psMinimizeLMChi2") >= 5) {
            p_psImagePrint  (psTraceGetDestination(), Alpha, "alpha guess");
            p_psVectorPrint (psTraceGetDestination(), Beta, "beta guess");
            p_psVectorPrint (psTraceGetDestination(), Params, "params guess");
        }
        if (psTraceGetLevel (".psLib.dataManip.psMinimizeLMChi2") == 4) {
            p_psVectorPrintRow (psTraceGetDestination(), Params, "params guess");
        }
        # endif /* PS_NO_TRACE */

        // 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 (".psLib.dataManip.psMinimizeLMChi2", 4, "last chisq: %f, new chisq %f, delta: %f, rho: %f\n", min->value, Chisq, min->lastDelta, rho);
        # ifndef PS_NO_TRACE
        // dump some useful info if trace is defined
        if (psTraceGetLevel (".psLib.dataManip.psMinimizeLMChi2") >= 5) {
            p_psImagePrint  (psTraceGetDestination(), Alpha, "alpha guess");
            p_psVectorPrint (psTraceGetDestination(), Beta, "beta guess");
            p_psVectorPrint (psTraceGetDestination(), Params, "params guess");
        }
        # endif /* PS_NO_TRACE */

        /* 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.dataManip.psMinimizeLMChi2", 3, "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, 0.0);
    }      

    // free the internal temporary data
    psFree (alpha);
    psFree (Alpha);
    psFree (beta);
    psFree (Beta);
    psFree (Params);
    psFree (dy);

    if (min->iter == min->maxIter) {
      return (false);
    } 
    return (true);
}

// XXX EAM: this needs to respect the mask on params
// XXX EAM: check not NULL on alpha, beta, 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)
{

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

    // zero alpha and beta for summing below
    for (int j = 0; j < params->n; j++) {
        for (int 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 (int 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 (int j = 0; j < params->n; j++) {
	    if ((paramMask != NULL) && (paramMask->data.U8[j])) continue;
            weight = deriv->data.F32[j] * dy->data.F32[i];
            for (int 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 (int j = 1; j < params->n; j++) {
        for (int 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 (int 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);
}

// 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,
                           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 (".psLib.dataManip.psMinLM_GuessABP", 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);
    Beta  = psMatrixLUSolve (Beta, LUm, beta, LUv);
    Alpha = psMatrixInvert (Alpha, A, &det);

    # else
    // gauss-jordan version
    psTrace (".psLib.dataManip.psMinLM_GuessABP", 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);
    }

    psGaussJordan (Alpha, Beta);
    # 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];
    }

    # if (USE_LU_DECOMP)
    psFree (A);
    psFree (LUm);
    psFree (LUv);
    # endif

    return true;
}