# include "psphot.h"

bool psphotModelWithPSF_LMM (
    psMinimization *min,
    psImage *covar,
    psVector *params,
    psMinConstraint *constraint,
    pmSource *source,
    const psKernel *psf,
    psMinimizeLMChi2Func func)
{
    psTrace("psphot", 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(func, false);
    PS_ASSERT_PTR_NON_NULL(source, false);

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

    // this function has test values and current values for several things
    // the current value is in lower case
    // the test 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);

    psF32 Chisq = 0.0;
    psF32 lambda = 0.001;
    psF32 dLinear = 0.0;

    // generate PCM data storage structure
    pmPCMData *pcm = pmPCMDataAlloc (params, paramMask, source);

    // calculate initial alpha and beta, set chisq (min->value)
    min->value = psphotModelWithPSF_SetABX(alpha, beta, params, paramMask, pcm, source, psf, func);
    if (isnan(min->value)) {
        min->iter = min->maxIter;
        return(false);
    }
    // dump some useful info if trace is defined
    if (psTraceGetLevel("psphot") >= 6) {
        p_psImagePrint(psTraceGetDestination(), alpha, "alpha guess (0)");
        p_psVectorPrint(psTraceGetDestination(), beta, "beta guess (0)");
    }
    if (psTraceGetLevel("psphot") >= 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("psphot", 5, "Iteration number %d.  (max iterations is %d).\n", min->iter, min->maxIter);
        psTrace("psphot", 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("psphot") >= 6) {
            p_psImagePrint(psTraceGetDestination(), Alpha, "alpha guess (1)");
            p_psVectorPrint(psTraceGetDestination(), Beta, "beta guess (1)");
        }
        if (psTraceGetLevel("psphot") >= 5) {
            p_psVectorPrint(psTraceGetDestination(), Params, "params guess (1)");
        }

        // calculate Chisq for new guess, update Alpha & Beta
        Chisq = psphotModelWithPSF_SetABX(Alpha, Beta, Params, paramMask, pcm, source, psf, 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("psphot", 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("psphot") >= 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) / (source->pixels->numCols*source->pixels->numRows - 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;

	    // save the new convolved model image
	    psFree (source->modelFlux);
	    source->modelFlux = pmPCMDataSaveImage(pcm);
        } else {
            lambda *= 10.0;
        }
        min->iter++;
    }
    psTrace("psphot", 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(covar, Beta, Params, alpha, beta, params, paramMask, NULL, 0.0, NULL)) {
            psTrace ("psphot", 5, "failure to calculate covariance matrix\n");
        }
    }

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

    if (min->iter == min->maxIter) {
        psTrace("psphot", 3, "---- end (false) ----\n");
        return(false);
    }

    psTrace("psphot", 3, "---- end (true) ----\n");
    return(true);
}

psF32 psphotModelWithPSF_SetABX(
    psImage  *alpha,
    psVector *beta,
    const psVector *params,
    const psVector *paramMask,
    pmPCMData *pcm,
    const pmSource *source,
    const psKernel *psf,
    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(source, NAN);
    PS_ASSERT_IMAGE_NON_NULL(source->pixels, NAN);
    PS_ASSERT_IMAGE_NON_NULL(source->weight, NAN);
    PS_ASSERT_IMAGE_NON_NULL(source->maskObj, NAN);

    PS_ASSERT_VECTOR_TYPE(params, PS_TYPE_F32, false);
    if (paramMask) {
        PS_ASSERT_VECTOR_TYPE(paramMask, PS_TYPE_MASK, false);
    }

    // generate the model and derivative images for this parameter set

    // storage for model derivatives
    psVector *deriv = psVectorAlloc(params->n, PS_TYPE_F32);

    // working vector to store local coordinate
    psVector *coord = psVectorAlloc(2, PS_TYPE_F32);

    psImageInit (pcm->model, 0.0);
    for (int n = 0; n < params->n; n++) {
	if (!pcm->dmodels->data[n]) continue;
	psImageInit (pcm->dmodels->data[n], 0.0);
    }

    // fill in the coordinate and value entries
    for (psS32 i = 0; i < source->pixels->numRows; i++) {
        for (psS32 j = 0; j < source->pixels->numCols; j++) {

	    // XXX can we skip some of the data points where the model
	    // is not going to be fitted??

            // skip masked points
	    // XXX probably should not skipped masked points: 
	    // XXX skip if convolution of unmasked pixels will not see this pixel
            // if (source->maskObj->data.U8[i][j]) {
	    // continue;
	    // }

            // skip zero-weight points
	    // XXX why is this not masked?
            // if (source->weight->data.F32[i][j] == 0) {
	    // continue;
	    // }
            // skip nan value points
	    // XXX why is this not masked?
            // if (!isfinite(source->pixels->data.F32[i][j])) {
	    // continue;
	    // }

            // Convert i/j to image space:
            coord->data.F32[0] = (psF32) (j + source->pixels->col0);
            coord->data.F32[1] = (psF32) (i + source->pixels->row0);

	    pcm->model->data.F32[i][j] = func (deriv, params, coord);

	    for (int n = 0; n < params->n; n++) {
	      if ((paramMask != NULL) && (paramMask->data.U8[n])) { continue; }
	      psImage *dmodel = pcm->dmodels->data[n];
	      dmodel->data.F32[i][j] = deriv->data.F32[n];
	    }
        }
    }
    psFree(coord);
    psFree(deriv);


    // convolve model and dmodel arrays with PSF
    psImageConvolveDirect (pcm->modelConv, pcm->model, psf);
    for (int n = 0; n < pcm->dmodels->n; n++) {
      if (pcm->dmodels->data[n] == NULL) continue;
      psImage *dmodel = pcm->dmodels->data[n];
      psImage *dmodelConv = pcm->dmodelsConv->data[n];
      psImageConvolveDirect (dmodelConv, dmodel, psf);
    }

    // XXX TEST : SAVE IMAGES
# if (1) 
    psphotSaveImage (NULL, psf->image, "psf.fits");
    psphotSaveImage (NULL, pcm->model, "model.fits");
    psphotSaveImage (NULL, pcm->modelConv, "modelConv.fits");
    psphotSaveImage (NULL, source->pixels, "obj.fits");
    psphotSaveImage (NULL, source->maskObj, "mask.fits");
    psphotSaveImage (NULL, source->weight, "weight.fits");
# endif

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

    for (psS32 i = 0; i < source->pixels->numRows; i++) {
        for (psS32 j = 0; j < source->pixels->numCols; j++) {
	  // XXX are we doing the right thing with the mask?
            // skip masked points
            if (source->maskObj->data.U8[i][j]) {
                continue;
            }
            // skip zero-weight points
            if (source->weight->data.F32[i][j] == 0) {
                continue;
            }
            // skip nan value points
            if (!isfinite(source->pixels->data.F32[i][j])) {
                continue;
            }

	    float ymodel  = pcm->modelConv->data.F32[i][j];
	    float yweight = 1.0 / source->weight->data.F32[i][j];
	    float delta = ymodel - source->pixels->data.F32[i][j];

	    chisq += PS_SQR(delta) * yweight;

	    if (isnan(delta))
	      psAbort("nan in delta");
	    if (isnan(chisq))
	      psAbort("nan in chisq");

	    for (psS32 n1 = 0; n1 < params->n; n1++) {
	      if ((paramMask != NULL) && (paramMask->data.U8[n1])) {
                continue;
	      }
	      psImage *dmodel = pcm->dmodelsConv->data[n1];
	      float weight = dmodel->data.F32[i][j] * yweight;
	      for (psS32 n2 = 0; n2 <= n1; n2++) {
                if ((paramMask != NULL) && (paramMask->data.U8[n2])) {
		  continue;
                }
		dmodel = pcm->dmodelsConv->data[n2];
                alpha->data.F32[n1][n2] += weight * dmodel->data.F32[i][j];
	      }
	      beta->data.F32[n1] += 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;
            }
        }
    }

    return(chisq);
}

static void pmPCMDataFree (pmPCMData *pcm) {

    if (pcm == NULL) return;

    psFree (pcm->model);
    psFree (pcm->modelConv);
    psFree (pcm->dmodels);
    psFree (pcm->dmodelsConv);
    return;
}

pmPCMData *pmPCMDataAlloc (
    const psVector *params,
    const psVector *paramMask,
    pmSource *source) {

    pmPCMData *pcm = (pmPCMData *) psAlloc(sizeof(pmPCMData));
    psMemSetDeallocator(pcm, (psFreeFunc) pmPCMDataFree);

    // Allocate storage images for raw model and derivative images
    pcm->model = psImageCopy (NULL, source->pixels, PS_TYPE_F32);
    pcm->dmodels = psArrayAlloc (params->n);
    for (psS32 n = 0; n < params->n; n++) {
      if ((paramMask != NULL) && (paramMask->data.U8[n])) { continue; }
      pcm->dmodels->data[n] = psImageCopy (NULL, source->pixels, PS_TYPE_F32);
    }

    // Allocate storage images for convolved model and derivative images
    pcm->modelConv = psImageCopy (NULL, source->pixels, PS_TYPE_F32);
    pcm->dmodelsConv = psArrayAlloc (params->n);
    for (psS32 n = 0; n < params->n; n++) {
      if ((paramMask != NULL) && (paramMask->data.U8[n])) { continue; }
      pcm->dmodelsConv->data[n] = psImageCopy (NULL, source->pixels, PS_TYPE_F32);
    }

    return pcm;
}

psImage *pmPCMDataSaveImage (pmPCMData *pcm) {

    psImage *model = psImageCopy (NULL, pcm->modelConv, PS_TYPE_F32);
   
    return model;
}

/*
 *
 * we have a function func(param; value)
 
 * basic LMM:
 
 - fill in the data (x, y)
 
 chisq = SetABX (alpha, beta, params, paramMask, x, y, dy, func)
 
 while () {
   GuessABP (Alpha, Beta, Params, alpha, beta, params, paramMask, checkLimits, lambda)
   dLinear = dLinear(Beta, beta, lambda);
   chisq = SetABX (alpha, beta, params, paramMask, x, y, dy, func)
   convergence tests...
 }
 
 

 ** GuessABP:

      f_c = sum_i (kern_i * func (x_i; p_o))

df_c/dp_o = d/dp_o [sum_i (kern_i * func (x_i; p_o))]

df_c/dp_o = sum_i (d/dp_o [kern_i * func (x_i; p_o)])

df_c/dp_o = sum_i (kern_i * d/dp_o [func (x_i; p_o)])

- generate image arrays for func, dfunc/dp_j (not masked)
- convolve each with psf
- measure delta = f_conv - data
- etc
*/


