/** @file  pmStack.c
 *
 *  This file will perform image combination of several images of the
 *  same field, produce a list of questionable pixels, then tag some
 *  of those pixels as cosmic rays.
 *
 *  @author Paul Price, IfA
 *  @author GLG, MHPCC
 *
 *  @version $Revision: 1.48.2.1 $ $Name: not supported by cvs2svn $
 *  @date $Date: 2009-02-19 17:59:50 $
 *  Copyright 2004-2007 Institute for Astronomy, University of Hawaii
 *
 */

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

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

#include "pmHDU.h"
#include "pmFPA.h"
#include "pmReadoutStack.h"
#include "pmConceptsAverage.h"

#include "pmStack.h"

#define PIXEL_LIST_BUFFER 100           // Number of entries to add to pixel list at a time
#define PIXEL_MAP_BUFFER 2              // Number of entries to add to pixel map at a time
//#define ADD_VARIANCE                    // Allow additional variance (besides variance factor)?
#define NUM_DIRECT_STDEV 5              // For less than this number of values, measure stdev directly


//#define TESTING                         // Enable test output
//#define TEST_X 843-1                     // x coordinate to examine
//#define TEST_Y 813-1                     // y coordinate to examine
//#define TEST_RADIUS 0                    // Radius to examine


// Data structure for use as a buffer when combining pixels
// Use of this structure means we don't have to do an allocation in the combination function for each pixel
typedef struct {
    psVector *pixels;                   // Pixel values
    psVector *variances;                // Pixel variances
    psVector *weights;                  // Pixel weightings
    psVector *exps;                     // Pixel exposures
    psVector *sources;                  // Pixel sources (which image did they come from?)
    psVector *limits;                   // Rejection limits
    psVector *suspects;                 // Pixel is suspect?
    psVector *sort;                     // Buffer for sorting (to get a robust estimator of the standard dev)
} combineBuffer;

static void combineBufferFree(combineBuffer *buffer)
{
    psFree(buffer->pixels);
    psFree(buffer->variances);
    psFree(buffer->weights);
    psFree(buffer->exps);
    psFree(buffer->sources);
    psFree(buffer->limits);
    psFree(buffer->suspects);
    psFree(buffer->sort);
    return;
}

static combineBuffer *combineBufferAlloc(long numImages // Number of images that will be combined
    )
{
    combineBuffer *buffer = psAlloc(sizeof(combineBuffer));
    psMemSetDeallocator(buffer, (psFreeFunc)combineBufferFree);

    buffer->pixels = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->variances = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->weights = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->exps = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->sources = psVectorAlloc(numImages, PS_TYPE_U16);
    buffer->limits = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->suspects = psVectorAlloc(numImages, PS_TYPE_U8);
    buffer->sort = psVectorAlloc(numImages, PS_TYPE_F32);
    return buffer;
}


// Deallocator for the stack data
static void stackDataFree(pmStackData *data)
{
    psFree(data->readout);
    psFree(data->reject);
    psFree(data->inspect);
    return;
}


// Determine a mean value and variance for the combination
// Not using psVectorStats because it assumes that the "weights" are errors, and weights by 1/error^2
static bool combinationMeanVariance(float *mean, // Mean value, to return
                                    float *var, // Variance value, to return
                                    float *exp, // Exposure time, to return
                                    float *expWeight,          // Weighted exposure time, to return
                                    const psVector *values, // Values to combine
                                    const psVector *variances, // Pixel variances to combine
                                    const psVector *exps,      // Exposure times to combine
                                    const psVector *weights // Weights to apply
                                    )
{
    assert(mean);
    assert(values && weights);
    assert(values->n == weights->n);
    assert((var && variances) || !var);
    assert(!variances || variances->n == values->n);
    assert(values->type.type == PS_TYPE_F32);
    assert(!values || values->type.type == PS_TYPE_F32);
    assert(weights->type.type == PS_TYPE_F32);

    // We're not using the input pixel variances to generate a weighted average for the pixel flux (because
    // that introduces systematic biases), so the variance of the output pixel value should simply be:
    //     simga^2 = sum(weight_i^2 * sigma_i^2) / (sum(weight_i))^2
    // This reduces, when the weights are all identically unity, to:
    //     variance_combination = sum(variance_i) / N^2
    // and if the variances are all equal:
    //     variance_combination = variance_individual / N
    // which makes sense --- the standard deviation of the combination is reduced by a factor of sqrt(N).

    float sumValueWeight = 0.0;         // Sum of the value multiplied by the weight
    float sumVarianceWeight = 0.0;     // Sum of the pixel variances multiplied by the global weights
    float sumWeight = 0.0;              // Sum of the image weights
    float sumExp = 0.0;                 // Sum of the exposure time
    float sumExpWeight = 0.0;           // Sum of the exposure time multiplied by the global weights
    int numGood = 0;                    // Number of good exposures
    for (int i = 0; i < values->n; i++) {
        sumValueWeight += values->data.F32[i] * weights->data.F32[i];
        sumWeight += weights->data.F32[i];
        if (variances) {
            sumVarianceWeight += variances->data.F32[i] * PS_SQR(weights->data.F32[i]);
        }
        if (exps) {
            sumExp += exps->data.F32[i];
            sumExpWeight += exps->data.F32[i] * weights->data.F32[i];
            numGood++;
        }
    }

    if (sumWeight <= 0) {
        return false;
    }

    *mean = sumValueWeight / sumWeight;
    if (var) {
        *var = sumVarianceWeight / PS_SQR(sumWeight);
    }
    if (exp) {
        *exp = sumExp;
    }
    if (expWeight) {
        *expWeight = sumExpWeight;
    }
    return true;
}

