/** @file  pmReadoutCombine.c
 *
 *  This file will contain a module which will combine multiple readout images.
 *
 *  @author GLG, MHPCC
 *
 *  @version $Revision: 1.5 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2006-02-06 21:03:25 $
 *
 *  Copyright 2004 Maui High Performance Computing Center, University of Hawaii
 *
 */

#include<stdio.h>
#include<math.h>
#include "pslib.h"
#include "pmReadoutCombine.h"

/******************************************************************************
DetermineNumBits(data): This routine takes an enum psStatsOptions as an
argument and returns the number of non-zero bits.
 *****************************************************************************/
static psS32 DetermineNumBits(psStatsOptions data)
{
    psS32 i;
    psU64 tmpData = data;
    psS32 numBits = 0;

    for (i=0;i<(8 * sizeof(psStatsOptions));i++) {
        if (0x0001 & tmpData) {
            numBits++;
        }
        tmpData = tmpData >> 1;
    }
    return(numBits);
}

/******************************************************************************
XXX: Must add support for S16 and S32 types.  F32 currently supported.
 *****************************************************************************/
psImage *pmReadoutCombine(psImage *output,
                          const psList *inputs,
                          psCombineParams *params,
                          const psVector *zero,
                          const psVector *scale,
                          bool applyZeroScale,
                          psF32 gain,
                          psF32 readnoise)
{
    PS_ASSERT_PTR_NON_NULL(inputs, NULL);
    PS_ASSERT_PTR_NON_NULL(params, NULL);
    PS_ASSERT_PTR_NON_NULL(params->stats, NULL);
    if (zero != NULL) {
        PS_ASSERT_VECTOR_TYPE(zero, PS_TYPE_F32, NULL);
        //        PS_ASSERT_VECTOR_TYPE_S16_S32_F32(zero, NULL);
    }
    if (scale != NULL) {
        PS_ASSERT_VECTOR_TYPE(scale, PS_TYPE_F32, NULL);
        //        PS_ASSERT_VECTOR_TYPE_S16_S32_F32(scale, NULL);
    }
    if ((zero != NULL) && (scale != NULL)) {
        PS_ASSERT_VECTOR_TYPE_EQUAL(zero, scale, NULL);
        // PS_ASSERT_VECTOR_TYPE_S16_S32_F32(scale, NULL);
    }

    psStats *stats = params->stats;
    psS32 i;
    psS32 j;
    psS32 maxInputCols = 0;
    psS32 maxInputRows = 0;
    psS32 minInputCols = PS_MAX_S32;
    psS32 minInputRows = PS_MAX_S32;
    psListElem *tmpInput = NULL;
    pmReadout *tmpReadout = NULL;
    psS32 numInputs = 0;
    psS32 tmpI;
    psElemType outputType = PS_TYPE_F32;

    if (1 < DetermineNumBits(params->stats->options)) {
        psError(PS_ERR_UNKNOWN, true,
                "Multiple statistical options have been requested.  Returning NULL.\n");
        return(NULL);
    }

    //
    // We step through each readout in the input image list.  If any readout is
    // NULL, empty, or has the wrong type, we generate an error and return
    // NULL.  We determine how big of an output image is needed to combine
    // these input images.  We do this by taking the
    //     max(readout->col0 + readout->numCols + image->col0 + image->numCols)
    //     max(readout->row0 + readout->numRows + image->row0 + image->numRows)
    // We then compare that to
    //     output->col0 + output->numCols
    //     output->row0 + output->numRows
    // to determine if the output image actually stores that pixel.  A similar
    // thing is done for the minimum row and column.
    //
    tmpInput = (psListElem *) inputs->head;
    while (NULL != tmpInput) {
        tmpReadout = (pmReadout *) tmpInput->data;
        PS_ASSERT_READOUT_NON_NULL(tmpReadout, output);
        PS_ASSERT_READOUT_NON_EMPTY(tmpReadout, output);
        PS_ASSERT_READOUT_TYPE(tmpReadout, PS_TYPE_F32, output);

        outputType = tmpReadout->image->type.type;

        minInputRows = PS_MIN(minInputRows,
                              (tmpReadout->row0 + tmpReadout->image->row0));
        tmpI = tmpReadout->row0 +
               tmpReadout->image->row0 +
               tmpReadout->image->numRows;
        maxInputRows = PS_MAX(maxInputRows, tmpI);

        minInputCols = PS_MIN(minInputCols,
                              (tmpReadout->col0 + tmpReadout->image->col0));
        tmpI = tmpReadout->col0 +
               tmpReadout->image->col0 +
               tmpReadout->image->numCols;
        maxInputCols = PS_MAX(maxInputCols, tmpI);
        tmpInput = tmpInput->next;
        numInputs++;
    }

    // We ensure that the zero vector is of the proper size.
    if (zero != NULL) {
        PS_ASSERT_VECTOR_TYPE(zero, PS_TYPE_F32, NULL);
        if (numInputs > zero->n) {
            psError(PS_ERR_UNKNOWN, true, "zero vector has incorrect size (%d).  Returning NULL.\n", zero->n);
            return(NULL);
        } else if (numInputs < zero->n) {
            psLogMsg(__func__, PS_LOG_WARN,
                     "WARNING: the zero vector too many elements (%d)\n", zero->n);
        }
    }

    // We ensure that the scale vector is of the proper size.
    if (scale != NULL) {
        PS_ASSERT_VECTOR_TYPE(scale, PS_TYPE_F32, NULL);
        if (numInputs > scale->n) {
            psError(PS_ERR_UNKNOWN, true, "scale vector has incorrect size (%d).  Returning NULL.\n", scale->n);
            return(NULL);
        } else if (numInputs < scale->n) {
            psLogMsg(__func__, PS_LOG_WARN,
                     "WARNING: the scale vector has too many elements (%d)\n", scale->n);
        }
    }

    // At this point, the following variables have been computed:
    // maxInputRows: the largest input row value, in output image space.
    // maxInputCols: the largest input column value, in output image space.
    // minInputRows: the smallest input row value, in output image space.
    // minInputCols: the smallest input column value, in output image space.
    //
    if (output == NULL) {
        output = psImageAlloc(maxInputCols-minInputCols,
                              maxInputRows-minInputRows, outputType);
        *(psS32 *) &(output->col0) = minInputCols;
        *(psS32 *) &(output->row0) = minInputRows;
    } else {
        PS_ASSERT_IMAGE_TYPE(output, PS_TYPE_F32, NULL);
        if (((output->col0 + output->numCols) < maxInputCols) ||
                ((output->row0 + output->numRows) < maxInputRows)) {
            psError(PS_ERR_UNKNOWN, true,
                    "Output image (%d, %d) is too small to hold combined images.  Returning NULL.\n",
                    output->row0 + output->numRows,
                    output->col0 + output->numCols);
            return(NULL);
        }

        if ((output->col0 > minInputCols) || (output->row0 > minInputRows)) {
            psError(PS_ERR_UNKNOWN, true,
                    "Output image offset is larger then input image offset.  Returning NULL.\n");
            return(NULL);
        }
    }

    psVector *tmpPixels = psVectorAlloc(numInputs, PS_TYPE_F32);
    psVector *tmpPixelErrors = psVectorAlloc(numInputs, PS_TYPE_F32);
    psVector *tmpPixelMask = psVectorAlloc(numInputs, PS_TYPE_U8);
    psVector *tmpPixelMaskNKeep = psVectorAlloc(numInputs, PS_TYPE_U8);
    psVector *outRowLower = psVectorAlloc(numInputs, PS_TYPE_U32);
    psVector *outRowUpper = psVectorAlloc(numInputs, PS_TYPE_U32);
    psVector *outColLower = psVectorAlloc(numInputs, PS_TYPE_U32);
    psVector *outColUpper = psVectorAlloc(numInputs, PS_TYPE_U32);
    pmReadout **tmpReadouts = (pmReadout **) psAlloc(numInputs * sizeof(pmReadout *));

    // For each input readout, we create a pointer to that readout in
    // "tmpReadouts[]", and we store the min/max pixel indices for that
    // readout, in output image coordinates, in the psVectors
    // (outRowLower, outColLower, outRowUpper, outColUpper).
    i = 0;
    tmpInput = (psListElem *) inputs->head;
    while (NULL != tmpInput) {
        tmpReadouts[i] = (pmReadout *) tmpInput->data;
        outRowLower->data.U32[i] = tmpReadouts[i]->row0 + tmpReadouts[i]->image->row0;
        outColLower->data.U32[i] = tmpReadouts[i]->col0 + tmpReadouts[i]->image->col0;
        outRowUpper->data.U32[i] = tmpReadouts[i]->row0 +
                                   tmpReadouts[i]->image->row0 +
                                   tmpReadouts[i]->image->numRows;
        outColUpper->data.U32[i] = tmpReadouts[i]->col0 +
                                   tmpReadouts[i]->image->col0 +
                                   tmpReadouts[i]->image->numCols;
        tmpInput = tmpInput->next;
        i++;
    }

    // We loop through each pixel in the output image.  We loop through each
    // input readout.  We determine if that output pixel is contained in the
    // image from that readout.  If so, we save it in psVector tmpPixels.
    // If not, we set a mask for that element in tmpPixels.  Then, we mask off
    // pixels not between fracLow and fracHigh.  Then we call the vector
    // stats routine on those pixels/mask.  Then we set the output pixel value
    // to the result of the stats call.

    for (i = output->row0; i < (output->row0 + output->numRows) ; i++) {
        for (j = output->col0; j < (output->col0 + output->numCols) ; j++) {
            for (psS32 r = 0; r < numInputs ; r++) {
                //  printf("[%d][%d]: [%d][%d] to [%d][%d]\n", i, j, outRowLower->data.U32[r], outColLower->data.U32[r], outRowUpper->data.U32[r], outColUpper->data.U32[r]);
                if ((outRowLower->data.U32[r] <= i) &&
                        (outColLower->data.U32[r] <= j) &&
                        (outRowUpper->data.U32[r] > i) &&
                        (outColUpper->data.U32[r] > j)) {

                    psS32 imageRow = i - (tmpReadouts[r]->row0 +
                                          tmpReadouts[r]->image->row0);
                    psS32 imageCol = j - (tmpReadouts[r]->col0 +
                                          tmpReadouts[r]->image->col0);

                    if ((NULL == tmpReadouts[r]->mask) ||
                            !(params->maskVal && tmpReadouts[r]->mask->data.U8[imageRow][imageCol])) {
                        tmpPixels->data.F32[r] = tmpReadouts[r]->image->data.F32[imageRow][imageCol];
                        tmpPixelMask->data.U8[r] = 0;
                    } else {
                        tmpPixels->data.F32[r] = 0.0;
                        tmpPixelMask->data.U8[r] = 1;
                    }
                } else {
                    tmpPixels->data.F32[r] = 0.0;
                    tmpPixelMask->data.U8[r] = 1;
                }
                // printf("readout[%d], image [%d][%d] is %f\n", r, i, j, tmpPixels->data.F32[r]);
            }
            // At this point, we have scanned all input readouts for this
            // one output pixel.
            //            for (psS32 r = 0; r < numInputs ; r++) printf("(0)tmpPixels->data.F32[%d] is %f\n", r, tmpPixels->data.F32[r]);

            // Determine how many pixels lie between fracLow and fracHigh.
            psS32 pixelCount = 0;
            for (psS32 r = 0; r < numInputs ; r++) {
                if (tmpPixelMask->data.U8[r] == 0) {
                    if ((params->fracLow <= tmpPixels->data.F32[r]) &&
                            (params->fracHigh >= tmpPixels->data.F32[r])) {
                        pixelCount++;
                    }
                }
            }

            // If more than params->nKeep pixels lie between the valid range,
            // then loop through the pixels, and mask away any pixels outside
            // that range.
            if (pixelCount >= params->nKeep) {
                for (psS32 r = 0; r < numInputs ; r++) {
                    if (tmpPixelMask->data.U8[r] == 0) {
                        if ((params->fracLow <= tmpPixels->data.F32[r]) &&
                                (params->fracHigh >= tmpPixels->data.F32[r])) {
                            tmpPixelMaskNKeep->data.U8[r] = 0;
                        } else {
                            tmpPixelMaskNKeep->data.U8[r] = 1;
                        }
                    }
                }
            }

            if ((gain > 0.0) && (readnoise >= 0.0)) {
                psF32 x;
                psF32 sigma;
                if (applyZeroScale == true) {
                    for (psS32 r = 0; r < numInputs ; r++) {
                        if (zero != NULL) {
                            x = zero->data.F32[r];
                        } else {
                            x = 0.0;
                        }
                        if (scale != NULL) {
                            x+= tmpPixels->data.F32[r] * scale->data.F32[r];
                        } else {
                            x+= tmpPixels->data.F32[r];
                        }
                        sigma = sqrtf((readnoise*readnoise) + gain * x) / gain;

                        tmpPixelErrors->data.F32[r] = sigma;
                        tmpPixels->data.F32[r]= x;
                    }
                } else {
                    for (psS32 r = 0; r < numInputs ; r++) {
                        x= tmpPixels->data.F32[r];

                        if (zero != NULL) {
                            sigma = zero->data.F32[r];
                        } else {
                            sigma = 0.0;
                        }
                        if (scale != NULL) {
                            sigma+= tmpPixels->data.F32[r] * scale->data.F32[r];
                        } else {
                            sigma+= tmpPixels->data.F32[r];
                        }
                        sigma = sqrtf((readnoise*readnoise) + (gain * sigma)) / gain;

                        tmpPixelErrors->data.F32[r] = sigma;
                        tmpPixels->data.F32[r]= x;
                    }
                }
                // Calculate the specified statistic on the stack of pixels.
                //                for (psS32 r = 0; r < numInputs ; r++) printf("(1)tmpPixels->data.F32[%d] is %f\n", r, tmpPixels->data.F32[r]);
                psStats *rc = psVectorStats(stats,
                                            tmpPixels,
                                            tmpPixelErrors,
                                            tmpPixelMaskNKeep,
                                            1);
                if (rc == NULL) {
                    psError(PS_ERR_UNKNOWN, false, "psVectorStats(): could not perform requested statistical operation.  Returning NULL.\n");
                    return(NULL);
                }
            } else {
                if (scale != NULL) {
                    for (psS32 r = 0; r < numInputs ; r++) {
                        tmpPixels->data.F32[r]*= scale->data.F32[r];
                    }
                }

                // We add the zero vector, if non-NULL.
                if (zero != NULL) {
                    for (psS32 r = 0; r < numInputs ; r++) {
                        tmpPixels->data.F32[r]+= zero->data.F32[r];
                    }
                }

                // Calculate the specified statistic on the stack of pixels.
                //                for (psS32 r = 0; r < numInputs ; r++) printf("(2)tmpPixels->data.F32[%d] is %f\n", r, tmpPixels->data.F32[r]);
                psStats *rc = psVectorStats(stats,
                                            tmpPixels,
                                            NULL,
                                            tmpPixelMaskNKeep,
                                            1);
                if (rc == NULL) {
                    psError(PS_ERR_UNKNOWN, false, "psVectorStats(): could not perform requested statistical operation.  Returning NULL.\n");
                    return(NULL);
                }
            }


            // Set the pixel value in the output image to the stat value.
            double statValue;
            if (!p_psGetStatValue(stats, &statValue)) {
                psError(PS_ERR_UNKNOWN, true, "Could not determine stats value.  Returning NULL.\n");
                return(NULL);
            } else {
                output->data.F32[i-output->row0][j-output->col0] = (psF32) statValue;
            }
        }
    }

    psFree(tmpPixels);
    psFree(tmpPixelErrors);
    psFree(tmpPixelMask);
    psFree(tmpPixelMaskNKeep);
    psFree(outRowLower);
    psFree(outRowUpper);
    psFree(outColLower);
    psFree(outColUpper);
    psFree(tmpReadouts);

    return(output);
}


