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

#include <stdio.h>
#include <math.h>
#include <string.h>
#include "pslib.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(__func__, 4, "---- %s() begin ----\n", __func__);
    psFree(tmp->params);
    psFree(tmp->dparams);
    psTrace(__func__, 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(__func__, 3, "---- %s() begin ----\n", __func__);
    pmModel *tmp = (pmModel *) psAlloc(sizeof(pmModel));

    tmp->type = type;
    tmp->chisq = 0.0;
    tmp->nIter = 0;
    tmp->radiusTMP = 0;
    tmp->status = PM_MODEL_UNTRIED;

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

    psMemSetDeallocator(tmp, (psFreeFunc) modelFree);
    psTrace(__func__, 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->status = model->status;
    new->radiusTMP = model->radiusTMP;

    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(__func__, 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(__func__, 3, "---- %s() end ----\n", __func__);
    return(tmpF);
}

bool AddOrSubModel(psImage *image,
                   psImage *mask,
                   pmModel *model,
                   bool center,
                   bool sky,
                   bool add
                      )
{
    psTrace(__func__, 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 skyValue = params->data.F32[0];
    psF32 pixelValue;

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

            // Convert i/j to image coord space:
            // 'center' option changes meaning of i,j:
            if (center) {
                imageCol = j - 0.5*image->numCols + model->params->data.F32[2];
                imageRow = i - 0.5*image->numRows + model->params->data.F32[3];
            } else {
                imageCol = j + image->col0;
                imageRow = i + image->row0;
            }

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

            // set the appropriate pixel value for this coordinate
            if (sky) {
                pixelValue = modelFunc (NULL, params, x);
            } else {
                pixelValue = modelFunc (NULL, params, x) - skyValue;
            }


            // add or subtract the value
            if (add
               ) {
                image->data.F32[i][j] += pixelValue;
            }
            else {
                image->data.F32[i][j] -= pixelValue;
            }
        }
    }
    psFree(x);
    psTrace(__func__, 3, "---- %s(true) end ----\n", __func__);
    return(true);
}

/******************************************************************************
 *****************************************************************************/
bool pmModelAdd(psImage *image,
                psImage *mask,
                pmModel *model,
                bool center,
                bool sky)
{
    psTrace(__func__, 3, "---- %s() begin ----\n", __func__);
    psBool rc = AddOrSubModel(image, mask, model, center, sky, true);
    psTrace(__func__, 3, "---- %s(%d) end ----\n", __func__, rc);
    return(rc);
}

/******************************************************************************
 *****************************************************************************/
bool pmModelSub(psImage *image,
                psImage *mask,
                pmModel *model,
                bool center,
                bool sky)
{
    psTrace(__func__, 3, "---- %s() begin ----\n", __func__);
    psBool rc = AddOrSubModel(image, mask, model, center, sky, false);
    psTrace(__func__, 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);
}