// Return the median and standard deviation for the pixels
// Not using psVectorStats because it has additional allocations which slow things down
static bool combinationMedianStdev(float *median, // Median value, to return
                                   float *stdev, // Standard deviation value, to return
                                   const psVector *values, // Values to combine
                                   psVector *sortBuffer // Buffer for sorting
                                   )
{
    assert(values);
    assert(values->type.type == PS_TYPE_F32);
    assert(sortBuffer && sortBuffer->nalloc >= values->n && sortBuffer->type.type == PS_TYPE_F32);

    int num = values->n;                // Number of values
    sortBuffer = psVectorSortIndex(sortBuffer, values);
    if (!sortBuffer) {
        *median = NAN;
        *stdev = NAN;
        return false;
    }

    if (num == 3) {
        // Attempt to measure standard deviation with only three values (and one of those possibly corrupted)
        *median = values->data.F32[sortBuffer->data.S32[1]];
        if (stdev) {
            float diff1 = values->data.F32[sortBuffer->data.S32[0]] - *median;
            float diff2 = values->data.F32[sortBuffer->data.S32[2]] - *median;
            // This factor of sqrt(2) might not be exact, but it's about right
            *stdev = M_SQRT2 * PS_MIN(fabsf(diff1), fabsf(diff2));
        }
    } else {
        *median = num % 2 ? values->data.F32[sortBuffer->data.S32[num / 2]] :
            (values->data.F32[sortBuffer->data.S32[num / 2 - 1]] +
             values->data.F32[sortBuffer->data.S32[num / 2]]) / 2.0;
        if (stdev) {
            if (num <= NUM_DIRECT_STDEV) {
                // If there are not many values, the direct standard deviation is better
                double sum = 0.0;
                for (int i = 0; i < num; i++) {
                    sum += PS_SQR(values->data.F32[sortBuffer->data.S32[i]] - *median);
                }
                *stdev = sqrt(sum / (double)(num - 1));
            } else {
                // Standard deviation from the interquartile range
                *stdev = 0.74 * (values->data.F32[sortBuffer->data.S32[(int)(0.75 * num)]] -
                                 values->data.F32[sortBuffer->data.S32[(int)(0.25 * num)]]);
            }
        }
    }

    return true;
}

// Return the weighted Olympic mean for the pixels
static float combinationWeightedOlympic(const psVector *values, // Values to combine
                                        const psVector *weights, // Weights to combine
                                        float frac, // Fraction to discard
                                        psVector *sortBuffer // Buffer for sorting
                                        )
{
    int numGood = values->n;            // Number of good values

    int numBad = frac * numGood + 0.5;  // Number of bad values
    int low = numBad / 2, high = low + numGood - numBad; // Indices (modulo masked pixels)

    sortBuffer = psVectorSortIndex(sortBuffer, values);

    double sumValues = 0.0, sumWeight = 0.0; // Sums for weighted mean
    for (int i = 0, j = 0; i < values->n; i++) {
        int index = sortBuffer->data.S32[i]; // Index of interest
        j++;
        if (j > high) {
            break;
        }
        if (j <= low) {
            continue;
        }
        sumValues += values->data.F32[index] * weights->data.F32[index];
        sumWeight += weights->data.F32[index];
    }

    return sumValues / sumWeight;
}

// Mark a pixel for inspection
// Value in pixel doesn't seem to agree with the stack, so need to look closer
static inline void combineMarkInspect(const psArray *inputs, // Stack data
                                      int x, int y, // Pixel
                                      int source // Source image index
                                      )
{
#ifdef TESTING
    if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
        fprintf(stderr, "Marking image %d, pixel %d,%d for inspection\n", source, x, y);
    }
#endif
    pmStackData *data = inputs->data[source]; // Stack data of interest
    if (!data) {
        psWarning("Can't find input data for source %d", source);
        return;
    }
    data->inspect = psPixelsAdd(data->inspect, data->inspect->nalloc, x, y);
    return;
}

// Mark a pixel for rejection
// Cannot possibly inspect this pixel and confirm that it's good.
// e.g., Only a single input
static inline void combineMarkReject(const psArray *inputs, // Stack data
                                     int x, int y, // Pixel
                                     int source // Source image index
                                     )
{
#ifdef TESTING
    if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
        fprintf(stderr, "Marking pixel image %d, pixel %d,%d for rejection\n", source, x, y);
    }
#endif
    pmStackData *data = inputs->data[source]; // Stack data of interest
    if (!data) {
        psWarning("Can't find input data for source %d", source);
        return;
    }
    data->reject = psPixelsAdd(data->reject, data->reject->nalloc, x, y);
    return;
}


