/** @file  pmAstrometryDistortion.c
*
*  @brief This file defines the basic types for measuring the focal-plane distortion.
*
*  @ingroup AstroImage
*
*  @author EAM, IfA
*
*  @version $Revision: 1.4 $ $Name: not supported by cvs2svn $
*  @date $Date: 2006-06-21 03:21:16 $
*
*  Copyright 2006 Institute for Astronomy, University of Hawaii
*/

/******************************************************************************/
/*  INCLUDE FILES                                                             */
/******************************************************************************/

//#include <stdio.h>
//#include <string.h>
//#include <math.h>
//#include <unistd.h>   // for unlink
#include <pslib.h>
#include "pmFPA.h"
#include "pmAstrometryObjects.h"
#include "pmAstrometryDistortion.h"

static void pmAstromGradientFree (pmAstromGradient *grad)
{

    if (grad == NULL)
        return;

    return;
}

pmAstromGradient *pmAstromGradientAlloc ()
{

    pmAstromGradient *gradient = psAlloc (sizeof(pmAstromGradient));
    psMemSetDeallocator(gradient, (psFreeFunc) pmAstromGradientFree);

    // XXX init any of the data values?

    return (gradient);
}

psArray *pmAstromMeasureGradients(psArray *grads, psArray *rawstars, psArray *refstars, psArray *matches, psMetadata *config)
{

    if (grads == NULL) {
        grads = psArrayAlloc (100);
    }

    // determine range
    // XXX for the moment, I'm hard-wiring this for megacam
    // XXX we need to get the chip dimensions from the metadata somewhere
    int Nx = 2;
    int Ny = 2;
    int DX = 2080 / Nx;
    int DY = 4600 / Ny;

    psPolynomial2D *local = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, 1, 1);
    local->mask[1][1] = 1;

    // measure gradient for fractional chip regions
    for (int nx = 0; nx < Nx; nx++) {
        for (int ny = 0; ny < Ny; ny++) {
            int Xmin = nx*DX;
            int Xmax = Xmin + DX;
            int Ymin = ny*DY;
            int Ymax = Ymin + DY;

            psVector *L  = psVectorAlloc (100, PS_TYPE_F32);
            psVector *M  = psVectorAlloc (100, PS_TYPE_F32);
            psVector *dP = psVectorAlloc (100, PS_TYPE_F32);
            psVector *dQ = psVectorAlloc (100, PS_TYPE_F32);
            int Npts = 0;

            // XXX this is a bit inefficient: first sorting by X or Y could speed this up.
            // XXX or assigning to a segment in a single pass first
            // select the stars within this chip region
            for (int i = 0; i < matches->n; i++) {

                pmAstromMatch *match = matches->data[i];

                pmAstromObj *raw = rawstars->data[match->raw];

                if (raw->chip->x < Xmin)
                    continue;
                if (raw->chip->x > Xmax)
                    continue;
                if (raw->chip->y < Ymin)
                    continue;
                if (raw->chip->y > Ymax)
                    continue;

                pmAstromObj *ref = refstars->data[match->ref];

                L->data.F32[Npts] = ref->FP->x;
                M->data.F32[Npts] = ref->FP->y;

                dP->data.F32[Npts] = ref->TP->x - raw->TP->x;
                dQ->data.F32[Npts] = ref->TP->y - raw->TP->y;

                psVectorExtend (L, 100, 1);
                psVectorExtend (M, 100, 1);
                psVectorExtend (dP, 100, 1);
                psVectorExtend (dQ, 100, 1);
                Npts++;
            }
            if (Npts < 5)
                continue;

            // fit the collection of positions and offsets with a local 1st order gradient
            // 3hi/3lo sigma clipping on the rflux vs metric fit
            psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);
            psVector *mask = psVectorAlloc (Npts, PS_TYPE_MASK);
            mask->n = Npts;
            psVectorInit (mask, 0);

            pmAstromGradient *grad = pmAstromGradientAlloc ();

            psVectorStats (stats, L, NULL, NULL, 0);
            grad->FP.x = stats->sampleMedian;

            psVectorStats (stats, M, NULL, NULL, 0);
            grad->FP.y = stats->sampleMedian;

            psVectorClipFitPolynomial2D (local, stats, mask, 0xff, dP, NULL, L, M);

            grad->dTPdL.x = local->coeff[1][0];
            grad->dTPdM.x = local->coeff[0][1];

            psVectorClipFitPolynomial2D (local, stats, mask, 0xff, dQ, NULL, L, M);

            grad->dTPdL.y = local->coeff[1][0];
            grad->dTPdM.y = local->coeff[0][1];

            psArrayAdd (grads, 100, grad);
        }
    }
    return grads;
}

