/*****************************************************************************
    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 40
#define DATA_WIDTH 1
#define NUM_PARAMS 3

// The function is f(x) = (A * e^(lambda x)) + b

float myFunc(const psVector *restrict myData,
             const psVector *restrict myParams)
{
    float x = myData->data.F32[0];
    float A = myParams->data.F32[0];
    float lambda = myParams->data.F32[1];
    float b = myParams->data.F32[2];
    float tmp = A * exp(-lambda * x) + b;

    //    printf("--------- myFunc((%.1f) %.1f %.1f %.1f) is %.1f ---------\n",
    //           x, A, lambda, b, tmp);

    return(tmp);
}

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

    if (whichParamDeriv == 0) {
        tmp = exp(-lambda * x);
    } else
        if (whichParamDeriv == 1) {
            tmp = -x * A * exp(-lambda * x);
        } else
            if (whichParamDeriv == 2) {
                tmp = 1.0;
            }

    //    printf("--------- myFuncDeriv((%.1f) %.1f %.1f %.1f, %d) is %.1f ---------\n",
    //           x, A, lambda, myParams->data.F32[2], 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;
    psVector *theParams = NULL;
    psS32 currentId = psMemGetId();
    psS32 testStatus = true;
    psS32 memLeaks = 0;

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

    // Build the data.
    initialGuess->data.F32[0] = 1.0;
    initialGuess->data.F32[1] = 0.0;
    initialGuess->data.F32[2] = 0.0;

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

    printf("========== Setting Data Values ==========\n");
    data->data.F32[0] = 6.013392;
    data->data.F32[1] = 5.515377;
    data->data.F32[2] = 5.261095;
    data->data.F32[3] = 4.777455;
    data->data.F32[4] = 4.451353;
    data->data.F32[5] = 3.904903;
    data->data.F32[6] = 3.504387;
    data->data.F32[7] = 3.414999;
    data->data.F32[8] = 3.242736;
    data->data.F32[9] = 3.122204;
    data->data.F32[10] = 2.837632;
    data->data.F32[11] = 2.534700;
    data->data.F32[12] = 2.439173;
    data->data.F32[13] = 2.380830;
    data->data.F32[14] = 2.316090;
    data->data.F32[15] = 2.060826;
    data->data.F32[16] = 1.945679;
    data->data.F32[17] = 1.914126;
    data->data.F32[18] = 1.759509;
    data->data.F32[19] = 1.665067;
    data->data.F32[20] = 1.737926;
    data->data.F32[21] = 1.575523;
    data->data.F32[22] = 1.525066;
    data->data.F32[23] = 1.409607;
    data->data.F32[24] = 1.395209;
    data->data.F32[25] = 1.416887;
    data->data.F32[26] = 1.376042;
    data->data.F32[27] = 1.260955;
    data->data.F32[28] = 1.289627;
    data->data.F32[29] = 1.422672;
    data->data.F32[30] = 1.228287;
    data->data.F32[31] = 1.199176;
    data->data.F32[32] = 1.189989;
    data->data.F32[33] = 0.930076;
    data->data.F32[34] = 1.224608;
    data->data.F32[35] = 1.147376;
    data->data.F32[36] = 1.114001;
    data->data.F32[37] = 1.195118;
    data->data.F32[38] = 1.269581;
    data->data.F32[39] = 1.061982;

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