/** @file  pmSourceFitModel.c
 *
 *  fit single source models to image pixels
 *
 *  @author EAM, IfA
 *  @author GLG, MHPCC
 *
 *  @version $Revision: 1.7 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-06-21 03:21:16 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

#include <stdio.h>
#include <math.h>
#include <string.h>
#include <pslib.h>
#include "pmHDU.h"
#include "pmFPA.h"
#include "pmPeaks.h"
#include "pmMoments.h"
#include "pmGrowthCurve.h"
#include "pmModel.h"
#include "pmPSF.h"
#include "pmSource.h"
#include "pmModelGroup.h"
#include "pmSourceFitModel.h"

// save as static values so they may be set externally
static psF32 PM_SOURCE_FIT_MODEL_NUM_ITERATIONS = 15;
static psF32 PM_SOURCE_FIT_MODEL_TOLERANCE = 0.1;
static bool PM_PSF_POISSON_WEIGHTS = true;

bool pmSourceFitModelInit (float nIter, float tol, bool poissonErrors)
{

    PM_SOURCE_FIT_MODEL_NUM_ITERATIONS = nIter;
    PM_SOURCE_FIT_MODEL_TOLERANCE = tol;
    PM_PSF_POISSON_WEIGHTS = poissonErrors;

    return true;
}

bool pmSourceFitModel (pmSource *source,
                       pmModel *model,
                       pmSourceFitMode mode)
{
    psTrace(__func__, 3, "---- %s() begin ----\n", __func__);
    PS_ASSERT_PTR_NON_NULL(source, false);
    PS_ASSERT_PTR_NON_NULL(source->moments, false);
    PS_ASSERT_PTR_NON_NULL(source->peak, false);
    PS_ASSERT_PTR_NON_NULL(source->pixels, false);
    PS_ASSERT_PTR_NON_NULL(source->mask, false);
    PS_ASSERT_PTR_NON_NULL(source->weight, false);

    // XXX EAM : is it necessary for the mask & weight to exist?  the
    //           tests below could be conditions (!NULL)

    psBool fitStatus = true;
    psBool onPic     = true;
    psBool rc        = true;

    psVector *params = model->params;
    psVector *dparams = model->dparams;
    psVector *paramMask = NULL;

    pmModelFunc modelFunc = pmModelFunc_GetFunction (model->type);

    psF32 dSky = source->moments->dSky;

    // maximum number of valid pixels
    psS32 nPix = source->pixels->numRows * source->pixels->numCols;

    // construct the coordinate and value entries
    psArray *x = psArrayAlloc(nPix);
    psVector *y = psVectorAlloc(nPix, PS_TYPE_F32);
    psVector *yErr = psVectorAlloc(nPix, PS_TYPE_F32);

    nPix = 0;
    for (psS32 i = 0; i < source->pixels->numRows; i++) {
        for (psS32 j = 0; j < source->pixels->numCols; j++) {
            // skip masked points
            if (source->mask->data.U8[i][j]) {
                continue;
            }
            // skip zero-weight points
            if (source->weight->data.F32[i][j] == 0) {
                continue;
            }

            psVector *coord = psVectorAlloc(2, PS_TYPE_F32);
            coord->n = 2;

            // Convert i/j to image space:
            coord->data.F32[0] = (psF32) (j + source->pixels->col0);
            coord->data.F32[1] = (psF32) (i + source->pixels->row0);
            x->data[nPix] = (psPtr *) coord;
            y->data.F32[nPix] = source->pixels->data.F32[i][j];
            // psMinimizeLMChi2 takes wt = 1/dY^2
            if (PM_PSF_POISSON_WEIGHTS) {
                yErr->data.F32[nPix] = 1.0 / source->weight->data.F32[i][j];
            } else {
                yErr->data.F32[nPix] = 1.0 / dSky;
            }
            // XXX EAM : suggestion from RHL is to use the local sky as weight
            // to avoid the bias from systematic errors
            // here we would just use the source sky variance
            nPix++;
        }
    }
    x->n = nPix;
    y->n = nPix;
    yErr->n = nPix;

    // XXX EAM : the new minimization API supplies the constraints as a struct
    // XXX the chisq if a fcn of source flux. the minimization criterion should
    // probably  be scaled somehow by flux.  measure the trend?  it depends on
    // the about of systematic error in the models themselves...
    psMinimization *myMin = psMinimizationAlloc(PM_SOURCE_FIT_MODEL_NUM_ITERATIONS,
                            PM_SOURCE_FIT_MODEL_TOLERANCE);
    psMinConstrain *constrain = psMinConstrainAlloc();

    // set parameter mask based on fitting mode
    paramMask = psVectorAlloc (params->n, PS_TYPE_U8);
    paramMask->n = params->n;
    psVectorInit (paramMask, 1);

    int nParams = 0;
    switch (mode) {
    case PM_SOURCE_FIT_NORM:
        // NORM-only model fits only source normalization (Io)
        nParams = 1;
        paramMask->data.U8[1] = 0;
        break;
    case PM_SOURCE_FIT_PSF:
        // PSF model only fits x,y,Io
        nParams = 3;
        paramMask->data.U8[1] = 0;
        paramMask->data.U8[2] = 0;
        paramMask->data.U8[3] = 0;
        break;
    case PM_SOURCE_FIT_EXT:
        // EXT model fits all params (except sky)
        nParams = params->n - 1;
        psVectorInit (paramMask, 0);
        paramMask->data.U8[0] = 1;
        break;
    case PM_SOURCE_FIT_PSF_AND_SKY:
        nParams = 4;
        psAbort ("pmSourceFitModel", "PSF + SKY not currently available");
        break;
    case PM_SOURCE_FIT_EXT_AND_SKY:
        nParams = params->n;
        psAbort ("pmSourceFitModel", "EXT + SKY not currently available");
        break;
    default:
        psAbort ("pmSourceFitModel", "invalid fitting mode");
    }
    constrain->paramMask = paramMask;

    if (nPix <  nParams + 1) {
        psTrace (".pmObjects.pmSourceFitModel", 4, "insufficient valid pixels\n");
        psTrace(__func__, 3, "---- %s(false) end ----\n", __func__);
        model->status = PM_MODEL_BADARGS;
        psFree (x);
        psFree (y);
        psFree (yErr);
        psFree (myMin);
        psFree (constrain);
        psFree (paramMask);
        return(false);
    }

    // Set the parameter range checks
    pmModelLimits modelLimits = pmModelLimits_GetFunction (model->type);
    modelLimits (&constrain->paramDelta, &constrain->paramMin, &constrain->paramMax);

    // force the floating parameters to fall within the contraint ranges
    for (int i = 0; i < params->n; i++) {
        if (constrain->paramMask->data.U8[i])
            continue;
        params->data.F32[i] = PS_MIN(constrain->paramMax->data.F32[i], params->data.F32[i]);
        params->data.F32[i] = PS_MAX(constrain->paramMin->data.F32[i], params->data.F32[i]);
    }

    psImage *covar = psImageAlloc (params->n, params->n, PS_TYPE_F64);

    psTrace (".pmObjects.pmSourceFitModel", 5, "fitting function\n");

    fitStatus = psMinimizeLMChi2(myMin, covar, params, constrain, x, y, yErr, modelFunc);
    for (int i = 0; i < dparams->n; i++) {
        if ((paramMask != NULL) && paramMask->data.U8[i])
            continue;
        dparams->data.F32[i] = sqrt(covar->data.F64[i][i]);
    }

    // save the resulting chisq, nDOF, nIter
    model->chisq = myMin->value;
    model->nIter = myMin->iter;
    model->nDOF  = y->n - nParams;

    // get the Gauss-Newton distance for fixed model parameters
    if (paramMask != NULL) {
        psVector *delta = psVectorAlloc (params->n, PS_TYPE_F64);
        delta->n = params->n;
        psMinimizeGaussNewtonDelta(delta, params, NULL, x, y, yErr, modelFunc);
        for (int i = 0; i < dparams->n; i++) {
            if (!paramMask->data.U8[i])
                continue;
            dparams->data.F32[i] = delta->data.F64[i];
        }
        psFree (delta);
    }

    // set the model success or failure status
    if (!fitStatus) {
        model->status = PM_MODEL_NONCONVERGE;
    } else {
        model->status = PM_MODEL_SUCCESS;
    }

    // models can go insane: reject these
    onPic &= (params->data.F32[2] >= source->pixels->col0);
    onPic &= (params->data.F32[2] <  source->pixels->col0 + source->pixels->numCols);
    onPic &= (params->data.F32[3] >= source->pixels->row0);
    onPic &= (params->data.F32[3] <  source->pixels->row0 + source->pixels->numRows);
    if (!onPic) {
        model->status = PM_MODEL_OFFIMAGE;
    }

    source->mode |= PM_SOURCE_MODE_FITTED;

    psFree(x);
    psFree(y);
    psFree(yErr);
    psFree(myMin);
    psFree(covar);
    psFree(constrain->paramMask);
    psFree(constrain->paramMin);
    psFree(constrain->paramMax);
    psFree(constrain->paramDelta);
    psFree(constrain);

    rc = (onPic && fitStatus);
    psTrace(__func__, 3, "---- %s(%d) end ----\n", __func__, rc);
    return(rc);
}

/********************* Source Model Set Functions ***************************/

