/******************************************************************************
 * this file defines the RGAUSS source shape model (XXX need a better name!).  Note that these
 * model functions are loaded by pmModelGroup.c using 'include', and thus need no 'include'
 * statements of their own.  The models use a psVector to represent the set of parameters, with
 * the sequence used to specify the meaning of the parameter.  The meaning of the parameters
 * may thus vary depending on the specifics of the model.  All models which are used a PSF
 * representations share a few parameters, for which # define names are listed in pmModel.h:
 
   power-law with fitted slope
   1 / (1 + z + z^alpha)
 
 * PM_PAR_SKY 0   - local sky : note that this is unused and may be dropped in the future
 * PM_PAR_I0 1    - central intensity
 * PM_PAR_XPOS 2  - X center of object
 * PM_PAR_YPOS 3  - Y center of object
 * PM_PAR_SXX 4   - X^2 term of elliptical contour (sqrt(2) / SigmaX)
 * PM_PAR_SYY 5   - Y^2 term of elliptical contour (sqrt(2) / SigmaY)
 * PM_PAR_SXY 6   - X*Y term of elliptical contour
 * PM_PAR_7   7   - power-law slope (alpha)
 *****************************************************************************/

# define PM_MODEL_FUNC       pmModelFunc_RGAUSS
# define PM_MODEL_FLUX       pmModelFlux_RGAUSS
# define PM_MODEL_GUESS      pmModelGuess_RGAUSS
# define PM_MODEL_LIMITS     pmModelLimits_RGAUSS
# define PM_MODEL_RADIUS     pmModelRadius_RGAUSS
# define PM_MODEL_FROM_PSF   pmModelFromPSF_RGAUSS
# define PM_MODEL_FIT_STATUS pmModelFitStatus_RGAUSS

psF64 PS_MODEL_FUNC (psVector *deriv,
                     const psVector *params,
                     const psVector *pixcoord)
{
    psF32 *PAR = params->data.F32;

    psF32 X  = pixcoord->data.F32[0] - PAR[PM_PAR_XPOS];
    psF32 Y  = pixcoord->data.F32[1] - PAR[PM_PAR_YPOS];
    psF32 px = X / PAR[PM_PAR_SXX];
    psF32 py = Y / PAR[PM_PAR_SYY];
    psF32 z  = PS_SQR(px) + PS_SQR(py) + X*Y*PAR[PM_PAR_SXY];

    psF32 p  = pow(z, PAR[PM_PAR_7] - 1.0);
    psF32 r  = 1.0 / (1 + z + z*p);
    psF32 f  = PAR[PM_PAR_I0]*r + PAR[PM_PAR_SKY];

    if (deriv != NULL) {
        psF32 *dPAR = deriv->data.F32;

        // note difference from a pure gaussian: q = params->data.F32[PM_PAR_I0]*r
        psF32 t = PAR[PM_PAR_I0]*r*r;
        psF32 q = t*(1 + PAR[PM_PAR_7]*p);

        dPAR[PM_PAR_SKY] = +1.0;
        dPAR[PM_PAR_I0] = +r;
        dPAR[PM_PAR_XPOS] = q*(2.0*px/PAR[PM_PAR_SXX] + Y*PAR[PM_PAR_SXY]);
        dPAR[PM_PAR_YPOS] = q*(2.0*py/PAR[PM_PAR_SYY] + X*PAR[PM_PAR_SXY]);
        dPAR[PM_PAR_SXX] = +2.0*q*px*px/PAR[PM_PAR_SXX];
        dPAR[PM_PAR_SYY] = +2.0*q*py*py/PAR[PM_PAR_SYY];
        dPAR[PM_PAR_SXY] = -q*X*Y;
        dPAR[PM_PAR_7] = -5.0*t*log(z)*p*z;
    }
    return(f);
}

