/*
 * measure shutter correction:
 
 * input  : collection of shutter correction exposures (pre-processed)
 * output : a shutter correction image
 
 * the measurement could be performed on any focal-plane unit at a time. for
 * GPC, the obvious scale is to measure the effect on the entire focal plane at
 * once, with a single reference point in the field.  this is a little more
 * complex than just measuring the effect for a single 2D image array.  the
 * reference point and the detailed analysis points need to be defined for the
 * entire hierarchy rather than just as coordinate pairs or regions.  a
 * pmFPAview would be a natural element with which to define these points, but
 * at the moment, the pmFPAview structure defines a band in the CCD, not a
 * coordinate.  An option is to instead specify the reference locations as a
 * pmFPAview coupled with a psRegion, though we need to be careful not to
 * over-specify the pixels (ie, conflict between pmFPAview and psRegion).
 
 * at each point in an image with exposure time T, we measure f(k;T) = F(k;T) /
 * F(0;T) where k is the coordinate of the point of interest, 0 is the reference
 * coordinate, and F(k;T) is the measured number of counts at the point of
 * interest in this image.  given a collection of f(k;T) values, we need to
 * determine the model f(k;T) = A(k) (T + dTk) / (T + dTo) where dTk is the shutter
 * error at the given position, dTo is the shutter error at the reference
 * position, and A(k) is the scaling factor for the given position.
 
 - load the configuration information
 - load the input file information, build pmFPAfile's for all inputs
 
 - determine the reference point and detailed analysis regions from the config information
 - for each image
 -- measure the reference point counts
 - for each analysis region:
 -- measure shutter parameters (dTo, dTk, A):
 --- for each image:
 ---- measure counts at the region
 ---- divide by the reference counts
 --- linear extrapolation to find f(inf) = A(k)
 --- linear extrapolation to find f(0) = A(k) dTk / dTo
 --- linear interpolation to find coordinate where f(dTo) = A (1 + dTk/dTo) / 2
 --- non-linear fit of T, f(T) to f(k;T) = A(k) (T + dTk) / (T + dTo)
 - use the collection of dTo values to choose a best value for dTo (median)
 - for each image pixel
 -- divide by the reference counts
 -- generate the vectors T, f(T)
 -- linear fit of T, f(T) to f(k;T) = A(k) (T + dTk) / (T + dTo) using dTo above
 -- save dTk, A(k) in output image pixels
 -- apply dTk, A(k) to measure residual images
 -- generate residual FITS/JPEG images
 
*/

#if HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <math.h>
#include <strings.h>
#include <pslib.h>

#include "pmShutterCorrection.h"
#include "psVectorBracket.h"

static void pmShutterCorrParsFree (pmShutterCorrPars *pars)
{
    if (pars == NULL)
        return;
    return;
}

pmShutterCorrPars *pmShutterCorrParsAlloc ()
{

    pmShutterCorrPars *pars = (pmShutterCorrPars *) psAlloc(sizeof(pmShutterCorrPars));
    psMemSetDeallocator(pars, (psFreeFunc) pmShutterCorrParsFree);
    pars->scale  = 0.0;
    pars->offset = 0.0;
    pars->offref = 0.0;
    return (pars);
}

// use interpolation to guess shutter correction parameters given a set of exposures times and normalized
// counts (divided by the reference counts for each image)
pmShutterCorrPars *pmShutterCorrectionGuess (psVector *exptime, psVector *counts)
{

    psVector *tmpX = NULL;
    psVector *tmpY = NULL;

    pmShutterCorrPars *pars = pmShutterCorrParsAlloc ();
    psPolynomial1D *line = psPolynomial1DAlloc (PS_POLYNOMIAL_ORD, 1);

    int N = exptime->n;
    // ASSERT (N >> 5)

    // NOTE: vectors must be sorted on input.  It is expensive to sort or check this here, but
    // it is easy to arrange by sorting the images before generating these vectors.

    // choose the highest exptime point as the guess for the scale:
    // XXX we could examine the top 2 or 3 values and decide if we
    // extended exptime enough or median clip.
    pars->scale = counts->data.F32[N-1];

    // fit a line to the lowest three points and extrapolate to 0.0
    tmpX = psVectorAlloc (2, PS_TYPE_F32);
    tmpY = psVectorAlloc (2, PS_TYPE_F32);
    tmpX->n = tmpY->n = 2;

    tmpX->data.F32[0] = exptime->data.F32[0];
    tmpX->data.F32[1] = exptime->data.F32[1];
    tmpY->data.F32[0] = counts->data.F32[0];
    tmpY->data.F32[1] = counts->data.F32[1];

    // fit a line and extrapolate the fit to 0.0
    psVectorFitPolynomial1D (line, NULL, 0, tmpY, NULL, tmpX);
    float ratio = psPolynomial1DEval (line, 0.0) / pars->scale;

    // XXX we need a sanity check:
    // if the mean value of the three points is higher than pars->scale,
    // then the slope should be negative.
    // if the mean value of the three points is lower than pars->scale,
    // then the slope should be positive.

    // find two points bracketing the value counts = A (1 + dTk/dTo) / 2 = pars->scale (1 + ratio) / 2
    float value = pars->scale * (1 + ratio) / 2.0;

    int Np;
    if (ratio < 1.0) {
        Np = psVectorBracket (counts, value, true);
    } else {
        Np = psVectorBracketDescend (counts, value, true);
    }
    int Nm = (Np == 0) ? 1 : Np - 1;

    tmpX->data.F32[0] = counts->data.F32[Nm];
    tmpX->data.F32[1] = counts->data.F32[Np];
    tmpY->data.F32[0] = exptime->data.F32[Nm];
    tmpY->data.F32[1] = exptime->data.F32[Np];

    // fit a line and extrapolate the fit to counts = A (1 + dTk/dTo) : exptime = dTo
    line = psVectorFitPolynomial1D (line, NULL, 0, tmpY, NULL, tmpX);
    pars->offref = psPolynomial1DEval (line, value);
    pars->offset = ratio * pars->offref;

    psFree (line);
    psFree (tmpX);
    psFree (tmpY);

    return (pars);
}