// Extract vectors for simple combination/rejection operations
static void combineExtract(int *num,                        // Number of good pixels
                           bool *suspect,                   // Any suspect pixels?
                           combineBuffer *buffer, // Buffer with vectors
                           psImage *image, // Combined image, for output
                           psImage *mask, // Combined mask, for output
                           psImage *variance, // Combined variance map, for output
                           const psArray *inputs, // Stack data
                           const psVector *weights, // Global (single value) weights for data, or NULL
                           const psVector *exps,    // Exposures for data, or NULL
                           const psVector *addVariance, // Additional variance for data
                           const psVector *reject, // Indices of pixels to reject, or NULL
                           int x, int y, // Coordinates of interest; frame of output image
                           psImageMaskType maskVal, // Value to mask
                           psImageMaskType maskSuspect // Value to suspect
                           )
{
    // Rudimentary error checking
    assert(buffer);
    assert(image);
    assert(mask);
    assert(inputs);

    psVector *pixelData = buffer->pixels; // Values for the pixel of interest
    psVector *pixelVariances = variance ? buffer->variances : NULL; // Variances for the pixel of interest
    psVector *pixelWeights = buffer->weights; // Image weights for the pixel of interest
    psVector *pixelExps = buffer->exps;       // Exposure times
    psVector *pixelSources = buffer->sources; // Sources for the pixel of interest
    psVector *pixelLimits = buffer->limits; // Limits for the pixel of interest
    psVector *pixelSuspects = buffer->suspects; // Is the pixel suspect?

    if (suspect) {
        *suspect = false;
    }

    // Extract the pixel and mask data
    int numGood = 0;                    // Number of good pixels
    for (int i = 0, j = 0; i < inputs->n; i++) {
        // Check if this pixel has been rejected.  Assumes that the rejection pixel list is sorted --- it
        // should be because of how pixelMapGenerate works
        if (reject && reject->data.U16[j] == i) {
            j++;
            continue;
        }

        pmStackData *data = inputs->data[i]; // Stack data of interest
        if (!data) {
            continue;
        }

        int xIn = x - data->readout->col0, yIn = y - data->readout->row0; // Coordinates on input readout
        psImage *mask = data->readout->mask; // Mask of interest
        if (mask->data.PS_TYPE_IMAGE_MASK_DATA[yIn][xIn] & maskVal) {
            continue;
        }

        pixelSuspects->data.U8[numGood] = mask->data.PS_TYPE_IMAGE_MASK_DATA[yIn][xIn] & maskSuspect ?
            true : false;

        psImage *image = data->readout->image; // Image of interest
        psImage *variance = data->readout->variance; // Variance map of interest
        pixelData->data.F32[numGood] = image->data.F32[yIn][xIn];
        if (variance) {
            pixelVariances->data.F32[numGood] = variance->data.F32[yIn][xIn] * addVariance->data.F32[i];
        }
        pixelWeights->data.F32[numGood] = data->weight;
        pixelExps->data.F32[numGood] = data->exp;
        pixelSources->data.U16[numGood] = i;
        numGood++;
    }
    pixelData->n = numGood;
    if (variance) {
        pixelVariances->n = numGood;
    }
    pixelWeights->n = numGood;
    pixelSources->n = numGood;
    pixelLimits->n = numGood;
    pixelSuspects->n = numGood;
    *num = numGood;

#ifdef TESTING
    if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
        for (int i = 0; i < numGood; i++) {
            fprintf(stderr, "Input %d, pixel %d,%d (%" PRIu16 "): %f %f (%f) %f %f %d\n",
                    i, x, y, pixelSources->data.U16[i], pixelData->data.F32[i], pixelVariances->data.F32[i],
                    addVariance->data.F32[i], pixelWeights->data.F32[i], pixelExps->data.F32[i],
                    pixelSuspects->data.U8[i]);
        }
    }
#endif

    return;
}


// Combine pixels
static void combinePixels(psImage *image, // Combined image, for output
                          psImage *mask, // Combined mask, for output
                          psImage *variance, // Combined variance map, for output
                          psImage *exp,   // Exposure map (time), for output
                          psImage *expnum,       // Exposure map (number) for output
                          psImage *expweight,    // Exposure map (weighted time) for output
                          int num,      // Number of good pixels
                          combineBuffer *buffer, // Buffer with vectors
                          int x, int y, // Coordinates of interest; frame of output image
                          psImageMaskType bad, // Value for bad pixels
                          bool safe,           // Safe combination?
                          float invTotalWeight    // Inverse of total weight for all inputs
                          )
{
    psVector *pixelData = buffer->pixels; // Values for the pixel of interest
    psVector *pixelVariances = variance ? buffer->variances : NULL; // Variances for the pixel of interest
    psVector *pixelWeights = buffer->weights; // Image weights for the pixel of interest
    psVector *pixelExps = buffer->exps;       // Exposure times

    // Default option is that the pixel is bad
    float imageValue = NAN, varianceValue = NAN; // Value for combined image and variance map
    psImageMaskType maskValue = bad;    // Value for combined mask
    float expValue = 0.0, expWeightValue = NAN; // Exposure value (straight, and weighted)

    switch (num) {
      case 0: {
          // Nothing to combine: it's bad
#ifdef TESTING
          if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
              fprintf(stderr, "No inputs to combine, pixel %d,%d is bad.\n", x, y);
          }
#endif
          break;
      }
      case 1: {
          // Accept the single pixel unless we have to be safe
          if (!safe) {
              imageValue = pixelData->data.F32[0];
              if (variance) {
                  varianceValue = pixelVariances->data.F32[0];
              }
              if (exp) {
                  expValue = pixelExps->data.F32[0];
                  expWeightValue = pixelExps->data.F32[0];
              }
              maskValue = 0;
#ifdef TESTING
              if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                  fprintf(stderr, "Single input to combine, safety off, pixel %d,%d --> %f\n",
                          x, y, imageValue);
              }
#endif
          }
#ifdef TESTING
          else if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
              fprintf(stderr, "Single input to combine, safety on, pixel %d,%d is bad.\n", x, y);
          }
#endif
          break;
      }
      case 2: {
          // Automatically accept the mean of the pixels only if we're not playing safe
          if (!safe) {
              if (combinationMeanVariance(&imageValue, &varianceValue, &expValue, &expWeightValue,
                                          pixelData, pixelVariances, pixelExps, pixelWeights)) {
                  maskValue = 0;
#ifdef TESTING
                  if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                      fprintf(stderr, "Two inputs to combine using unsafe, pixel %d,%d --> %f %f\n",
                              x, y, imageValue, varianceValue);
                  }
#endif
              }
          }