/* This function measures the robust median at each of the minimum and maximum
 * coordinates and determines the difference and mean of the two values. The size
 * of the box used to make the measurement at each point is specified by the
 * configuration variable FRINGE_SQUARE_RADIUS. From the collection of
 * differences, the robust median is calculated, and returned as part of the
 * fringe statistics. For each fringe point, the values of delta and midValue are
 * also assigned and available to the user on return.
 */

psStats *pmFringeStats(
    psArray *fringePoints,
    psImage *image,
    psMetadata *config)
{
    PS_ASSERT_PTR_NON_NULL(fringePoints, NULL);
    //    for (psS32 i = 0 ; i < fringePoints->n ; i++) {
    //        if (!psMemCheckFringePoint((pmFringePoint *) fringePoints->data[i])) {
    //            psError(PS_ERR_UNKNOWN, true, "fringePoints->data[%d] is not of type pmFringePoint.\n");
    //            return(NULL);
    //        }
    //    }
    PS_ASSERT_IMAGE_NON_NULL(image, NULL);
    PS_ASSERT_IMAGE_NON_EMPTY(image, NULL);
    PS_ASSERT_IMAGE_TYPE(image, PS_TYPE_F32, NULL);
    PS_ASSERT_PTR_NON_NULL(config, NULL);

    psBool rc;
    psF32 frSquareRadius = psMetadataLookupF32(&rc, config, "FRINGE_SQUARE_RADIUS");
    if (!rc) {
        psError(PS_ERR_UNKNOWN, true, "Could not determing the fringe radius from the metadata.\n");
    }

    psRegion minRegion;
    psRegion maxRegion;
    psStats *minStats = psStatsAlloc(PS_STAT_FITTED_MEAN);
    psStats *maxStats = psStatsAlloc(PS_STAT_FITTED_MEAN);
    psStats *diffStats = psStatsAlloc(PS_STAT_ROBUST_MEDIAN);
    psVector *diffs = psVectorAlloc(fringePoints->n, PS_TYPE_F32);

    //
    // Loop through each fringe point.  Determine the robust mean around
    // the min and max for that fringe point.  Save the difference between
    // those two numbers in psVector diffs.
    //
    // XXX: Ensure you have the radius correct.  (add 1 to x1 and y1?)
    //
    for (psS32 i = 0 ; i < fringePoints->n ; i++) {
        pmFringePoint *fp = (pmFringePoint *) fringePoints->data[i];
        minRegion.x0 = fp->xMin - frSquareRadius;
        minRegion.x1 = fp->xMin + frSquareRadius;
        minRegion.y0 = fp->yMin - frSquareRadius;
        minRegion.y1 = fp->yMin + frSquareRadius;
        psImage *minSubImage = psImageSubset(image, minRegion);
        minStats = psImageStats(minStats, minSubImage, NULL, 0);

        maxRegion.x0 = fp->xMax - frSquareRadius;
        maxRegion.x1 = fp->xMax + frSquareRadius;
        maxRegion.y0 = fp->yMax - frSquareRadius;
        maxRegion.y1 = fp->yMax + frSquareRadius;
        psImage *maxSubImage = psImageSubset(image, maxRegion);
        maxStats = psImageStats(maxStats, maxSubImage, NULL, 0);

        if ((minStats == NULL) || (maxStats == NULL)) {
            psError(PS_ERR_UNKNOWN, true, "Could not determine robust mean on subimage.\n");
            psFree(minStats);
            psFree(maxStats);
            return(NULL);
        }

        fp->midValue = 0.5 * (maxStats->fittedMean + minStats->fittedMean);
        fp->delta = maxStats->fittedMean - minStats->fittedMean;
        diffs->data.F32[i] = fp->delta;
    }
    psFree(minStats);
    psFree(maxStats);

    diffStats = psVectorStats(diffStats, diffs, NULL, NULL, 0);
    psFree(diffs);
    if (diffStats == NULL) {
        psError(PS_ERR_UNKNOWN, true, "Could not determine fitted median of the differences.\n");
        return(NULL);
    }
    return(diffStats);
}



