/** @file  pmModel.c
 *
 *  Functions to define and manipulate object models
 *
 *  @author GLG, MHPCC
 *  @author EAM, IfA
 *
 *  @version $Revision: 1.13 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2007-06-20 02:22:26 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <math.h>
#include <string.h>
#include <pslib.h>
#include "pmResiduals.h"
#include "pmModel.h"

/* The following types and functions need to be known by pmModel.c and are defined in
   pmModelGroup.h.  The use of pmSource and other types in pmModelGroup.h prevent us from
   directly including it here  ***/

typedef psMinimizeLMChi2Func pmModelFunc;
typedef bool (*pmModelFitStatusFunc)(pmModel *model);
pmModelFunc pmModelFunc_GetFunction (pmModelType type);
pmModelFitStatusFunc pmModelFitStatusFunc_GetFunction (pmModelType type);
int pmModelParameterCount(pmModelType type);

static void modelFree(pmModel *tmp)
{
    psTrace("psModules.objects", 4, "---- %s() begin ----\n", __func__);
    psFree(tmp->params);
    psFree(tmp->dparams);
    psTrace("psModules.objects", 4, "---- %s() end ----\n", __func__);
}

/******************************************************************************
pmModelAlloc(): Allocate the pmModel structure, along with its parameters,
and initialize the type member.  Initialize the params to 0.0.
XXX EAM: simplifying code with pmModelParameterCount
*****************************************************************************/
pmModel *pmModelAlloc(pmModelType type)
{
    psTrace("psModules.objects", 3, "---- %s() begin ----\n", __func__);

    pmModel *tmp = (pmModel *) psAlloc(sizeof(pmModel));
    psMemSetDeallocator(tmp, (psFreeFunc) modelFree);

    tmp->type = type;
    tmp->chisq = 0.0;
    tmp->nDOF  = 0;
    tmp->nIter = 0;
    tmp->radiusFit = 0;
    tmp->flags = PM_MODEL_STATUS_NONE;
    tmp->residuals = NULL;              // XXX should the model own this memory?

    psS32 Nparams = pmModelParameterCount(type);
    if (Nparams == 0) {
        psError(PS_ERR_UNKNOWN, true, "Undefined pmModelType");
        return(NULL);
    }

    tmp->params  = psVectorAlloc(Nparams, PS_TYPE_F32);
    tmp->dparams = psVectorAlloc(Nparams, PS_TYPE_F32);

    for (psS32 i = 0; i < tmp->params->n; i++) {
        tmp->params->data.F32[i] = 0.0;
        tmp->dparams->data.F32[i] = 0.0;
    }

    psTrace("psModules.objects", 3, "---- %s() end ----\n", __func__);
    return(tmp);
}

// copy model to a new structure
pmModel *pmModelCopy (pmModel *model)
{

    pmModel *new = pmModelAlloc (model->type);

    new->chisq     = model->chisq;
    new->nDOF      = model->nDOF;
    new->nIter     = model->nIter;
    new->flags     = model->flags;
    new->radiusFit = model->radiusFit;

    // XXX note that model->residuals is just a reference
    new->residuals = model->residuals;

    for (int i = 0; i < new->params->n; i++) {
        new->params->data.F32[i]  = model->params->data.F32[i];
        new->dparams->data.F32[i] = model->dparams->data.F32[i];
    }

    return (new);
}

/******************************************************************************
    pmModelEval(source, level, row): evaluates the model function at the specified coords.

    NOTE: The coords are in subImage source->pixel coords, not image coords.

    XXX: Use static vectors for x (NO: needs to be thread safe)
*****************************************************************************/
psF32 pmModelEval(pmModel *model, psImage *image, psS32 col, psS32 row)
{
    psTrace("psModules.objects", 3, "---- %s() begin ----\n", __func__);
    PS_ASSERT_PTR_NON_NULL(image, false);
    PS_ASSERT_PTR_NON_NULL(model, false);
    PS_ASSERT_PTR_NON_NULL(model->params, false);

    // Allocate the x coordinate structure and convert row/col to image space.
    //
    psVector *x = psVectorAlloc(2, PS_TYPE_F32);
    x->data.F32[0] = (psF32) (col + image->col0);
    x->data.F32[1] = (psF32) (row + image->row0);
    psF32 tmpF;
    pmModelFunc modelFunc;

    modelFunc = pmModelFunc_GetFunction (model->type);
    tmpF = modelFunc (NULL, model->params, x);
    psFree(x);
    psTrace("psModules.objects", 3, "---- %s() end ----\n", __func__);
    return(tmpF);
}