// these static variables are used by one pass of pmSourceFitSet to store
// data for a model set.  If we are going to make psphot thread-safe, these
// will have to go in a structure of their own.
// sky, p1.1, p1.2, p1.3,... p1.n, p2.1, p2.2,
// nPar = nSrc*(nOnePar - 1) + 1
static pmModelFunc mFunc;
static int nPar;
static psVector *onePar;
static psVector *oneDeriv;

bool pmModelFitSetInit (pmModelType type)
{

    mFunc = pmModelFunc_GetFunction (type);
    nPar  = pmModelParameterCount (type);

    onePar = psVectorAlloc (nPar, PS_DATA_F32);
    oneDeriv = psVectorAlloc (nPar, PS_DATA_F32);
    onePar->n = oneDeriv->n = nPar;

    return true;
}

void pmModelFitSetClear (void)
{

    psFree (onePar);
    psFree (oneDeriv);
    return;
}

psF32 pmModelFitSet(psVector *deriv,
                    const psVector *params,
                    const psVector *x)
{

    psF32 value;
    psF32 model;

    psF32 *PAR = onePar->data.F32;
    psF32 *dPAR = oneDeriv->data.F32;

    psF32 *pars = params->data.F32;
    psF32 *dpars = (deriv == NULL) ? NULL : deriv->data.F32;

    int nSrc = (params->n - 1) / (nPar - 1);

    PAR[0] = model = pars[0];
    for (int i = 0; i < nSrc; i++) {
        int nOff = i*nPar - i;
        for (int n = 1; n < nPar; n++) {
            PAR[n] = pars[nOff + n];
        }
        if (deriv == NULL) {
            value = mFunc (NULL, onePar, x);
        } else {
            value = mFunc (oneDeriv, onePar, x);
            for (int n = 1; n < nPar; n++) {
                dpars[nOff + n] = dPAR[n];
            }
        }
        model += value;
    }
    if (deriv != NULL) {
        dpars[0] = dPAR[0]*2.0;
    }
    return (model);
}