bool PM_MODEL_LIMITS  (psVector **beta_lim, psVector **params_min, psVector **params_max)
{

    *beta_lim   = psVectorAlloc (8, PS_TYPE_F32);
    *params_min = psVectorAlloc (8, PS_TYPE_F32);
    *params_max = psVectorAlloc (8, PS_TYPE_F32);

    beta_lim[0][0].data.F32[PM_PAR_SKY] = 1000;
    beta_lim[0][0].data.F32[PM_PAR_I0] = 3e6;
    beta_lim[0][0].data.F32[PM_PAR_XPOS] = 5;
    beta_lim[0][0].data.F32[PM_PAR_YPOS] = 5;
    beta_lim[0][0].data.F32[PM_PAR_SXX] = 0.5;
    beta_lim[0][0].data.F32[PM_PAR_SYY] = 0.5;
    beta_lim[0][0].data.F32[PM_PAR_SXY] = 0.5;
    beta_lim[0][0].data.F32[PM_PAR_7] = 0.5;

    params_min[0][0].data.F32[PM_PAR_SKY] = -1000;
    params_min[0][0].data.F32[PM_PAR_I0] = 0;
    params_min[0][0].data.F32[PM_PAR_XPOS] = -100;
    params_min[0][0].data.F32[PM_PAR_YPOS] = -100;
    params_min[0][0].data.F32[PM_PAR_SXX] = 0.5;
    params_min[0][0].data.F32[PM_PAR_SYY] = 0.5;
    params_min[0][0].data.F32[PM_PAR_SXY] = -5.0;
    params_min[0][0].data.F32[PM_PAR_7] = 1.25;

    params_max[0][0].data.F32[PM_PAR_SKY] = 1e5;
    params_max[0][0].data.F32[PM_PAR_I0] = 1e8;
    params_max[0][0].data.F32[PM_PAR_XPOS] = 1e4;  // this should be set by image dimensions!
    params_max[0][0].data.F32[PM_PAR_YPOS] = 1e4;  // this should be set by image dimensions!
    params_max[0][0].data.F32[PM_PAR_SXX] = 100.0;
    params_max[0][0].data.F32[PM_PAR_SYY] = 100.0;
    params_max[0][0].data.F32[PM_PAR_SXY] = +5.0;
    params_max[0][0].data.F32[PM_PAR_7] = 4.0;

    return (TRUE);
}

bool PS_MODEL_GUESS  (psModel *model, psSource *source)
{
    pmMoments *moments = source->moments;
    pmPeak    *peak    = source->peak;
    psF32     *PAR     = model->params->data.F32;

    psEllipseAxes axes;
    psEllipseShape shape;
    psEllipseMoments tmpMoments;

    // XXX fix this stuff : should be using correct ellipse relationships...
    tmpMoments.x2 = PS_SQR(source->moments->Sx);
    tmpMoments.y2 = PS_SQR(source->moments->Sy);
    tmpMoments.xy = source->moments->Sxy;

    axes = psEllipseMomentsToAxes(tmpMoments);
    shape = psEllipseAxesToShape(axes);

    PAR[PM_PAR_SKY] = moments->Sky;
    PAR[PM_PAR_I0] = moments->Peak - moments->Sky;
    PAR[PM_PAR_XPOS] = peak->x;
    PAR[PM_PAR_YPOS] = peak->y;
    PAR[PM_PAR_SXX]  = PS_MAX(0.5, moments->Sx);
    PAR[PM_PAR_SYY]  = PS_MAX(0.5, moments->Sy);
    PAR[PM_PAR_SXY]  = shape.sxy;
    PAR[PM_PAR_7]    = 2.0;
    return(true);
}

psF64 PS_MODEL_FLUX (const psVector *params)
{
    float norm, z;
    psEllipseShape shape;

    psF32 *PAR = params->data.F32;

    shape.sx  = PAR[PM_PAR_SXX] / sqrt(2.0);
    shape.sy  = PAR[PM_PAR_SYY] / sqrt(2.0);
    shape.sxy = PAR[PM_PAR_SXY];

    // Area is equivalent to 2 pi sigma^2
    psEllipseAxes axes = psEllipseShapeToAxes (shape);
    psF64 Area = 2.0 * M_PI * axes.major * axes.minor;

    // the area needs to be multiplied by the integral of f(z)
    norm = 0.0;

    # define DZ 0.25

    float f0 = 1.0;
    float f1, f2;
    for (z = DZ; z < 50; z += DZ) {
        f1 = 1.0 / (1 + z + pow(z, PAR[PM_PAR_7]));
        z += DZ;
        f2 = 1.0 / (1 + z + pow(z, PAR[PM_PAR_7]));
        norm += f0 + 4*f1 + f2;
        f0 = f2;
    }
    norm *= DZ / 3.0;

    psF64 Flux = PAR[PM_PAR_I0] * Area * norm;

    return(Flux);
}