static bool AddOrSubModel(psImage *image,
                          psImage *mask,
                          pmModel *model,
                          pmModelOpMode mode,
                          bool add,
                          psMaskType maskVal
    )
{
    psTrace("psModules.objects", 3, "---- %s() begin ----\n", __func__);

    PS_ASSERT_PTR_NON_NULL(model, false);
    PS_ASSERT_IMAGE_NON_NULL(image, false);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, false);

    psVector *x = psVectorAlloc(2, PS_TYPE_F32);
    psVector *params = model->params;
    pmModelFunc modelFunc = pmModelFunc_GetFunction (model->type);
    psS32 imageCol;
    psS32 imageRow;
    psF32 pixelValue;

    // save original values; restore before returning
    // use the true source position for the residual model
    // the PSF model has presumably already been set for this coordinate
    float xPos    = params->data.F32[PM_PAR_XPOS];
    float yPos    = params->data.F32[PM_PAR_YPOS];
    float IoSave  = params->data.F32[PM_PAR_I0];
    float skySave = params->data.F32[PM_PAR_SKY];

    if (mode & PM_MODEL_OP_NORM) {
        params->data.F32[PM_PAR_I0] = 1.0;
    }
    if (!(mode & PM_MODEL_OP_SKY)) {
        params->data.F32[PM_PAR_SKY] = 0.0;
    }
    if (mode & PM_MODEL_OP_CENTER) {
        params->data.F32[PM_PAR_XPOS] = image->col0 + 0.5*image->numCols;
        params->data.F32[PM_PAR_YPOS] = image->row0 + 0.5*image->numRows;
    }

    // use these values for this realization
    float xCenter  = params->data.F32[PM_PAR_XPOS];
    float yCenter  = params->data.F32[PM_PAR_YPOS];
    float Io       = params->data.F32[PM_PAR_I0];

    int xBin = 1;
    int yBin = 1;
    float xResidCenter = 0.0;
    float yResidCenter = 0.0;

    psImageInterpolateOptions *Ro = NULL;
    psImageInterpolateOptions *Rx = NULL;
    psImageInterpolateOptions *Ry = NULL;
    if (model->residuals && (mode & (PM_MODEL_OP_RES0 | PM_MODEL_OP_RES1))) {
        Ro = psImageInterpolateOptionsAlloc(
            PS_INTERPOLATE_BILINEAR,
            model->residuals->Ro, NULL, NULL, 0, 0.0, 0.0, 1, 0, 0.0);
        Rx = psImageInterpolateOptionsAlloc(
            PS_INTERPOLATE_BILINEAR,
            model->residuals->Rx, NULL, NULL, 0, 0.0, 0.0, 1, 0, 0.0);
        Ry = psImageInterpolateOptionsAlloc(
            PS_INTERPOLATE_BILINEAR,
            model->residuals->Ry, NULL, NULL, 0, 0.0, 0.0, 1, 0, 0.0);

        xBin = model->residuals->xBin;
        yBin = model->residuals->yBin;
        xResidCenter = model->residuals->xCenter;
        yResidCenter = model->residuals->yCenter;
    }

    for (psS32 iy = 0; iy < image->numRows; iy++) {
        for (psS32 ix = 0; ix < image->numCols; ix++) {
            if ((mask != NULL) && (mask->data.U8[iy][ix] & maskVal))
                continue;

            // Convert i/j to image coord space:
            // XXX should we use using 0.5 pixel offset?
            imageCol = ix + image->col0;
            imageRow = iy + image->row0;

            x->data.F32[0] = (float) imageCol;
            x->data.F32[1] = (float) imageRow;

            pixelValue = 0.0;

            // add in the desired components for this coordinate
            if (mode & PM_MODEL_OP_FUNC) {
                pixelValue += modelFunc (NULL, params, x);
            }

            // get the contribution from the residual model
            if (Ro) {
                // fractional image position
                float ox = xBin*(imageCol + 0.5 - xCenter) + xResidCenter;
                float oy = yBin*(imageRow + 0.5 - yCenter) + yResidCenter;

                if (mode & PM_MODEL_OP_RES0) {
                    psU8 mflux = 0;
                    double Fo = 0.0;
                    psImageInterpolate (&Fo, NULL, &mflux, ox, oy, Ro);
                    if (!mflux && isfinite(Fo)) {
                        pixelValue += Io*Fo;
                    }
                }
                if (mode & PM_MODEL_OP_RES1) {
                    psU8 mflux = 0;
                    double Fx = 0.0;
                    double Fy = 0.0;
                    psImageInterpolate (&Fx, NULL, &mflux, ox, oy, Rx);
                    psImageInterpolate (&Fy, NULL, &mflux, ox, oy, Ry);
                    if (!mflux && isfinite(Fx) && isfinite(Fy)) {
                        pixelValue += Io*(xPos*Fx + yPos*Fy);
                    }
                }
            }

            // add or subtract the value
            if (add) {
                image->data.F32[iy][ix] += pixelValue;
            } else {
                image->data.F32[iy][ix] -= pixelValue;
            }
        }
    }

    // restore original values
    params->data.F32[PM_PAR_I0]   = IoSave;
    params->data.F32[PM_PAR_SKY]  = skySave;
    params->data.F32[PM_PAR_XPOS] = xPos;
    params->data.F32[PM_PAR_YPOS] = yPos;

    psFree(x);
    psFree(Ro);
    psFree(Rx);
    psFree(Ry);
    psTrace("psModules.objects", 3, "---- %s(true) end ----\n", __func__);
    return(true);
}

/******************************************************************************
 *****************************************************************************/
bool pmModelAdd(psImage *image,
                psImage *mask,
                pmModel *model,
                pmModelOpMode mode,
                psMaskType maskVal)
{
    psTrace("psModules.objects", 3, "---- %s() begin ----\n", __func__);
    psBool rc = AddOrSubModel(image, mask, model, mode, true, maskVal);
    psTrace("psModules.objects", 3, "---- %s(%d) end ----\n", __func__, rc);
    return(rc);
}

/******************************************************************************
 *****************************************************************************/
bool pmModelSub(psImage *image,
                psImage *mask,
                pmModel *model,
                pmModelOpMode mode,
                psMaskType maskVal)
{
    psTrace("psModules.objects", 3, "---- %s() begin ----\n", __func__);
    psBool rc = AddOrSubModel(image, mask, model, mode, false, maskVal);
    psTrace("psModules.objects", 3, "---- %s(%d) end ----\n", __func__, rc);
    return(rc);
}

// check for success of model fit
bool pmModelFitStatus (pmModel *model)
{

    bool status;

    pmModelFitStatusFunc statusFunc = pmModelFitStatusFunc_GetFunction (model->type);
    status = statusFunc (model);

    return (status);
}