#ifdef TESTING
          else {
              if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                  fprintf(stderr, "Two inputs to combine, safety on, pixel %d,%d is bad\n", x, y);
              }
          }
#endif
          break;
      }
      default: {
          // Can combine without too much worrying
          if (!combinationMeanVariance(&imageValue, &varianceValue, &expValue, &expWeightValue,
                                       pixelData, pixelVariances, pixelExps, pixelWeights)) {
              break;
          }
          maskValue = 0;
#ifdef TESTING
          if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
              fprintf(stderr, "Combined inputs, pixel %d,%d --> %f %f\n", x, y, imageValue, varianceValue);
          }
#endif
          break;
      }
    }

    image->data.F32[y][x] = imageValue;
    mask->data.PS_TYPE_IMAGE_MASK_DATA[y][x] = maskValue;
    if (variance) {
        variance->data.F32[y][x] = varianceValue;
    }
    if (exp) {
        exp->data.F32[y][x] = expValue;
    }
    if (expnum) {
        expnum->data.PS_TYPE_IMAGE_MASK_DATA[y][x] = num;
    }
    if (expweight) {
        expweight->data.F32[y][x] = expWeightValue * invTotalWeight;
    }

    return;
}


// Test pixels to be combined
// Returns false to repeat without suspect pixels
static bool combineTest(int num,      // Number of good pixels
                        bool suspect, // Does the stack contain suspect pixels?
                        psArray *inputs,       // Original inputs (for flagging)
                        combineBuffer *buffer, // Buffer with vectors
                        int x, int y, // Coordinates of interest; frame of output image
                        float iter, // Number of rejection iterations per input
                        float rej, // Number of standard deviations at which to reject
                        float sys,    // Relative systematic error
                        float olympic,// Fraction of values to discard (Olympic weighted mean)
                        bool useVariance, // Use variance for rejection when combining?
                        bool safe    // Combine safely?
                        )
{
    if (iter <= 0) {
        return true;
    }

    int numIter = PS_MAX(iter * num, 1); // Number of iterations

#ifdef TESTING
    if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
        fprintf(stderr, "Testing pixel %d,%d: %d %f %f %f %d %d\n",
                x, y, numIter, rej, sys, olympic, useVariance, safe);
    }
#endif

    psVector *pixelData = buffer->pixels; // Values for the pixel of interest
    psVector *pixelWeights = buffer->weights; // Is the pixel suspect?
    psVector *pixelVariances = buffer->variances; // Variances for the pixel of interest
    psVector *pixelSources = buffer->sources; // Sources for the pixel of interest
    psVector *pixelSuspects = buffer->suspects; // Is the pixel suspect?
    psVector *pixelLimits = buffer->limits; // Is the pixel suspect?

    // Set up rejection limits
    float rej2 = PS_SQR(rej); // Rejection level squared
    if (num > 2 && useVariance) {
        // Convert rejection limits --- saves doing it later multiple times
        // Using squared rejection limit because it's cheaper than sqrts
        double sumWeights = 0.0;
        for (int i = 0; i < num; i++) {
            sumWeights += pixelWeights->data.F32[i];
        }
        for (int i = 0; i < num; i++) {
            // Systematic error contributes to the rejection level
            float sysVar = PS_SQR(sys * pixelData->data.F32[i]);
#ifdef TESTING
            // Correct variance for comparison against weighted mean including itself
            float compare = 1.0 - pixelWeights->data.F32[i] / sumWeights;
            if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                fprintf(stderr, "Variance %d (%d), pixel %d,%d: %f %f %f\n", i, pixelSources->data.U16[i],
                        x, y, pixelVariances->data.F32[i], sysVar, compare);
            }
#endif
            pixelLimits->data.F32[i] = rej2 * (pixelVariances->data.F32[i] + sysVar);
        }
    }

    int maskIndex = 0;                  // Index of pixel to mask
    int totalClipped = 0;               // Total number of pixels clipped
    for (int i = 0; i < numIter && maskIndex >= 0; i++) {
        maskIndex = -1;                 // Nothing to reject

        switch (num) {
          case 0:
            break;
          case 1:
            if (i == 0 && safe) {
                combineMarkReject(inputs, x, y, pixelSources->data.U16[0]);
            }
            break;
          case 2: {
              if (useVariance) {
                  // Use variance to check that the two are consistent
                  float diff = 0.5 * (pixelData->data.F32[0] - pixelData->data.F32[1]); // Mean flux difference
                  float var1 = pixelVariances->data.F32[0]; // Variance of first
                  float var2 = pixelVariances->data.F32[1]; // Variance of second
                  // Systematic error contributes to the rejection level
                  var1 += PS_SQR(sys * pixelData->data.F32[0]);
                  var2 += PS_SQR(sys * pixelData->data.F32[1]);

                  float sigma2 = var1 + var2; // Combined variance
                  if (PS_SQR(diff) > rej2 * sigma2) {
                      // Not consistent: don't believe either!
                      if (i == 0 && suspect) {
                          combineMarkReject(inputs, x, y, pixelSources->data.U16[0]);
                          combineMarkReject(inputs, x, y, pixelSources->data.U16[1]);
                      } else {
                          combineMarkInspect(inputs, x, y, pixelSources->data.U16[0]);
                          combineMarkInspect(inputs, x, y, pixelSources->data.U16[1]);
                      }
#ifdef TESTING
                      if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                          fprintf(stderr, "Flagged both inputs for pixel %d,%d (%f > %f x %f\n)",
                                  x, y, diff, rej, sqrtf(sigma2));
                      }
#endif
                  }
              } else if (i == 0 && safe) {
                  // Can't test them, and we want to be safe, so reject
                  combineMarkReject(inputs, x, y, pixelSources->data.U16[0]);
                  combineMarkReject(inputs, x, y, pixelSources->data.U16[1]);
              }
              break;
          }
#if 0
          case 3: {
              // Want to be a bit careful on the rejection than for a larger number of inputs
              if (!useVariance) {
                  return combineTestGeneral(num, suspect, inputs, buffer, x, y, numIter, rej, sys,
                                            olympic, useVariance, safe, allowSuspect);
              }

              // Differences between pixel values
              float diff01 = pixelData->data.F32[0] - pixelData->data.F32[1];
              float diff12 = pixelData->data.F32[1] - pixelData->data.F32[2];
              float diff20 = pixelData->data.F32[2] - pixelData->data.F32[0];
              // Variance for each pixel
              float var0 = pixelVariances->data.F32[0] + PS_SQR(sys * pixelData->data.F32[0]);
              float var1 = pixelVariances->data.F32[1] + PS_SQR(sys * pixelData->data.F32[1]);
              float var2 = pixelVariances->data.F32[2] + PS_SQR(sys * pixelData->data.F32[2]);
              // Errors in pixel differences
              float err01 = var0 + var1;
              float err12 = var1 + var2;
              float err20 = var2 + var0;

#ifdef TESTING
              if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                  fprintf(stderr, "Diff 0-1: %f %f\n", diff01, err01);
                  fprintf(stderr, "Diff 1-2: %f %f\n", diff12, err12);
                  fprintf(stderr, "Diff 2-0: %f %f\n", diff20, err20);
              }
#endif

              int badPairs = 0;         // Number of bad pairs
              bool bad01 = false, bad12 = false, bad20 = false; // Pair is bad?
              if (PS_SQR(diff01) > rej2 * err01) {
                  bad01 = true;
                  badPairs++;
              }
              if (PS_SQR(diff12) > rej2 * err12) {
                  bad12 = true;
                  badPairs++;
              }
              if (PS_SQR(diff20) > rej2 * err20) {
                  bad20 = true;
                  badPairs++;
              }

              if (badPairs > 0 && allowSuspect && suspect) {
                  return false;
              }

              switch (badPairs) {
                case 0:
                  // Nothing to worry about!
                  break;
                case 1:
                  // Can't tell which image is bad, so be sure to get it
                  if (bad01) {
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[0]);
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[1]);
                      break;
                  }
                  if (bad12) {
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[1]);
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[2]);
                      break;
                  }
                  if (bad20) {
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[2]);
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[0]);
                      break;
                  }
                  psAbort("Should never get here");
                case 2:
                  if (bad01 && bad12) {
                      // 2 and 0 are good
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[1]);
                      break;
                  }
                  if (bad12 && bad20) {
                      // 0 and 1 are good
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[2]);
                      break;
                  }
                  if (bad20 && bad01) {
                      // 1 and 2 are good
                      combineMarkInspect(inputs, x, y, pixelSources->data.U16[0]);
                      break;
                  }
                  psAbort("Should never get here");
                case 3:
                  // Everything's bad
                  combineMarkInspect(inputs, x, y, pixelSources->data.U16[0]);
                  combineMarkInspect(inputs, x, y, pixelSources->data.U16[1]);
                  combineMarkInspect(inputs, x, y, pixelSources->data.U16[2]);
                  break;
              }
              break;
          }