/*
i:         0           1               2
n:         1  2  3  4  5  6  1  2  3  4  5  6  1  2  3  4  5  6
i*6 + n: 0 1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18
*/

bool pmSourceFitSet (pmSource *source,
                     psArray *modelSet,
                     pmSourceFitMode mode)
{
    psTrace(__func__, 3, "---- %s() begin ----\n", __func__);
    PS_ASSERT_PTR_NON_NULL(source, false);
    PS_ASSERT_PTR_NON_NULL(source->moments, false);
    PS_ASSERT_PTR_NON_NULL(source->peak, false);
    PS_ASSERT_PTR_NON_NULL(source->pixels, false);
    PS_ASSERT_PTR_NON_NULL(source->mask, false);
    PS_ASSERT_PTR_NON_NULL(source->weight, false);

    psBool fitStatus = true;
    psBool onPic     = true;
    psBool rc        = true;

    // base values on first model
    pmModel *model = modelSet->data[0];

    // set the static variables
    pmModelFitSetInit (model->type);

    int nSrc = modelSet->n;
    int nPar = model->params->n - 1;  // number of object parameters (excluding sky)

    // define parameter vectors for source set
    psVector *params = psVectorAlloc (nSrc*nPar + 1, PS_TYPE_F32);
    psVector *dparams = psVectorAlloc (nSrc*nPar + 1, PS_TYPE_F32);
    params->n = params->nalloc;
    dparams->n = dparams->nalloc;

    pmModelLimits modelLimits = pmModelLimits_GetFunction (model->type);

    // define limits for a single-source model
    psVector *oneDelta;
    psVector *oneParMin;
    psVector *oneParMax;
    modelLimits (&oneDelta, &oneParMin, &oneParMax);

    psMinConstrain *constrain = psMinConstrainAlloc();
    constrain->paramMin = psVectorAlloc (nSrc*nPar + 1, PS_TYPE_F32);
    constrain->paramMax = psVectorAlloc (nSrc*nPar + 1, PS_TYPE_F32);
    constrain->paramMask = psVectorAlloc (nSrc*nPar + 1, PS_TYPE_U8);
    constrain->paramDelta = psVectorAlloc (nSrc*nPar + 1, PS_TYPE_F32);
    constrain->paramMin->n = constrain->paramMin->nalloc;
    constrain->paramMax->n = constrain->paramMax->nalloc;
    constrain->paramMask->n = constrain->paramMask->nalloc;
    constrain->paramDelta->n = constrain->paramDelta->nalloc;

    // set the parameter guesses and constraints for the multiple models
    // first the values for the single sky parameter
    params->data.F32[0] = model->params->data.F32[0];
    constrain->paramMin->data.F32[0]   = oneParMin->data.F32[0];
    constrain->paramMax->data.F32[0]   = oneParMax->data.F32[0];
    constrain->paramDelta->data.F32[0] = oneDelta->data.F32[0];

    // next, the values for the source parameters
    for (int i = 0; i < nSrc; i++) {
        model = modelSet->data[i];
        for (int n = 1; n < nPar + 1; n++) {
            params->data.F32[i*nPar + n] = model->params->data.F32[n];
            constrain->paramDelta->data.F32[i*nPar + n] = oneDelta->data.F32[n];
            constrain->paramMin->data.F32[i*nPar + n]   = oneParMin->data.F32[n];
            constrain->paramMax->data.F32[i*nPar + n]   = oneParMax->data.F32[n];
        }
    }
    psFree (oneDelta);
    psFree (oneParMin);
    psFree (oneParMax);

    if (psTraceGetLevel(__func__) >= 5) {
        for (int i = 0; i < params->n; i++) {
            fprintf (stderr, "%d %f %d\n", i, params->data.F32[i], constrain->paramMask->data.U8[i]);
        }
    }

    // set the parameter masks based on the fitting mode
    int nParams = 0;
    switch (mode) {
    case PM_SOURCE_FIT_NORM:
        // NORM-only model fits only source normalization (Io)
        nParams = nSrc;
        psVectorInit (constrain->paramMask, 1);
        for (int i = 0; i < nSrc; i++) {
            constrain->paramMask->data.U8[1 + i*nPar] = 0;
        }
        break;
    case PM_SOURCE_FIT_PSF:
        // PSF model only fits x,y,Io
        nParams = 3*nSrc;
        psVectorInit (constrain->paramMask, 1);
        for (int i = 0; i < nSrc; i++) {
            constrain->paramMask->data.U8[1 + i*nPar] = 0;
            constrain->paramMask->data.U8[2 + i*nPar] = 0;
            constrain->paramMask->data.U8[3 + i*nPar] = 0;
        }
        break;
    case PM_SOURCE_FIT_EXT:
        // EXT model fits all params (except sky)
        nParams = nPar*nSrc;
        psVectorInit (constrain->paramMask, 0);
        constrain->paramMask->data.U8[0] = 1;
        break;
    case PM_SOURCE_FIT_PSF_AND_SKY:
        nParams = 1 + 3*nSrc;
        psAbort ("pmSourceFitModel", "PSF + SKY not currently available");
        break;
    case PM_SOURCE_FIT_EXT_AND_SKY:
        nParams = 1 + nPar*nSrc;
        psAbort ("pmSourceFitModel", "EXT + SKY not currently available");
        break;
    default:
        psAbort ("pmSourceFitModel", "invalid fitting mode");
    }

    // force the floating parameters to fall within the contraint ranges
    for (int i = 0; i < params->n; i++) {
        if (constrain->paramMask->data.U8[i])
            continue;
        params->data.F32[i] = PS_MIN(constrain->paramMax->data.F32[i], params->data.F32[i]);
        params->data.F32[i] = PS_MAX(constrain->paramMin->data.F32[i], params->data.F32[i]);
    }

    // maximum number of valid pixels
    psS32 nPix = source->pixels->numRows * source->pixels->numCols;

    // local sky variance
    psF32 dSky = source->moments->dSky;

    // construct the coordinate and value entries
    psArray *x = psArrayAlloc(nPix);
    psVector *y = psVectorAlloc(nPix, PS_TYPE_F32);
    psVector *yErr = psVectorAlloc(nPix, PS_TYPE_F32);

    nPix = 0;
    for (psS32 i = 0; i < source->pixels->numRows; i++) {
        for (psS32 j = 0; j < source->pixels->numCols; j++) {
            // skip masked points
            if (source->mask->data.U8[i][j]) {
                continue;
            }
            // skip zero-weight points
            if (source->weight->data.F32[i][j] == 0) {
                continue;
            }
            psVector *coord = psVectorAlloc(2, PS_TYPE_F32);
            coord->n = 2;

            // Convert i/j to image space:
            coord->data.F32[0] = (psF32) (j + source->pixels->col0);
            coord->data.F32[1] = (psF32) (i + source->pixels->row0);
            x->data[nPix] = (psPtr *) coord;
            y->data.F32[nPix] = source->pixels->data.F32[i][j];
            // psMinimizeLMChi2 takes wt = 1/dY^2
            if (PM_PSF_POISSON_WEIGHTS) {
                yErr->data.F32[nPix] = 1.0 / source->weight->data.F32[i][j];
            } else {
                yErr->data.F32[nPix] = 1.0 / dSky;
            }
            nPix++;
        }
    }
    if (nPix <  nParams + 1) {
        psTrace (__func__, 4, "insufficient valid pixels\n");
        psTrace(__func__, 3, "---- %s() end : fail pixels ----\n", __func__);
        model->status = PM_MODEL_BADARGS;
        psFree (x);
        psFree (y);
        psFree (yErr);
        psFree (params);
        psFree (dparams);
        psFree(constrain->paramMask);
        psFree(constrain->paramMin);
        psFree(constrain->paramMax);
        psFree(constrain->paramDelta);
        psFree(constrain);
        pmModelFitSetClear ();
        return(false);
    }
    x->n = nPix;
    y->n = nPix;
    yErr->n = nPix;

    psMinimization *myMin = psMinimizationAlloc(PM_SOURCE_FIT_MODEL_NUM_ITERATIONS,
                            PM_SOURCE_FIT_MODEL_TOLERANCE);

    psImage *covar = psImageAlloc (params->n, params->n, PS_TYPE_F64);

    psTrace (__func__, 5, "fitting function\n");
    fitStatus = psMinimizeLMChi2(myMin, covar, params, constrain, x, y, yErr, pmModelFitSet);
    if (!fitStatus) {
        // psError(PS_ERR_UNKNOWN, false, "Failed to fit model (%d)\n", nSrc);
        psTrace(__func__, 4, "Failed to fit model (%d)\n", nSrc);
    }

    // parameter errors from the covariance matrix
    for (int i = 0; i < dparams->n; i++) {
        if ((constrain->paramMask != NULL) && constrain->paramMask->data.U8[i])
            continue;
        dparams->data.F32[i] = sqrt(covar->data.F64[i][i]);
    }

    // get the Gauss-Newton distance for fixed model parameters
    if (constrain->paramMask != NULL) {
        psVector *delta = psVectorAlloc (params->n, PS_TYPE_F64);
        delta->n = params->n;
        psMinimizeGaussNewtonDelta(delta, params, NULL, x, y, yErr, pmModelFitSet);
        for (int i = 0; i < dparams->n; i++) {
            if (!constrain->paramMask->data.U8[i])
                continue;
            dparams->data.F32[i] = delta->data.F64[i];
        }
        psFree (delta);
    }

    // assign back the parameters to the models
    for (int i = 0; i < nSrc; i++) {
        model = modelSet->data[i];
        model->params->data.F32[0] = params->data.F32[0];
        for (int n = 1; n < nPar + 1; n++) {
            model->params->data.F32[n] = params->data.F32[i*nPar + n];
            model->dparams->data.F32[n] = dparams->data.F32[i*nPar + n];
        }
        // save the resulting chisq, nDOF, nIter
        // these are not unique for any one source
        model->chisq = myMin->value;
        model->nIter = myMin->iter;
        model->nDOF  = y->n - nParams;

        // set the model success or failure status
        model->status = fitStatus ? PM_MODEL_SUCCESS : PM_MODEL_NONCONVERGE;

        // models can go insane: reject these
        onPic &= (model->params->data.F32[2] >= source->pixels->col0);
        onPic &= (model->params->data.F32[2] <  source->pixels->col0 + source->pixels->numCols);
        onPic &= (model->params->data.F32[3] >= source->pixels->row0);
        onPic &= (model->params->data.F32[3] <  source->pixels->row0 + source->pixels->numRows);
        if (!onPic) {
            model->status = PM_MODEL_OFFIMAGE;
        }
    }

    source->mode |= PM_SOURCE_MODE_FITTED;

    psFree(x);
    psFree(y);
    psFree(yErr);
    psFree(myMin);
    psFree(covar);
    psFree(constrain->paramMask);
    psFree(constrain->paramMin);
    psFree(constrain->paramMax);
    psFree(constrain->paramDelta);
    psFree(constrain);
    psFree(params);
    psFree(dparams);

    // free static memory used by pmModelFitSet
    pmModelFitSetClear ();

    rc = (onPic && fitStatus);
    psTrace (__func__, 5, "onPic: %d, fitStatus: %d, nIter: %d, chisq: %f, nDof: %d\n", onPic, fitStatus, model->nIter, model->chisq, model->nDOF);
    psTrace(__func__, 3, "---- %s end : status %d ----\n", __func__, rc);
    return(rc);
}