/**
 *
 * The input array fluxLevels consists of Ni vectors, one per mosaic image.
 * Each vector consists of Nj elements, each a measurement of the input
 * flat-field image flux levels. All of these vectors must be constructed with
 * the same number of elements, or the function will return an error. If a chip
 * is missing from a particular image, that element should be set to NaN. The
 * vector chipGains supplies initial guesses for the chip gains. If the vector
 * contains the values 0.0 or NaN for any of the elements, the gain is set to the
 * mean of the valid values. If the vector length does not match the number of
 * chips, an warning is raised, all chip gain guesses will be set to 1.0, and the
 * vector length modified to match the number of chips defined by the supplied
 * fluxLevels. The sourceFlux input vector must be allocated (not NULL), but the
 * routine will set the vector length to the number of source images regardless
 * of the initial state of the vector. All vectors used by this function must be
 * of type PS_DATA_F64.
 *
 
fluxLevels(i, j): for each flat field image i, this psArray contains a vector
with an elemenmt for each chip j.  So, fluxLevels(i, j) corresponds to the
measured flux M_(i, j) for flat image i, chip j.
 
chipGains[]: has j elements, one for each chip.
 
 
They have the observed flux levels for each chip of each image.  They want to
solve for the actual flux levels and the gain of each chip.
 
Okay, they want to solve for source fluxes and chip gains.
 
 *
 */