#endif
          default: {
              if (useVariance) {
                  float median = combinationWeightedOlympic(pixelData, pixelWeights,
                                                            olympic, buffer->sort); // Median for stack
#ifdef TESTING
                  if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                      fprintf(stderr, "Flag with variance pixel %d,%d: median = %f\n", x, y, median);
                  }
#endif
                  float worst = -INFINITY; // Largest deviation
                  for (int j = 0; j < num; j++) {
                      float diff = pixelData->data.F32[j] - median; // Difference from expected
#ifdef TESTING
                      if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                          fprintf(stderr, "Testing input %d for pixel %d,%d: %f\n", j, x, y, diff);
                      }
#endif

                      // Comparing squares --- cheaper than lots of sqrts
                      // pixelVariances includes the rejection limit, from above
                      float diff2 = PS_SQR(diff); // Square difference
                      if (diff2 > pixelLimits->data.F32[j]) {
                          float dev = diff2 / pixelLimits->data.F32[j]; // Deviation
                          if (dev > worst) {
                              worst = dev;
                              maskIndex = j;
                          }
                      }
                  }
              } else {
                  float median = NAN, stdev = NAN;  // Median and stdev of the combination, for rejection
                  combinationMedianStdev(&median, &stdev, pixelData, buffer->sort);
                  float limit = rej * stdev; // Rejection limit
#ifdef TESTING
                  if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                      fprintf(stderr,
                              "Flag without variance pixel %d,%d; median = %f, stdev = %f, limit = %f\n",
                              x, y, median, stdev, limit);
                  }