// linear fit to the counts and exptime, given a value for offref
pmShutterCorrPars *pmShutterCorrectionLinFit (psVector *exptime, psVector *counts, psVector *cntError, float offref)
{

    // this step is identical for all pixels: do it once and save?
    psVector *x = psVectorAlloc (exptime->n, PS_TYPE_F32);
    psVector *y = psVectorAlloc (exptime->n, PS_TYPE_F32);
    x->n = y->n = exptime->n;

    for (int i = 0; i < exptime->n; i++) {
        float value = 1.0 / (exptime->data.F32[i] + offref);
        x->data.F32[i] = exptime->data.F32[i] * value;
        y->data.F32[i] = value;
    }

    psPolynomial2D *line = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, 1, 1);

    // mask out the terms we will not fit
    line->mask[0][0] = 1;
    line->mask[1][1] = 1;
    line->coeff[0][0] = 0;
    line->coeff[1][1] = 0;

    psVectorFitPolynomial2D (line, NULL, 0, counts, cntError, x, y);

    pmShutterCorrPars *pars = pmShutterCorrParsAlloc ();

    pars->offref = offref;
    pars->scale  = line->coeff[1][0];
    pars->offset = line->coeff[0][1] / line->coeff[1][0];

    psFree (x);
    psFree (y);
    psFree (line);

    return (pars);
}

static psF32 pmShutterCorrectionModel(psVector *deriv,
                                      const psVector *params,
                                      const psVector *x)
{
    psF32 A  = params->data.F32[0];
    psF32 p  = x->data.F32[0] + params->data.F32[1];
    psF32 q  = 1.0 / (x->data.F32[0] + params->data.F32[2]);
    psF32 f  = A*p*q;

    if (deriv != NULL) {
        deriv->data.F32[0] = p*q;
        deriv->data.F32[1] = A*q;
        deriv->data.F32[2] = -A*p*q*q;
    }
    return(f);
}

// non-linear fit to the counts and exptime, given a guess for the three parameters
pmShutterCorrPars *pmShutterCorrectionFullFit (psVector *exptime, psVector *counts, psVector *cntError, pmShutterCorrPars *guess)
{

    psMinimization *myMin = psMinimizationAlloc(15, 0.1);

    psVector *params = psVectorAlloc (3, PS_TYPE_F32);
    params->data.F32[0] = guess->scale;
    params->data.F32[1] = guess->offset;
    params->data.F32[2] = guess->offref;
    params->n = 3;

    // XXX for the moment, don't set any constraints
    // psMinConstrain *constrain = psMinConstrainAlloc();
    // constrain->paramDelta = psVectorAlloc (3, PS_TYPE_F32);
    // constrain->paramMin   = psVectorAlloc (3, PS_TYPE_F32);
    // constrain->paramMax   = psVectorAlloc (3, PS_TYPE_F32);
    // constrain->paramMask  = NULL;
    psMinConstrain *constrain = NULL;

    // XXX ignore covariance matrix for the moment
    // psImage *covar = psImageAlloc (params->n, params->n, PS_TYPE_F64);
    psImage *covar = NULL;

    // construct the coordinate and value entries (y is counts)
    psArray *x = psArrayAlloc(exptime->n);
    x->n = exptime->n;

    for (int i = 0; i < exptime->n; i++) {
        psVector *coord = psVectorAlloc(1, PS_TYPE_F32);
        coord->n = 1;
        coord->data.F32[0] = exptime->data.F32[i];
        x->data[i] = coord;
    }

    psMinimizeLMChi2(myMin, covar, params, constrain, x, counts, cntError, pmShutterCorrectionModel);

    pmShutterCorrPars *pars = pmShutterCorrParsAlloc ();

    pars->scale  = params->data.F32[0];
    pars->offset = params->data.F32[1];
    pars->offref = params->data.F32[2];

    psFree (myMin);
    psFree (params);
    psFree (x);

    return (pars);
}