bool pmFlatNormalization(
    psVector *sourceFlux,
    psVector *chipGains,
    psArray *fluxLevels)
{
    PS_ASSERT_PTR_NON_NULL(fluxLevels, false);
    psS32 numImages = fluxLevels->n;
    psS32 numChips = ((psVector *) fluxLevels->data[0])->n;
    for (psS32 i = 0 ; i < numImages ; i++) {
        psVector *tmpVec = (psVector *) fluxLevels->data[i];
        PS_ASSERT_VECTOR_NON_NULL(tmpVec, false);
        PS_ASSERT_VECTOR_TYPE(tmpVec, PS_TYPE_F64, false);
        PS_ASSERT_VECTOR_SIZE(tmpVec, numChips, false);
    }

    //
    // Ensure that *localChipGains points to a vector of the same length as numImages.
    //
    PS_ASSERT_PTR_NON_NULL(chipGains, false);
    PS_ASSERT_VECTOR_TYPE(chipGains, PS_TYPE_F64, false);
    psVector *localChipGains = chipGains;
    if (numChips != chipGains->n) {
        psLogMsg(__func__, PS_LOG_WARN, "WARNING: the chipGains vector length does not match the number of chips.\n");
        localChipGains = psVectorAlloc(numChips, PS_TYPE_F64);
        psBool rc = psVectorInit(localChipGains, 1.0);
        if (rc == false) {
            printf("XXX: gen error\n");
        }
    }

    //
    // If the chipGains vector contains the values 0.0 or NaN for any of the elements,
    // the gain is set to the mean of the valid values.
    //
    psBool meanFlag = false;
    psVector *chipGainsMask = psVectorAlloc(chipGains->n, PS_TYPE_U8);
    for (psS32 i = 0 ; i < chipGains->n ; i++) {
        if ((fabs(chipGains->data.F64[i]) < FLT_EPSILON) ||
                (isnan(chipGains->data.F64[i]))) {
            chipGainsMask->data.U8[i] = 1;
            meanFlag = true;
        }
    }
    // Must calculate the mean.
    if (meanFlag == true) {
        psStats *stats = psStatsAlloc(PS_STAT_SAMPLE_MEAN);
        stats = psVectorStats(stats, chipGains, NULL, chipGainsMask, 1);
        if (stats == NULL) {
            printf("XXX: gen error\n");
        }
        psF64 mean;
        psBool rc = p_psGetStatValue(stats, &mean);
        if (rc == false) {
            printf("XXX: gen error\n");
        }
        // Set the gain to this mean for chips with a gain of 0.0 or NAN

        for (psS32 i = 0 ; i < chipGains->n ; i++) {
            if ((fabs(chipGains->data.F64[i]) < FLT_EPSILON) ||
                    (isnan(chipGains->data.F64[i]))) {
                chipGains->data.F64[i] = mean;
            }
        }
    }

    //
    // Assert that sourceFlux is non-NULL, correct type, correct size.
    //
    PS_ASSERT_PTR_NON_NULL(sourceFlux, false);
    PS_ASSERT_VECTOR_TYPE(sourceFlux, PS_TYPE_F64, false);
    psVectorRealloc(sourceFlux, numImages);

    //    psFree(psVector);
    if (numImages != chipGains->n) {
        psFree(localChipGains);
    }

    return(true);
}