// define this function so it never returns Inf or NaN
// return the radius which yields the requested flux
psF64 PS_MODEL_RADIUS (const psVector *params, psF64 flux)
{
    psF64 z, f, p;
    psEllipseShape shape;

    psF32 *PAR = params->data.F32;

    if (flux <= 0)
        return (1.0);
    if (PAR[PM_PAR_I0] <= 0)
        return (1.0);
    if (flux >= PAR[PM_PAR_I0])
        return (1.0);

    shape.sx  = PAR[PM_PAR_SXX] / sqrt(2.0);
    shape.sy  = PAR[PM_PAR_SYY] / sqrt(2.0);
    shape.sxy = PAR[PM_PAR_SXY];

    psEllipseAxes axes = psEllipseShapeToAxes (shape);
    psF64 sigma = axes.major;

    psF64 dz = 1.0 / (2.0 * sigma*sigma);
    psF64 limit = flux / PAR[PM_PAR_I0];

    // use the fact that f is monotonically decreasing
    z = 0;
    Nstep = 0;

    // choose a z value guaranteed to be beyond our limit
    float z0 = pow((1.0 / limit), (1.0 / 2.25));
    float z1 = (1.0 / limit) / PAR[PM_PAR_7];
    z1 = PS_MAX (z0, z1);
    z0 = 0.0;

    // perform a type of bisection to find the value
    float f0 = 1.0 / (1 + z0 + pow(z0, PAR[PM_PAR_7]));
    float f1 = 1.0 / (1 + z1 + pow(z1, PAR[PM_PAR_7]));
    while ((Nstep < 10) && (fabs(z1 - z0) > 0.5)) {
        z = 0.5*(z0 + z1);
        f = 1.0 / (1 + z + pow(z, PAR[PM_PAR_7]));
        if (f > limit) {
            z0 = z;
            f0 = f;
        } else {
            z1 = z;
            f1 = f;
        }
        Nstep ++;
    }
    psF64 radius = sigma * sqrt (2.0 * z);

    if (isnan(radius))
        psAbort("error in code: radius is NaN");

    return (radius);
}

bool PS_MODEL_FROM_PSF (psModel *modelPSF, psModel *modelFLT, pmPSF *psf)
{

    psF32 *out = modelPSF->params->data.F32;
    psF32 *in  = modelFLT->params->data.F32;

    // we require these two parameters to exist
    assert (psf->params_NEW->n > PM_PAR_YPOS);
    assert (psf->params_NEW->n > PM_PAR_XPOS);

    for (int i = 0; i < psf->params_NEW->n; i++) {
        if (psf->params_NEW->data[i] == NULL) {
            out[i] = in[i];
        } else {
            psPolynomial2D *poly = psf->params_NEW->data[i];
            out[i] = psPolynomial2DEval(poly, in[PM_PAR_XPOS], in[PM_PAR_YPOS]);
        }
    }

    // the 2D model for SXY actually fits SXY / (SXX^-2 + SYY^-2); correct here
    out[PM_PAR_SXY] = pmPSF_SXYtoModel (out);

    return(true);
}

bool PM_MODEL_FIT_STATUS (pmModel *model)
{

    psF32 dP;
    bool  status;

    psF32 *PAR  = model->params->data.F32;
    psF32 *dPAR = model->dparams->data.F32;

    dP = 0;
    dP += PS_SQR(dPAR[PM_PAR_SXX] / PAR[PM_PAR_SXX]);
    dP += PS_SQR(dPAR[PM_PAR_SYY] / PAR[PM_PAR_SYY]);
    dP = sqrt (dP);

    status = true;
    status &= (dP < 0.5);
    status &= (PAR[PM_PAR_I0] > 0);
    status &= ((dPAR[PM_PAR_I0]/PAR[PM_PAR_I0]) < 0.5);

    return status;
}

# undef PM_MODEL_FUNC
# undef PM_MODEL_FLUX
# undef PM_MODEL_GUESS
# undef PM_MODEL_LIMITS
# undef PM_MODEL_RADIUS
# undef PM_MODEL_FROM_PSF
# undef PM_MODEL_FIT_STATUS