#endif
                  float worst = -INFINITY; // Largest deviation
                  for (int j = 0; j < num; j++) {
                      float diff = fabsf(pixelData->data.F32[j] - median); // Difference from expected

                      if (diff > limit) {
                          float dev = diff / limit; // Deviation
                          if (dev > worst) {
                              worst = dev;
                              maskIndex = j;
                          }
                      }
                  }
              }
          }
        }

        // Do the actual rejection of the pixel
        if (maskIndex >= 0) {
            if (suspect) {
#ifdef TESTING
                if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                    fprintf(stderr, "Throwing out all suspect pixels for %d,%d\n", x, y);
                }
#endif
                // Throw out all suspect pixels
                int numGood = 0;        // Number of good pixels
                for (int j = 0; j < num; j++) {
                    if (pixelSuspects->data.U8[j]) {
                        combineMarkReject(inputs, x, y, pixelSources->data.U16[j]);
                        continue;
                    }
                    if (numGood == j) {
                        numGood++;
                        continue;
                    }
                    pixelData->data.F32[numGood] = pixelData->data.F32[j];
                    pixelWeights->data.F32[numGood] = pixelWeights->data.F32[j];
                    pixelSources->data.U16[numGood] = pixelSources->data.U16[j];
                    pixelLimits->data.F32[numGood] = pixelLimits->data.F32[j];
                    pixelVariances->data.F32[numGood] = pixelVariances->data.F32[j];
                    numGood++;
                }
                pixelData->n = numGood;
                pixelWeights->n = numGood;
                pixelSources->n = numGood;
                pixelLimits->n = numGood;
                pixelVariances->n = numGood;
                totalClipped += num - numGood;
                num = numGood;
                suspect = false;
            } else {
                // Throw out masked pixel
#ifdef TESTING
                if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                    fprintf(stderr, "Throwing out input %d for pixel %d,%d\n", maskIndex, x, y);
                }
#endif
                combineMarkInspect(inputs, x, y, pixelSources->data.U16[maskIndex]);
                int numGood = 0;        // Number of good pixels
                for (int j = 0; j < num; j++) {
                    if (j == maskIndex) {
                        continue;
                    }
                    if (numGood == j) {
                        numGood++;
                        continue;
                    }
                    pixelData->data.F32[numGood] = pixelData->data.F32[j];
                    pixelWeights->data.F32[numGood] = pixelWeights->data.F32[j];
                    pixelSources->data.U16[numGood] = pixelSources->data.U16[j];
                    pixelLimits->data.F32[numGood] = pixelLimits->data.F32[j];
                    pixelVariances->data.F32[numGood] = pixelVariances->data.F32[j];
                    numGood++;
                }
                pixelData->n = numGood;
                pixelWeights->n = numGood;
                pixelSources->n = numGood;
                pixelLimits->n = numGood;
                pixelVariances->n = numGood;
                totalClipped++;
                num--;
            }
        }
    }

    return true;
}


