/*****************************************************************************
    This routine must ensure that the psMinimizeChi2 function correctly
    minimizes an arbitrary function.
 *****************************************************************************/
#include <stdio.h>
#include "pslib.h"
#include "psTest.h"
#include "psMemory.h"
#include "psVector.h"
#include "psImage.h"
#include "psFunctions.h"
#include "psMinimize.h"
#include <math.h>
#define NUM_DATA 10
#define DATA_WIDTH 2
#define NUM_PARAMS 4

float myFunc(const psVector *restrict myData,
             const psVector *restrict myParams)
{
    float x = myData->data.F32[0];
    float y = myData->data.F32[1];
    float A = myParams->data.F32[0];
    float B = myParams->data.F32[1];
    float C = myParams->data.F32[2];
    float D = myParams->data.F32[3];
    float tmp = 0.0;

    tmp = (A * x) + (B*x*x) + (C*y*y) + (D*x*x*y);
    //    printf("--------- myFunc((%.1f %.1f) %.1f %.1f %.1f %.1f) is %.1f ---------\n",
    //           x, y, A, B, C, D, tmp);

    return(tmp);
}

float myFuncDeriv(const psVector *restrict myData,
                  const psVector *restrict myParams,
                  psS32 whichParamDeriv)
{
    float x = myData->data.F32[0];
    float y = myData->data.F32[1];
    float tmp = 0.0;

    if (whichParamDeriv == 0) {
        tmp = 1.0;
        tmp = x;
    } else
        if (whichParamDeriv == 1) {
            tmp = x;
            tmp = x*x;
        } else
            if (whichParamDeriv == 2) {
                tmp = y;
                tmp = y*y;
            } else
                if (whichParamDeriv == 3) {
                    tmp = x * x * y;
                }

    //    printf("--------- myFuncDeriv((%.1f %.1f) %.1f %.1f %.1f %.1f (%d)) is %.1f ---------\n",
    //            x, y, myParams->data.F32[0], myParams->data.F32[1],
    //            myParams->data.F32[2], myParams->data.F32[3], whichParamDeriv, tmp);
    return(tmp);
}


psS32 main()
{
    psImage *domain = NULL;
    psVector *data = NULL;
    psVector *errors = NULL;
    psVector *initialGuess = NULL;
    psVector *paramMask = NULL;
    psVector *tmpVecPtr = NULL;
    float chiSq = 0.0;
    psS32 i = 0;
    psS32 j = 0;
    psS32 currentId = psMemGetId();
    psS32 testStatus = true;
    psS32 memLeaks = 0;
    psVector *theParams = NULL;

    domain = psImageAlloc(DATA_WIDTH, NUM_DATA, PS_TYPE_F32);
    data = psVectorAlloc(NUM_DATA, PS_TYPE_F32);
    errors = psVectorAlloc(NUM_DATA, PS_TYPE_F32);
    initialGuess = psVectorAlloc(NUM_PARAMS, PS_TYPE_F32);
    paramMask = psVectorAlloc(NUM_PARAMS, PS_TYPE_F32);
    tmpVecPtr = psVectorAlloc(DATA_WIDTH, PS_TYPE_F32);


    //--------- myFunc((0.0 1.0) -8.6 -7.6 14.6 5.0) is 6.0 ---------

    initialGuess->data.F32[0] = -8.6;
    initialGuess->data.F32[1] = -7.6;
    initialGuess->data.F32[2] = 14.6;
    initialGuess->data.F32[3] = 5.0;

    // Build the data.
    initialGuess->data.F32[0] = 2.0;
    initialGuess->data.F32[1] = 3.0;
    initialGuess->data.F32[2] = 4.0;
    initialGuess->data.F32[3] = 5.0;


    for (i = 0; i<NUM_DATA; i++) {
        for (j=0; j<DATA_WIDTH; j++) {
            domain->data.F32[i][j] = (float) (i + j);
        }
    }
    for (i = 0; i<NUM_DATA; i++) {
        errors->data.F32[i] = 0.1;
    }

    //    printf("========== Setting Data Values ==========\n");
    for (i = 0; i<NUM_DATA; i++) {
        for (j=0; j<DATA_WIDTH; j++) {
            tmpVecPtr->data.F32[j] = domain->data.F32[i][j];
        }
        data->data.F32[i] = myFunc(tmpVecPtr, initialGuess);
    }
    //    printf("=========================================\n");

    for (i=0;i<NUM_PARAMS;i++) {
        initialGuess->data.F32[i]+= 1.0;
    }

    printPositiveTestHeader(stdout,
                            "psMinimize functions",
                            "psMinimizeChi2(): no masks");

    theParams = psMinimizeChi2(myFunc,
                               myFuncDeriv,
                               domain,
                               data,
                               errors,
                               initialGuess,
                               NULL,
                               &chiSq);
    for (i=0;i<NUM_PARAMS;i++) {
        printf("theParams[%d] is %.1f\n", i, theParams->data.F32[i]);
    }

    printf("chiSq is %f\n", chiSq);
    psMemCheckCorruption(1);
    printFooter(stdout,
                "psMinimize functions",
                "psMinimizeChi2(): no masks",
                testStatus);


    psMemCheckCorruption(1);
    psFree(domain);
    psFree(data);
    psFree(errors);
    psFree(initialGuess);
    psFree(paramMask);
    psFree(tmpVecPtr);

    psMemCheckCorruption(1);
    memLeaks = psMemCheckLeaks(currentId,NULL,stderr);
    if (0 != memLeaks) {
        psAbort(__func__,"Memory Leaks! (%d leaks)", memLeaks);
    }

    return (!testStatus);
}