bool pmAstromFitDistortion(pmFPA *fpa, psArray *grads, psMetadata *config)
{

    psPolynomial2D *localX = NULL;
    psPolynomial2D *localY = NULL;

    // assign the gradient elements to psVectors for fitting
    psVector *dPdL = psVectorAlloc (grads->n, PS_TYPE_F32);
    psVector *dQdL = psVectorAlloc (grads->n, PS_TYPE_F32);
    psVector *dPdM = psVectorAlloc (grads->n, PS_TYPE_F32);
    psVector *dQdM = psVectorAlloc (grads->n, PS_TYPE_F32);
    psVector *L = psVectorAlloc (grads->n, PS_TYPE_F32);
    psVector *M = psVectorAlloc (grads->n, PS_TYPE_F32);

    for (int i = 0; i < grads->n; i++) {

        pmAstromGradient *grad = grads->data[i];

        dPdL->data.F32[i] = grad->dTPdL.x;
        dQdL->data.F32[i] = grad->dTPdL.y;

        dPdM->data.F32[i] = grad->dTPdM.x;
        dQdM->data.F32[i] = grad->dTPdM.y;

        L->data.F32[i] = grad->FP.x;
        M->data.F32[i] = grad->FP.y;
    }
    dPdL->n = dQdL->n = dPdM->n = dQdM->n = grads->n;
    L->n = M->n = grads->n;

    psStats *stats = psStatsAlloc (PS_STAT_SAMPLE_MEDIAN | PS_STAT_SAMPLE_STDEV);
    psVector *mask = psVectorAlloc (grads->n, PS_TYPE_MASK);
    mask->n = grads->n;
    psVectorInit (mask, 0);

    // the order of the gradient fits need to be 1 less than the distortion term
    // determine the gradient order(s) from the fpa->toTP structure
    localX = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, fpa->toTangentPlane->x->nX-1, fpa->toTangentPlane->x->nY);
    localY = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, fpa->toTangentPlane->x->nX,   fpa->toTangentPlane->x->nY-1);

    // set masks based on fpa->toTP
    for (int i = 0; i < fpa->toTangentPlane->x->nX; i++) {
        for (int j = 0; j < fpa->toTangentPlane->x->nY; j++) {
            if (fpa->toTangentPlane->x->mask[i][j][0][0]) {
                localX->mask[i-1][j] = 1;
                localY->mask[i][j-1] = 1;
            }
        }
    }

    // fit the local gradients in each direction
    psVectorClipFitPolynomial2D (localX, stats, mask, 0xff, dPdL, NULL, L, M);
    psVectorClipFitPolynomial2D (localY, stats, mask, 0xff, dPdM, NULL, L, M);

    // XXX test: generate fit
    psVector *dPdLf = psPolynomial2DEvalVector (localX, L, M);
    psVector *dPdMf = psPolynomial2DEvalVector (localY, L, M);

    // update fpa->toTP distortion terms
    fpa->toTangentPlane->x->coeff[0][0][0][0] = 0;
    for (int i = 1; i < fpa->toTangentPlane->x->nX; i++) {
        if (!fpa->toTangentPlane->x->mask[i][0][0][0]) {
            fpa->toTangentPlane->x->coeff[i][0][0][0] = localX->coeff[i-1][0] / i;
        }
    }
    for (int j = 1; j < fpa->toTangentPlane->x->nY; j++) {
        if (!fpa->toTangentPlane->x->mask[0][j][0][0]) {
            fpa->toTangentPlane->x->coeff[0][j][0][0] = localY->coeff[0][j-1] / j;
        }
    }
    for (int i = 1; i < fpa->toTangentPlane->x->nX; i++) {
        for (int j = 1; j < fpa->toTangentPlane->x->nY; j++) {
            if (!fpa->toTangentPlane->x->mask[i][j][0][0]) {
                fpa->toTangentPlane->x->coeff[i][j][0][0] = 0.5*(localX->coeff[i-1][j] / i + localY->coeff[i][j-1] / j);
            }
        }
    }
    psFree (localX);
    psFree (localY);

    // the order of the gradient fits need to be 1 less than the distortion term
    // determine the gradient order(s) from the fpa->toTP structure
    localX = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, fpa->toTangentPlane->y->nX-1, fpa->toTangentPlane->y->nY);
    localY = psPolynomial2DAlloc (PS_POLYNOMIAL_ORD, fpa->toTangentPlane->y->nX,   fpa->toTangentPlane->y->nY-1);

    // set masks based on fpa->toTP
    for (int i = 0; i < fpa->toTangentPlane->y->nX; i++) {
        for (int j = 0; j < fpa->toTangentPlane->y->nY; j++) {
            if (fpa->toTangentPlane->y->mask[i][j][0][0]) {
                localX->mask[i-1][j] = 1;
                localY->mask[i][j-1] = 1;
            }
        }
    }

    // fit the local gradients in each direction
    psVectorClipFitPolynomial2D (localX, stats, mask, 0xff, dQdL, NULL, L, M);
    psVectorClipFitPolynomial2D (localY, stats, mask, 0xff, dQdM, NULL, L, M);

    // XXX test: generate fit
    psVector *dQdLf = psPolynomial2DEvalVector (localX, L, M);
    psVector *dQdMf = psPolynomial2DEvalVector (localY, L, M);

    FILE *f = fopen ("grads.dat", "w");
    for (int i = 0; i < dPdLf->n; i++) {
        fprintf (f, "%f %f  | %f %f  %f %f | %f %f  %f %f\n",
                 L->data.F32[i], M->data.F32[i],
                 dPdL->data.F32[i], dPdM->data.F32[i],
                 dQdL->data.F32[i], dQdM->data.F32[i],
                 dPdLf->data.F32[i], dPdMf->data.F32[i],
                 dQdLf->data.F32[i], dQdMf->data.F32[i]);
    }
    fclose (f);

    // update fpa->toTP distortion terms
    fpa->toTangentPlane->y->coeff[0][0][0][0] = 0;
    for (int i = 1; i < fpa->toTangentPlane->y->nX; i++) {
        if (!fpa->toTangentPlane->y->mask[i][0][0][0]) {
            fpa->toTangentPlane->y->coeff[i][0][0][0] = localX->coeff[i-1][0] / i;
        }
    }
    for (int j = 1; j < fpa->toTangentPlane->y->nY; j++) {
        if (!fpa->toTangentPlane->y->mask[0][j][0][0]) {
            fpa->toTangentPlane->y->coeff[0][j][0][0] = localY->coeff[0][j-1] / j;
        }
    }
    for (int i = 1; i < fpa->toTangentPlane->y->nX; i++) {
        for (int j = 1; j < fpa->toTangentPlane->y->nY; j++) {
            if (!fpa->toTangentPlane->y->mask[i][j][0][0]) {
                fpa->toTangentPlane->y->coeff[i][j][0][0] = 0.5*(localX->coeff[i-1][j] / i + localY->coeff[i][j-1] / j);
            }
        }
    }
    psFree (localX);
    psFree (localY);

    // XXX free unneeded structures
    psFree (dPdL);
    psFree (dPdM);
    psFree (dQdL);
    psFree (dQdM);
    psFree (L);
    psFree (M);
    psFree (stats);
    psFree (mask);

    return true;
}