// Ensure the input array of pmStackData is valid, and get some details out of it
static bool validateInputData(bool *haveVariances, // Do we have variance maps in the sky images?
                              int *num,    // Number of inputs
                              int *numCols, int *numRows, // Size of (sky) images
                              const psArray *input, // Input array of pmStackData to validate
                              const pmReadout *output, // Output readout
                              const pmReadout *exp    // Exposure map
    )
{
    PS_ASSERT_ARRAY_NON_NULL(input, false);
    *num = input->n;

    pmStackData *data = NULL;           // First image off the rank, used as a template
    for (int i = 0; !data && i < input->n; i++) {
        data = input->data[i];
    }
    PS_ASSERT_PTR_NON_NULL(data, false);
    assert(psMemGetDeallocator(data) == (psFreeFunc)stackDataFree); // Ensure it's the right type
    *haveVariances = false;
    PS_ASSERT_IMAGE_NON_NULL(data->readout->image, false);
    PS_ASSERT_IMAGE_TYPE(data->readout->image, PS_TYPE_F32, false);
    PS_ASSERT_IMAGE_NON_NULL(data->readout->mask, false);
    PS_ASSERT_IMAGE_TYPE(data->readout->mask, PS_TYPE_IMAGE_MASK, false);
    PS_ASSERT_IMAGES_SIZE_EQUAL(data->readout->image, data->readout->mask, false);
    *numCols = data->readout->image->numCols;
    *numRows = data->readout->image->numRows;
    if (data->readout->variance) {
        *haveVariances = true;
        PS_ASSERT_IMAGE_NON_NULL(data->readout->variance, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(data->readout->image, data->readout->variance, false);
        PS_ASSERT_IMAGE_TYPE(data->readout->variance, PS_TYPE_F32, false);
    }
    bool haveRejects = (data->reject != NULL); // Do we have rejected pixels?

    // Make sure the rest correspond with the first
    for (int i = 1; i < *num; i++) {
        pmStackData *data = input->data[i]; // Stack data for this input
        if (!data) {
            continue;
        }
        assert(psMemGetDeallocator(data) == (psFreeFunc)stackDataFree); // Ensure it's the right type
        if (!data->readout) {
            psError(PS_ERR_UNEXPECTED_NULL, true, "The readout is specified in some but not all inputs.");
            return false;
        }
        if ((haveRejects && !data->reject) || (data->reject && !haveRejects)) {
            psError(PS_ERR_UNEXPECTED_NULL, true,
                    "The rejected pixels are specified in some but not all inputs.");
            return false;
        }
        PS_ASSERT_IMAGE_NON_NULL(data->readout->image, false);
        PS_ASSERT_IMAGE_NON_NULL(data->readout->mask, false);
        PS_ASSERT_IMAGE_TYPE(data->readout->image, PS_TYPE_F32, false);
        PS_ASSERT_IMAGE_TYPE(data->readout->mask, PS_TYPE_IMAGE_MASK, false);
        PS_ASSERT_IMAGE_SIZE(data->readout->image, *numCols, *numRows, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(data->readout->image, data->readout->mask, false);
        if (*haveVariances) {
            PS_ASSERT_IMAGE_NON_NULL(data->readout->variance, false);
            PS_ASSERT_IMAGES_SIZE_EQUAL(data->readout->image, data->readout->variance, false);
            PS_ASSERT_IMAGE_TYPE(data->readout->variance, PS_TYPE_F32, false);
        }
    }

    PM_ASSERT_READOUT_NON_NULL(output, false);
    if (output->image) {
        PS_ASSERT_IMAGE_NON_NULL(output->image, false);
        PS_ASSERT_IMAGE_TYPE(output->image, PS_TYPE_F32, false);
        PS_ASSERT_IMAGE_NON_NULL(output->mask, false);
        PS_ASSERT_IMAGE_TYPE(output->mask, PS_TYPE_IMAGE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(output->image, output->mask, false);
    }

    if (exp) {
        PM_ASSERT_READOUT_NON_NULL(exp, false);
        if (exp->image) {
            PS_ASSERT_IMAGES_SIZE_EQUAL(exp->image, output->image, false);
        }
        if (exp->mask) {
            PS_ASSERT_IMAGES_SIZE_EQUAL(exp->mask, output->image, false);
        }
    }

    return true;
}


// Generate a "pixel map".
//
// A "pixel map" is an image-like structure containing a vector that contains the indices of images.  The idea
// is to provide a reverse lookup for an array of pixel lists, so that the image for which a pixel is flagged
// can be identified easily.
static psArray *pixelMapGenerate(const psArray *input, // Data to stack
                                 int minCols, int maxCols, int minRows, int maxRows // Bounds of interest
    )
{
    int numCols = maxCols - minCols + 1, numRows = maxRows - minRows + 1; // Size of map

    psArray *map = psArrayAlloc(numRows); // The pixel map
    for (int y = 0; y < numRows; y++) {
        map->data[y] = psArrayAlloc(numCols);
    }

    for (int i = 0; i < input->n; i++) {
        pmStackData *data = input->data[i];
        if (!data) {
            continue;
        }
        assert(data->reject);
        psPixels *pixels = data->reject; // The rejected pixels
        for (int j = 0; j < pixels->n; j++) {
            int x = pixels->data[j].x - minCols, y = pixels->data[j].y - minRows; // Coordinates of interest
            if (x < 0 || x >= numCols || y < 0 || y >= numRows) {
                continue;
            }
            psArray *columns = map->data[y]; // The columns for that row
            psVector *images = columns->data[x]; // The images for that column
            if (!images) {
                images = columns->data[x] = psVectorAllocEmpty(PIXEL_MAP_BUFFER, PS_TYPE_U16);
            }
            int size = images->n;       // Element number at which to add
            columns->data[x] = psVectorExtend(images, PIXEL_MAP_BUFFER, 1);
            images->data.U16[size] = i;
        }
    }

    return map;
}

// Query a "pixel map", by returning the list of image indices for a particular pixel.
static psVector *pixelMapQuery(const psArray *map, // Pixel map
                               int x0, int y0, // Offset into map
                               int x, int y // Coordinates of interest
    )
{
    // Adjust for offset
    x -= x0;
    y -= y0;

    assert(y >= 0 && y < map->n);
    psArray *colMap = map->data[y];     // Columns for that row
    assert(x >= 0 && x < colMap->n);
    return colMap->data[x];
}


//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

/// Constructor
pmStackData *pmStackDataAlloc(pmReadout *readout, float weight, float exp, float addVariance)
{
    pmStackData *data = psAlloc(sizeof(pmStackData)); // Stack data, to return
    psMemSetDeallocator(data, (psFreeFunc)stackDataFree);

    data->readout = psMemIncrRefCounter(readout);
    data->reject = NULL;
    data->inspect = NULL;
    data->weight = weight;
    data->exp = exp;
    data->addVariance = addVariance;

    return data;
}

/// Stack input images
bool pmStackCombine(pmReadout *combined, pmReadout *expmaps, psArray *input,
                    psImageMaskType maskVal, psImageMaskType maskSuspect,
                    psImageMaskType bad, int kernelSize,
                    float iter, float rej, float sys, float olympic,
                    bool useVariance, bool safe, bool rejection)
{
    bool haveVariances;                 // Do we have the variance maps?
    int num;                            // Number of inputs
    int numCols, numRows;               // Size of (sky) images
    if (!validateInputData(&haveVariances, &num, &numCols, &numRows, input, combined, expmaps)) {
        return false;
    }
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);
    PS_ASSERT_INT_POSITIVE(bad, false);
    if (isnan(rej)) {
        PS_ASSERT_FLOAT_EQUAL(iter, 0, false);
    } else {
        PS_ASSERT_FLOAT_LARGER_THAN(iter, 0, false);
        PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, false);
    }
    if (useVariance && !haveVariances) {
        psWarning("Unable to use variance in rejection if no variance maps supplied --- option turned off");
        useVariance = false;
    }

    psVector *addVariance = psVectorAlloc(num, PS_TYPE_F32); // Additional variance for each image
    psVector *weights = psVectorAlloc(num, PS_TYPE_F32); // Relative weighting for each image
    psVector *exps = psVectorAlloc(num, PS_TYPE_F32);    // Exposure times for each image
    psArray *stack = psArrayAlloc(num); // Stack of readouts
    float totalExpWeight = 0.0;           // Total value of all weighted exposure times
    float totalExp = 0.0;                 // Total exposure time
    for (int i = 0; i < num; i++) {
        pmStackData *data = input->data[i]; // Stack data for this input
        if (!data) {
            weights->data.F32[i] = 0.0;
            exps->data.F32[i] = NAN;
            continue;
        }
        weights->data.F32[i] = data->weight;
        exps->data.F32[i] = data->exp;
        totalExp += exps->data.F32[i];
        totalExpWeight += exps->data.F32[i] * weights->data.F32[i];
        pmReadout *ro = data->readout;  // Readout of interest
        stack->data[i] = psMemIncrRefCounter(ro);
        addVariance->data.F32[i] = ro->covariance ? psImageCovarianceFactor(ro->covariance) : 1.0;
#ifdef ADD_VARIANCE
        if (isfinite(data->addVariance)) {
            addVariance->data.F32[i] *= data->addVariance;
        }
#endif
        if (!rejection) {
            // Ensure pixels can be put on the appropriate list
            if (!data->inspect) {
                data->inspect = psPixelsAllocEmpty(PIXEL_LIST_BUFFER);
            }
            if (!data->reject) {
                data->reject = psPixelsAllocEmpty(PIXEL_LIST_BUFFER);
            }
        }
    }
    totalExpWeight = totalExp / totalExpWeight;    // Convert to inverse

    int minInputCols, maxInputCols, minInputRows, maxInputRows; // Smallest and largest values to combine
    int xSize, ySize;                   // Size of the output image
    if (!pmReadoutStackValidate(&minInputCols, &maxInputCols, &minInputRows, &maxInputRows, &xSize, &ySize,
                                stack)) {
        psError(psErrorCodeLast(), false, "Input stack is not valid.");
        psFree(stack);
        return false;
    }
    psFree(stack);
    pmReadoutUpdateSize(combined, minInputCols, minInputRows, xSize, ySize, true, true, bad);
    psTrace("psModules.imcombine", 1, "Have for combination [%d:%d,%d:%d] (%dx%d)\n",
            minInputCols, maxInputCols, minInputRows, maxInputRows, xSize, ySize);

    // Reduce combination area by the size of the kernel
    minInputCols += kernelSize;
    maxInputCols -= kernelSize;
    minInputRows += kernelSize;
    maxInputRows -= kernelSize;
    psTrace("psModules.imcombine", 1, "Combining on [%d:%d,%d:%d]\n",
            minInputCols, maxInputCols, minInputRows, maxInputRows);


    // Buffer for combination
    combineBuffer *buffer = combineBufferAlloc(num);

    // Pull the products out, allocate if necessary
    psImage *combinedImage = combined->image; // Combined image
    if (!combinedImage) {
        combined->image = psImageAlloc(numCols, numRows, PS_TYPE_F32);
        combinedImage = combined->image;
    }
    psImage *combinedMask = combined->mask; // Combined mask
    if (!combinedMask) {
        combined->mask = psImageAlloc(numCols, numRows, PS_TYPE_IMAGE_MASK);
        combinedMask = combined->mask;
    }

    psImage *combinedVariance = combined->variance; // Combined variance map
    if (haveVariances && !combinedVariance) {
        combined->variance = psImageAlloc(numCols, numRows, PS_TYPE_F32);
        combinedVariance = combined->variance;
    }

    psImage *exp = NULL, *expnum = NULL, *expweight = NULL; // Exposure map and exposure number
    if (expmaps) {
        if (!expmaps->image) {
            expmaps->image = psImageAlloc(numCols, numRows, PS_TYPE_F32);
        }
        exp = expmaps->image;

        if (!expmaps->mask) {
            expmaps->mask = psImageAlloc(numCols, numRows, PS_TYPE_IMAGE_MASK);
        }
        expnum = expmaps->mask;

        if (!expmaps->variance) {
            expmaps->variance = psImageAlloc(numCols, numRows, PS_TYPE_F32);
        }
        expweight = expmaps->variance;
    }

    // Set up rejection list
    psArray *pixelMap = NULL;           // Map of pixels to source
    if (rejection) {
        pixelMap = pixelMapGenerate(input, minInputCols, maxInputCols, minInputRows, maxInputRows);
    }

    // Combine each pixel
    for (int y = minInputRows; y < maxInputRows; y++) {
        for (int x = minInputCols; x < maxInputCols; x++) {
#ifdef TESTING
            if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                fprintf(stderr, "Combining pixel %d,%d: %x %x %f %f %f %f %d %d %d\n",
                        x, y, maskVal, bad, iter, rej, sys, olympic, useVariance, safe, rejection);
            }
#endif
            psVector *reject = NULL; // Images to reject for this pixel
            if (rejection) {
                reject = pixelMapQuery(pixelMap, minInputCols, minInputRows, x, y);
#ifdef TESTING
                if (PS_SQR(x - TEST_X) + PS_SQR(y - TEST_Y) <= PS_SQR(TEST_RADIUS)) {
                    fprintf(stderr, "Rejected inputs for pixel %d,%d: ", x, y);
                    if (!reject) {
                        fprintf(stderr, "<none>\n");
                    } else {
                        for (int i = 0; i < reject->n; i++) {
                            fprintf(stderr, "%d ", reject->data.U16[i]);
                        }
                        fprintf(stderr, "\n");
                    }
                }
#endif
            }

            int num;                    // Number of good pixels
            bool suspect;               // Suspect pixels in stack?
            combineExtract(&num, &suspect, buffer, combinedImage, combinedMask, combinedVariance,
                           input, weights, exps, addVariance, reject, x, y, maskVal, maskSuspect);
            combinePixels(combinedImage, combinedMask, combinedVariance, exp, expnum, expweight,
                          num, buffer, x, y, bad, safe, totalExpWeight);

            if (iter > 0) {
                combineTest(num, suspect, input, buffer, x, y, iter, rej, sys, olympic,
                            useVariance, safe);
            }
        }
    }

    psFree(pixelMap);
    psFree(weights);
    psFree(buffer);
    psFree(addVariance);


#ifndef PS_NO_TRACE
    if (!rejection && psTraceGetLevel("psModules.imcombine") >= 5) {
        for (int i = 0; i < num; i++) {
            pmStackData *data = input->data[i]; // Stack data for this input
            if (!data || !data->inspect) {
                continue;
            }
            psTrace("psModules.imcombine", 5, "Image %d: %ld pixels to inspect.\n", i, data->inspect->n);
        }
    }
#endif

    return true;
}

