/** @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 1085                     // x coordinate to examine
//#define TEST_Y 3371                     // y coordinate 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 *masks;                    // Pixel masks
    psVector *variances;                // Pixel variances
    psVector *weights;                  // Pixel weightings
    psVector *sources;                  // Pixel sources (which image did they come from?)
    psVector *limits;                   // Rejection limits
    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->masks);
    psFree(buffer->variances);
    psFree(buffer->weights);
    psFree(buffer->sources);
    psFree(buffer->limits);
    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->masks = psVectorAlloc(numImages, PS_TYPE_VECTOR_MASK);
    buffer->variances = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->weights = psVectorAlloc(numImages, PS_TYPE_F32);
    buffer->sources = psVectorAlloc(numImages, PS_TYPE_U16);
    buffer->limits = psVectorAlloc(numImages, PS_TYPE_F32);
    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
                                    const psVector *values, // Values to combine
                                    const psVector *variances, // Pixel variances 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
    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 (sumWeight <= 0) {
        return false;
    }

    *mean = sumValueWeight / sumWeight;
    if (var) {
        *var = sumVarianceWeight / PS_SQR(sumWeight);
    }
    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
                                   const psVector *masks, // Mask to apply
                                   psVector *sortBuffer // Buffer for sorting
                                   )
{
    assert(values);
    assert(!masks || values->n == masks->n);
    assert(values->type.type == PS_TYPE_F32);
    assert(!masks || masks->type.type == PS_TYPE_VECTOR_MASK);
    assert(sortBuffer && sortBuffer->nalloc >= values->n && sortBuffer->type.type == PS_TYPE_F32);

    // Need to filter out clipped values
    int num = 0;            // Number of valid values
    for (int i = 0; i < values->n; i++) {
        if (!masks || !masks->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
            sortBuffer->data.F32[num++] = values->data.F32[i];
        }
    }
    sortBuffer->n = num;
    if (!psVectorSortInPlace(sortBuffer)) {
        return false;
    }

    if (num == 3) {
        // Attempt to measure standard deviation with only three values (and one of those possibly corrupted)
        *median = sortBuffer->data.F32[1];
        if (stdev) {
            float diff1 = sortBuffer->data.F32[0] - *median;
            float diff2 = sortBuffer->data.F32[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 ? sortBuffer->data.F32[num / 2] :
            (sortBuffer->data.F32[num / 2 - 1] + sortBuffer->data.F32[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(sortBuffer->data.F32[i] - *median);
                }
                *stdev = sqrt(sum / (double)(num - 1));
            } else {
                // Standard deviation from the interquartile range
                *stdev = 0.74 * (sortBuffer->data.F32[(int)(0.75 * num)] -
                                 sortBuffer->data.F32[(int)(0.25 * num)]);
            }
        }
    }

    return true;
}

#if 0
// Return the weighted median for the pixels
// This does not appear to produce as clean images as the weighted Olympic mean
static float combinationWeightedMedian(const psVector *values, // Values to combine
                                       const psVector *weights, // Weights to combine
                                       const psVector *masks, // Mask to apply
                                       psVector *sortBuffer // Buffer for sorting
                                       )
{
    double sumWeight = 0.0;             // Sum of weights
    for (int j = 0; j < values->n; j++) {
        if (masks->data.PS_TYPE_VECTOR_MASK_DATA[j]) {
            continue;
        }
        sumWeight += weights->data.F32[j];
    }

    sortBuffer = psVectorSortIndex(sortBuffer, values);
    double target = sumWeight / 2.0;    // Target weight

    int dominant = -1;                  // Index of dominant value, if any
    double cumulativeWeight = 0.0;      // Sum of weights
    for (int j = 0; j < values->n; j++) {
        if (masks->data.PS_TYPE_VECTOR_MASK_DATA[j]) {
            continue;
        }
        int index = sortBuffer->data.S32[j];  // Index of value of interest
        float weight = weights->data.F32[index]; // Weight for value of interest
        if (weight >= target) {
            // Get the weighted median of the rest
            dominant = index;
            sumWeight -= weight;
            target = sumWeight / 2.0;
            continue;
        }
        cumulativeWeight += weight;
        if (cumulativeWeight >= target) {
            float median = values->data.F32[index]; // Weighted median median
            if (dominant != -1) {
                // In the case that a single value contains a disproportionate weight compared to the rest,
                // we use a weighted mean between that dominant value and the weighted median of the rest.
                return (values->data.F32[dominant] * weights->data.F32[dominant] + median * sumWeight) /
                    (weights->data.F32[dominant] + sumWeight);
            }
            return median;
        }
    }

    return NAN;
}
#endif

// Return the weighted Olympic mean for the pixels
static float combinationWeightedOlympic(const psVector *values, // Values to combine
                                        const psVector *weights, // Weights to combine
                                        const psVector *masks, // Mask to apply
                                        float frac, // Fraction to discard
                                        psVector *sortBuffer // Buffer for sorting
                                        )
{
    int numGood = 0;                    // Number of good values
    for (int i = 0; i < values->n; i++) {
        if (masks->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
            continue;
        }
        numGood++;
    }

    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
        if (masks->data.PS_TYPE_VECTOR_MASK_DATA[index]) {
            continue;
        }
        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
static inline void combineInspect(const psArray *inputs, // Stack data
                                  int x, int y, // Pixel
                                  int source // Source image index
                                  )
{
    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;
}

// Given a stack of images, combine with optional rejection.
// Pixels in the stack that are rejected are marked for subsequent inspection
static void combinePixels(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 *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 bad, // Value to give bad pixels
                          int numIter, // Number of rejection iterations
                          float rej, // Number of standard deviations at which to reject
                          float sys,    // Relative systematic error
                          float discard,// Fraction of values to discard (Olympic weighted mean)
                          bool useVariance, // Use variance for rejection when combining?
                          bool safe,    // Combine safely?
                          bool rejectInspect, // Reject values marked for inspection from combination?
                          combineBuffer *buffer // Buffer for combination; to avoid multiple allocations
                         )
{
    // Rudimentary error checking
    assert(image);
    assert(mask);
    assert(inputs);
    assert(numIter >= 0);
    assert(buffer);
    assert(addVariance);
    assert((useVariance && variance) || !useVariance);

    psVector *pixelData = buffer->pixels; // Values for the pixel of interest
    psVector *pixelMasks = buffer->masks; // Masks 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 *pixelSources = buffer->sources; // Sources for the pixel of interest
    psVector *pixelLimits = buffer->limits; // Limits for the pixel of interest
    psVector *sort = buffer->sort;      // Sort buffer

    // Extract the pixel and mask data
    int num = 0;                        // Number of good images
    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;
        }

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

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

    // The sensible thing to do varies according to how many good pixels there are.
    // 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
    switch (num) {
      case 0:
        // Nothing to combine: it's bad
#ifdef TESTING
    if (x == TEST_X && y == TEST_Y) {
        fprintf(stderr, "No inputs to combine, pixel is bad.\n");
    }
#endif
        break;
      case 1: {
          // Accept the single pixel unless we have to be safe
          if (!safe) {
#ifdef TESTING
              if (x == TEST_X && y == TEST_Y) {
                  fprintf(stderr, "Single input to combine, safety off.\n");
              }
#endif
              imageValue = pixelData->data.F32[0];
              if (variance) {
                  varianceValue = pixelVariances->data.F32[0];
              }
              maskValue = 0;
          }
#ifdef TESTING
          else if (x == TEST_X && y == TEST_Y) {
              fprintf(stderr, "Single input to combine, safety on, pixel is bad.\n");
          }
#endif
          break;
      }
      case 2: {
          // Accept the mean of the pixels only if we're going to reject based on the variance, or we're not
          // playing safe
          if (useVariance || !safe) {
              float mean, var;   // Mean and variance from combination
              if (combinationMeanVariance(&mean, &var, pixelData, pixelVariances, pixelWeights)) {
                  imageValue = mean;
                  varianceValue = var;
                  maskValue = 0;
#ifdef TESTING
                  if (x == TEST_X && y == TEST_Y) {
                      fprintf(stderr, "Two inputs to combine using variance/unsafe --> %f %f\n",
                              mean, var);
                  }
#endif
              }
          }
          if (useVariance && numIter > 0) {
              // 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) > PS_SQR(rej) * sigma2) {
                  // Not consistent: mark both for inspection
                  if (rejectInspect) {
                      imageValue = NAN;
                      varianceValue = NAN;
                      maskValue = bad;
                  } else {
                      combineInspect(inputs, x, y, pixelSources->data.U16[0]);
                      combineInspect(inputs, x, y, pixelSources->data.U16[1]);
                  }
#ifdef TESTING
                  if (x == TEST_X && y == TEST_Y) {
                      fprintf(stderr, "Both pixels marked for inspection (%f > %f x %f\n)",
                              diff, rej, sqrtf(sigma2));
                  }
#endif
              }
          }
          break;
      }
      default: {
          // Record the value derived with no clipping, because pixels rejected using the harsh clipping
          // applied in the first pass might later be accepted.
          float mean, var;           // Mean and variance of the combination
          if (!combinationMeanVariance(&mean, &var, pixelData, pixelVariances, pixelWeights)) {
              break;
          }
          imageValue = mean;
          varianceValue = var;
          maskValue = 0;
#ifdef TESTING
          if (x == TEST_X && y == TEST_Y) {
              fprintf(stderr, "Combined inputs: %f %f\n", mean, var);
          }
#endif

          // Prepare for clipping iteration
          if (numIter > 0) {
              pixelMasks->n = num;
              psVectorInit(pixelMasks, 0);
              if (useVariance) {
                  // Convert to rejection limits --- saves doing it later.
                  // Using squared rejection limit because it's cheaper than sqrts
                  float rej2 = PS_SQR(rej); // Rejection level squared
                  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 (x == TEST_X && y == TEST_Y) {
                          fprintf(stderr, "Variance %d (%d): %f %f %f\n", i, pixelSources->data.U16[i],
                                  pixelVariances->data.F32[i], sysVar, compare);
                      }
#endif

                      pixelLimits->data.F32[i] = rej2 * (pixelVariances->data.F32[i] + sysVar);
                  }
              }
          }

          // The clipping that follows is solely to identify suspect pixels.
          // These suspect pixels will be inspected in more detail by other functions.
          int numClipped = INT_MAX;     // Number of pixels clipped per iteration
          int totalClipped = 0;         // Total number of pixels clipped
          for (int i = 0; i < numIter && numClipped > 0 && num - totalClipped > 2; i++) {
              numClipped = 0;
              float median = NAN;       // Middle of distribution
              float limit = NAN;        // Rejection limit
              if (!useVariance) {
                  float stdev;  // Median and stdev of the combination, for rejection
                  if (!combinationMedianStdev(&median, useVariance ? NULL : &stdev,
                                              pixelData, pixelMasks, sort)) {
                      psWarning("Bad median/stdev at %d,%d", x, y);
                      break;
                  }
                  limit = rej * stdev;
#ifdef TESTING
                  if (x == TEST_X && y == TEST_Y) {
                      fprintf(stderr, "Rejecting without variance; rejection limit: %f\n", limit);
                  }
#endif
              } else {
#ifdef TESTING
                  if (x == TEST_X && y == TEST_Y) {
                      fprintf(stderr, "Rejecting with variance...\n");
                  }
#endif
                  median = combinationWeightedOlympic(pixelData, pixelWeights, pixelMasks, discard, sort);
              }

#ifdef TESTING
              if (x == TEST_X && y == TEST_Y) {
                  fprintf(stderr, "Median: %f\n", median);
              }
#endif


// Mask a pixel for inspection
#define MASK_PIXEL_FOR_INSPECTION() \
    pixelMasks->data.PS_TYPE_VECTOR_MASK_DATA[j] = 0xff; \
    if (!rejectInspect) { \
        combineInspect(inputs, x, y, pixelSources->data.U16[j]); \
    } \
    numClipped++; \
    totalClipped++;

              for (int j = 0; j < num; j++) {
                  if (pixelMasks->data.PS_TYPE_VECTOR_MASK_DATA[j]) {
                      continue;
                  }
                  float diff = pixelData->data.F32[j] - median; // Difference from expected
                  if (useVariance) {
                      // Comparing squares --- cheaper than lots of sqrts
                      // pixelVariances includes the rejection limit, from above
                      if (PS_SQR(diff) > pixelLimits->data.F32[j]) {
                          MASK_PIXEL_FOR_INSPECTION();
#ifdef TESTING
                          if (x == TEST_X && y == TEST_Y) {
                              fprintf(stderr, "Rejecting input %d based on variance: %f > %f\n",
                                      j, diff, sqrtf(pixelLimits->data.F32[j]));
                          }
#endif
                      }
                  } else if (fabsf(diff) > limit) {
                      MASK_PIXEL_FOR_INSPECTION();
#ifdef TESTING
                      if (x == TEST_X && y == TEST_Y) {
                          fprintf(stderr, "Rejecting input %d based on distribution: %f > %f\n",
                                  j, diff, limit);
                      }
#endif
                  }
              }
          }

          if (rejectInspect && totalClipped > 0) {
              // Get rid of the masked values
              // The alternative to this is to make combinationMeanVariance() accept a mask
              int good = 0;            // Index of good value
              for (int i = 0; i < num; i++) {
                  if (pixelMasks->data.PS_TYPE_VECTOR_MASK_DATA[i]) {
                      continue;
                  }
                  if (i != good) {
                      pixelData->data.F32[good] = pixelData->data.F32[i];
                      pixelVariances->data.F32[good] = pixelVariances->data.F32[i];
                      pixelWeights->data.F32[good] = pixelWeights->data.F32[i];
                      pixelData->data.F32[good] = pixelData->data.F32[i];
                  }
                  good++;
              }
              pixelData->n = good;
              pixelVariances->n = good;
              pixelWeights->n = good;
              if (combinationMeanVariance(&mean, &var, pixelData, pixelVariances, pixelWeights)) {
                  imageValue = mean;
                  varianceValue = var;
                  maskValue = 0;
              } else {
                  imageValue = NAN;
                  varianceValue = NAN;
                  maskValue = bad;
              }
          }

          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;
    }

    return;
}


// 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?
                              bool *haveRejects, // Do we have lists of rejected pixels?
                              int *num,    // Number of inputs
                              int *numCols, int *numRows, // Size of (sky) images
                              psArray *input // Input array of pmStackData to validate
    )
{
    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);
    }
    *haveRejects = (data->reject != NULL);

    // 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);
        }
    }

    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 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->addVariance = addVariance;

    return data;
}

/// Stack input images
bool pmStackCombine(pmReadout *combined, psArray *input, psImageMaskType maskVal, psImageMaskType bad,
                    int kernelSize, int numIter, float rej, float sys, float discard,
                    bool entire, bool useVariance, bool safe, bool rejectInspect)
{
    PS_ASSERT_PTR_NON_NULL(combined, false);
    bool haveVariances;                 // Do we have the variance maps?
    bool haveRejects;                   // Do we have lists of rejected pixels?
    int num;                            // Number of inputs
    int numCols, numRows;               // Size of (sky) images
    if (!validateInputData(&haveVariances, &haveRejects, &num, &numCols, &numRows, input)) {
        return false;
    }
    PS_ASSERT_INT_NONNEGATIVE(kernelSize, false);
    PS_ASSERT_INT_POSITIVE(bad, false);
    PS_ASSERT_INT_NONNEGATIVE(numIter, false);
    if (isnan(rej)) {
        PS_ASSERT_INT_EQUAL(numIter, 0, false);
    } else {
        PS_ASSERT_FLOAT_LARGER_THAN(rej, 0.0, false);
    }
    if (haveRejects) {
        // This is a subsequent combination, so expect that the image and mask already exist
        PS_ASSERT_IMAGE_NON_NULL(combined->image, false);
        PS_ASSERT_IMAGE_TYPE(combined->image, PS_TYPE_F32, false);
        PS_ASSERT_IMAGE_NON_NULL(combined->mask, false);
        PS_ASSERT_IMAGE_TYPE(combined->mask, PS_TYPE_IMAGE_MASK, false);
        PS_ASSERT_IMAGES_SIZE_EQUAL(combined->image, combined->mask, 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
    psArray *stack = psArrayAlloc(num); // Stack of readouts
    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;
            continue;
        }
        weights->data.F32[i] = data->weight;
        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 (!haveRejects && !data->inspect) {
            data->inspect = psPixelsAllocEmpty(PIXEL_LIST_BUFFER);
        }
    }

    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(PS_ERR_UNKNOWN, 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);

    if (haveRejects) {
        psImage *combinedImage = combined->image; // Combined image
        psImage *combinedMask = combined->mask; // Combined mask
        psImage *combinedVariance = combined->variance; // Combined variance map

        psArray *pixelMap = pixelMapGenerate(input, minInputCols, maxInputCols,
                                             minInputRows, maxInputRows); // Map of pixels to source
        psPixels *pixels = NULL;            // Total list of pixels, with no duplicates
        for (int i = 0; i < num; i++) {
            pmStackData *data = input->data[i]; // Stacking data; contains the list of pixels
            if (!data) {
                continue;
            }
            pixels = psPixelsConcatenate(pixels, data->reject);
        }
        pixels = psPixelsDuplicates(pixels, pixels);

        if (entire) {
            // Combine entire image
            for (int y = minInputRows; y < maxInputRows; y++) {
                for (int x = minInputCols; x < maxInputCols; x++) {
                    psVector *reject = pixelMapQuery(pixelMap, minInputCols, minInputRows,
                                                     x, y); // Reject these images
                    combinePixels(combinedImage, combinedMask, combinedVariance, input, weights,
                                  addVariance, reject, x, y, maskVal, bad, numIter, rej, sys, discard,
                                  useVariance, safe, rejectInspect, buffer);
                }
            }
        } else {
            // Only combine previously rejected pixels
            for (int i = 0; i < pixels->n; i++) {
                // Pixel coordinates are in the frame of the original image
                int x = pixels->data[i].x, y = pixels->data[i].y; // Coordinates of interest
                if (x < minInputCols || x >= maxInputCols || y < minInputRows || y >= maxInputRows) {
                    continue;
                }
                psVector *reject = pixelMapQuery(pixelMap, minInputCols, minInputRows,
                                                 x, y); // Reject these images
                combinePixels(combinedImage, combinedMask, combinedVariance, input, weights,
                              addVariance, reject, x, y, maskVal, bad, numIter, rej, sys, discard,
                              useVariance, safe, rejectInspect, buffer);
            }
        }
        psFree(pixels);
        psFree(pixelMap);
    } else {
        // 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;
        }

        for (int y = minInputRows; y < maxInputRows; y++) {
            for (int x = minInputCols; x < maxInputCols; x++) {
                combinePixels(combinedImage, combinedMask, combinedVariance, input, weights,
                              addVariance, NULL, x, y, maskVal, bad, numIter, rej, sys, discard,
                              useVariance, safe, rejectInspect, buffer);
            }
        }

#ifndef PS_NO_TRACE
        if (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
    }

    psFree(weights);
    psFree(buffer);

    return true;
}

