/** @file  pmModel.c
 *
 *  Functions to define and manipulate object models
 *
 *  @author GLG, MHPCC
 *  @author EAM, IfA
 *
 *  @version $Revision: 1.22 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2008-01-02 20:36:36 $
 *
 *  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 "pmHDU.h"
#include "pmFPA.h"
#include "pmResiduals.h"
#include "pmGrowthCurve.h"
#include "pmTrend2D.h"
#include "pmPSF.h"
#include "pmModel.h"
#include "pmModelClass.h"

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.
*****************************************************************************/
pmModel *pmModelAlloc(pmModelType type)
{
    psTrace("psModules.objects", 3, "---- %s() begin ----\n", __func__);

    pmModelClass *class = pmModelClassSelect (type);
    if (class == NULL) {
        psError(PS_ERR_UNKNOWN, true, "Undefined pmModelType");
        return(NULL);
    }

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

    tmp->type = type;
    tmp->chisq = 0.0;
    tmp->chisqNorm = 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 = pmModelClassParameterCount(type);
    assert (Nparams);

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

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

    tmp->modelFunc          = class->modelFunc;
    tmp->modelFlux          = class->modelFlux;
    tmp->modelRadius        = class->modelRadius;
    tmp->modelLimits        = class->modelLimits;
    tmp->modelGuess         = class->modelGuess;
    tmp->modelFromPSF       = class->modelFromPSF;
    tmp->modelParamsFromPSF = class->modelParamsFromPSF;
    tmp->modelFitStatus     = class->modelFitStatus;

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

bool psMemCheckModel(psPtr ptr)
{
    PS_ASSERT_PTR(ptr, false);
    return ( psMemGetDeallocator(ptr) == (psFreeFunc) modelFree);
}

// copy model to a new structure
pmModel *pmModelCopy (pmModel *model)
{
    PS_ASSERT_PTR_NON_NULL(model, NULL);

    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;

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

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

    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, NAN);
    PS_ASSERT_PTR_NON_NULL(model, NAN);
    PS_ASSERT_PTR_NON_NULL(model->params, NAN);

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

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

psF32 pmModelEvalWithOffset(pmModel *model, psImage *image, psS32 col, psS32 row, int dx, int dy)
{
    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 + dx);
    x->data.F32[1] = (psF32) (row + image->row0 + dy);
    psF32 tmpF;

    tmpF = model->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,
                          int dx,
                          int dy
    )
{
    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;

    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;

    psImage *myRo = NULL;
    psImage *myRx = NULL;
    psImage *myRy = NULL;
    psImageInterpolateOptions *Ro = NULL;
    psImageInterpolateOptions *Rx = NULL;
    psImageInterpolateOptions *Ry = NULL;
    if (model->residuals && (mode & (PM_MODEL_OP_RES0 | PM_MODEL_OP_RES1))) {
        // if the residual image and object image don't match,
        // supply an appropriately overlapped residual image
        psImage *inRo = model->residuals->Ro;
        psImage *inRx = model->residuals->Rx;
        psImage *inRy = model->residuals->Ry;
        xBin = model->residuals->xBin;
        yBin = model->residuals->yBin;
        xResidCenter = model->residuals->xCenter;
        yResidCenter = model->residuals->yCenter;
        if ((image->numCols != inRo->numCols) ||
            (image->numRows != inRo->numRows)) {
            myRo = psImageAlloc (image->numCols, image->numRows, PS_TYPE_F32);
            myRx = psImageAlloc (image->numCols, image->numRows, PS_TYPE_F32);
            myRy = psImageAlloc (image->numCols, image->numRows, PS_TYPE_F32);
            // Difference between input and desired centres
            int xDiff = (int)(inRo->numCols / 2) - (xPos - image->col0);
            int yDiff = (int)(inRo->numRows / 2) - (yPos - image->row0);
            xResidCenter -= xDiff;
            yResidCenter -= yDiff;
            for (int iy = 0; iy < myRo->numRows; iy++) {
                int jy = iy + yDiff;
                if ((jy < 0) || (jy >= inRo->numRows)) {
                    for (int ix = 0; ix < myRo->numCols; ix++) {
                        myRo->data.F32[iy][ix] = 0.0;
                        myRx->data.F32[iy][ix] = 0.0;
                        myRy->data.F32[iy][ix] = 0.0;
                    }
                    continue;
                }
                for (int ix = 0; ix < myRo->numCols; ix++) {
                    int jx = ix + xDiff;
                    if ((jx < 0) || (jx >= inRo->numCols)) {
                        myRo->data.F32[iy][ix] = 0.0;
                        myRx->data.F32[iy][ix] = 0.0;
                        myRy->data.F32[iy][ix] = 0.0;
                    } else {
                        myRo->data.F32[iy][ix] = inRo->data.F32[jy][jx];
                        myRx->data.F32[iy][ix] = inRx->data.F32[jy][jx];
                        myRy->data.F32[iy][ix] = inRy->data.F32[jy][jx];
                    }
                }
            }
        } else {
            myRo = psMemIncrRefCounter (inRo);
            myRx = psMemIncrRefCounter (inRx);
            myRy = psMemIncrRefCounter (inRy);
        }

        Ro = psImageInterpolateOptionsAlloc(PS_INTERPOLATE_BILINEAR, myRo, NULL, mask, 0, 0.0, 0.0, 1, 0, 0.0);
        Rx = psImageInterpolateOptionsAlloc(PS_INTERPOLATE_BILINEAR, myRx, NULL, NULL, 0, 0.0, 0.0, 1, 0, 0.0);
        Ry = psImageInterpolateOptionsAlloc(PS_INTERPOLATE_BILINEAR, myRy, NULL, NULL, 0, 0.0, 0.0, 1, 0, 0.0);

    }

    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 + dx;
            imageRow = iy + image->row0 + dy;

            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 += model->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;

                psU8 mflux = 0;
                if (mode & PM_MODEL_OP_RES0) {
                    double Fo = 0.0;
                    psImageInterpolate (&Fo, NULL, &mflux, ox, oy, Ro);
                    if (!mflux && isfinite(Fo)) {
                        pixelValue += Io*Fo;
                    }
                }
                // skip Rx,Ry if Ro is masked
                if (!mflux && (mode & PM_MODEL_OP_RES1)) {
                    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);
    psFree(myRo);
    psFree(myRx);
    psFree(myRy);
    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, 0.0, 0.0);
    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, 0.0, 0.0);
    psTrace("psModules.objects", 3, "---- %s(%d) end ----\n", __func__, rc);
    return(rc);
}

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

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