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

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

#include "pmSubtraction.h"
#include "pmSubtractionKernels.h"
#include "pmSubtractionStamps.h"

#include "pmSubtractionEquation.h"


//#define TESTING

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Private (file-static) functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Calculate the sum over a stamp product
static inline double calculateSumProduct(const psKernel *image1, // First image in multiplication
                                         const psKernel *image2, // Second image in multiplication
                                         const psKernel *weight, // Weight image
                                         int footprint // (Half-)Size of stamp
    )
{
    double sum = 0.0;                   // Sum of the image products
    for (int y = - footprint; y <= footprint; y++) {
        for (int x = - footprint; x <= footprint; x++) {
            sum += image1->kernel[y][x] * image2->kernel[y][x] / 1.0; // weight->kernel[y][x];
        }
    }
    return sum;
}

// Calculate a single element of the least-squares matrix, with the polynomial expansions in one direction
static inline bool calculateMatrixElement1(psImage *matrix, // Matrix to calculate
                                           int i, int j, // Coordinates of element
                                           const psKernel *image1, // First image in multiplication
                                           const psKernel *image2, // Second image in multiplication
                                           const psKernel *weight, // Weight image
                                           const psImage *polyValues, // Spatial polynomial values
                                           int numKernels, // Number of kernel basis functions
                                           int footprint, // (Half-)Size of stamp
                                           int spatialOrder, // Maximum order of spatial variation
                                           bool symmetric // Is the matrix symmetric?
    )
{
    double sum = calculateSumProduct(image1, image2, weight, footprint); // Sum of the image products
    if (!isfinite(sum)) {
        return false;
    }

    // Generate the pseudo-convolutions from the spatial polynomial terms
    for (int iyOrder = 0, iIndex = i; iyOrder <= spatialOrder; iyOrder++) {
        for (int ixOrder = 0; ixOrder <= spatialOrder - iyOrder; ixOrder++, iIndex += numKernels) {
            double convPoly = sum * polyValues->data.F64[iyOrder][ixOrder];

            assert(iIndex < matrix->numRows && j < matrix->numCols);

            matrix->data.F64[iIndex][j] = convPoly;
            if (symmetric) {

                assert(iIndex < matrix->numCols && j < matrix->numRows);

                matrix->data.F64[j][iIndex] = convPoly;
            }
        }
    }
    return true;
}

// Calculate a single element of the least-squares matrix, with the polynomial expansions in both directions
static inline bool calculateMatrixElement2(psImage *matrix, // Matrix to calculate
                                           int i, int j, // Coordinates of element
                                           const psKernel *image1, // First image in multiplication
                                           const psKernel *image2, // Second image in multiplication
                                           const psKernel *weight, // Weight image
                                           const psImage *polyValues, // Spatial polynomial values
                                           int numKernels, // Number of kernel basis functions
                                           int footprint, // (Half-)Size of stamp
                                           int spatialOrder, // Maximum order of spatial variation
                                           bool symmetric // Is the matrix symmetric?
    )
{
    double sum = calculateSumProduct(image1, image2, weight, footprint); // Sum of the image products
    if (!isfinite(sum)) {
        return false;
    }

    // Generate the pseudo-convolutions from the spatial polynomial terms
    for (int iyOrder = 0, iIndex = i; iyOrder <= spatialOrder; iyOrder++) {
        for (int ixOrder = 0; ixOrder <= spatialOrder - iyOrder; ixOrder++, iIndex += numKernels) {
            double iPoly = polyValues->data.F64[iyOrder][ixOrder]; // Value of polynomial
            for (int jyOrder = 0, jIndex = j; jyOrder <= spatialOrder; jyOrder++) {
                for (int jxOrder = 0; jxOrder <= spatialOrder - jyOrder; jxOrder++, jIndex += numKernels) {
                    double convPoly = sum * iPoly * polyValues->data.F64[jyOrder][jxOrder];

                    assert(iIndex < matrix->numRows && jIndex < matrix->numCols);

                    matrix->data.F64[iIndex][jIndex] = convPoly;
                    if (symmetric) {

                        assert(iIndex < matrix->numCols && jIndex < matrix->numRows);

                        matrix->data.F64[jIndex][iIndex] = convPoly;
                    }
                }
            }
        }
    }
    return true;
}

// Calculate the square part of the matrix derived from multiplying convolutions
static bool calculateMatrixSquare(psImage *matrix, // Matrix to calculate
                                  const psArray *convolutions1, // Convolutions for element 1
                                  const psArray *convolutions2, // Convolutions for element 2
                                  const psKernel *weight, // Weight image
                                  const psImage *polyValues, // Polynomial values
                                  int numKernels, // Number of kernel basis functions
                                  int spatialOrder, // Order of spatial variation
                                  int footprint // Half-size of stamp
                                  )
{
    bool symmetric = (convolutions1 == convolutions2 ? true : false); // Is matrix symmetric?

    for (int i = 0; i < numKernels; i++) {
        psKernel *iConv = convolutions1->data[i]; // Convolution for i-th element

        for (int j = (symmetric ? i : 0); j < numKernels; j++) {
            psKernel *jConv = convolutions2->data[j]; // Convolution for j-th element

            if (!calculateMatrixElement2(matrix, i, j, iConv, jConv, weight, polyValues, numKernels,
                                         footprint, spatialOrder, symmetric)) {
                psTrace("psModules.imcombine", 2, "Bad sumCC at %d, %d", i, j);
                return false;
            }
        }
    }

    return true;
}

// Calculate least-squares matrix and vector
static bool calculateMatrix(psImage *matrix, // Matrix to calculate
                            const pmSubtractionKernels *kernels, // Kernel components
                            const psArray *convolutions, // Convolutions of source with kernels
                            const psKernel *input, // Input stamp, or NULL
                            const psKernel *weight, // Weight stamp
                            const psImage *polyValues, // Spatial polynomial values
                            int footprint, // (Half-)Size of stamp
                            bool normAndBG // Calculate normalisation and background terms?
    )
{
    int numKernels = kernels->num;      // Number of kernel components
    int spatialOrder = kernels->spatialOrder; // Maximum order of spatial variation
    int numSpatial = PM_SUBTRACTION_POLYTERMS(spatialOrder); // Number of spatial variation terms
    int bgOrder = kernels->bgOrder;     // Maximum order of background fit
    int numBackground = normAndBG ? PM_SUBTRACTION_POLYTERMS(bgOrder) : 0; // Number of background terms
    int numTerms = numKernels * numSpatial + (normAndBG ? 1 + numBackground : 0); // Total number of terms
    assert(matrix);
    assert(matrix->numCols == matrix->numRows);
    assert(matrix->numCols == numTerms);
    assert(convolutions && convolutions->n == numKernels);
    assert(polyValues);
    assert(!normAndBG || input);        // If we want the normalisation and BG, then we need the input image

    // Square part of the matrix (convolution-convolution products)
    if (!calculateMatrixSquare(matrix, convolutions, convolutions, weight, polyValues, numKernels,
                               spatialOrder, footprint)) {
        return false;
    }

    // XXX To support higher-order background model than simply constant, the below code needs to be updated.
    if (normAndBG) {
        int normIndex = PM_SUBTRACTION_INDEX_NORM(kernels); // Index for normalisation
        int bgIndex = PM_SUBTRACTION_INDEX_BG(kernels); // Index in matrix for background

        for (int i = 0; i < numKernels; i++) {
            psKernel *conv = convolutions->data[i]; // Convolution for i-th element

            // Normalisation-convolution terms
            if (!calculateMatrixElement1(matrix, i, normIndex, conv, input, weight, polyValues, numKernels,
                                         footprint, spatialOrder, true)) {
                psTrace("psModules.imcombine", 2, "Bad sumIC at %d", i);
                return false;
            }

            // Background-convolution terms
            double sumC = 0.0;          // Sum of the convolution
            for (int y = - footprint; y <= footprint; y++) {
                for (int x = - footprint; x <= footprint; x++) {
                    sumC += conv->kernel[y][x] / 1.0; // weight->kernel[y][x];
                }
            }
            if (!isfinite(sumC)) {
                psTrace("psModules.imcombine", 2, "Bad sumC at %d", i);
                return false;
            }

            for (int yOrder = 0, index = i; yOrder <= spatialOrder; yOrder++) {
                for (int xOrder = 0; xOrder <= spatialOrder - yOrder; xOrder++, index += numKernels) {
                    double value = sumC * polyValues->data.F64[yOrder][xOrder];
                    matrix->data.F64[index][bgIndex] = value;
                    matrix->data.F64[bgIndex][index] = value;
                }
            }
        }

        // Background only, normalisation only, and background-normalisation terms
        double sum1 = 0.0;              // Sum of the weighting
        double sumI = 0.0;              // Sum of the input
        double sumII = 0.0;             // Sum of the input squared
        for (int y = - footprint; y <= footprint; y++) {
            for (int x = - footprint; x <= footprint; x++) {
                double invNoise2 = 1.0 / 1.0; // weight->kernel[y][x];
                double value = input->kernel[y][x] * invNoise2;
                sumI += value;
                sumII += value * input->kernel[y][x];
                sum1 += invNoise2;
            }
        }
        if (!isfinite(sumI)) {
            psTrace("psModules.imcombine", 2, "Bad sumI detected");
            return false;
        }
        if (!isfinite(sumII)) {
            psTrace("psModules.imcombine", 2, "Bad sumII detected");
            return false;
        }
        if (!isfinite(sum1)) {
            psTrace("psModules.imcombine", 2, "Bad sum1 detected");
            return false;
        }
        matrix->data.F64[normIndex][normIndex] = sumII;
        matrix->data.F64[bgIndex][bgIndex] = sum1;
        matrix->data.F64[normIndex][bgIndex] = sumI;
        matrix->data.F64[bgIndex][normIndex] = sumI;
    }

    return true;
}


// Calculate least-squares matrix and vector
static bool calculateVector(psVector *vector, // Vector to calculate, or NULL
                            const pmSubtractionKernels *kernels, // Kernel components
                            const psArray *convolutions, // Convolutions of source with kernels
                            const psKernel *input, // Input stamp, or NULL if !normAndBG
                            const psKernel *target, // Target stamp
                            const psKernel *weight, // Weight stamp
                            const psImage *polyValues, // Spatial polynomial values
                            int footprint, // (Half-)Size of stamp
                            bool normAndBG // Calculate normalisation and background terms?
    )
{
    int numKernels = kernels->num;      // Number of kernel components
    int spatialOrder = kernels->spatialOrder; // Maximum order of spatial variation
    int numSpatial = PM_SUBTRACTION_POLYTERMS(spatialOrder); // Number of spatial variation terms
    int bgOrder = kernels->bgOrder;     // Maximum order of background fit
    int numBackground = normAndBG ? PM_SUBTRACTION_POLYTERMS(bgOrder) : 0; // Number of background terms
    int numTerms = numKernels * numSpatial + (normAndBG ? 1 + numBackground : 0); // Total number of terms
    assert(vector && vector->n == numTerms);
    assert(convolutions && convolutions->n == numKernels);
    assert(target);
    assert(polyValues);
    assert(!normAndBG || input);       // If we want the normalisation and BG, then we need the input image

    // Convolution terms
    for (int i = 0; i < numKernels; i++) {
        psKernel *conv = convolutions->data[i]; // Convolution for i-th element
        double sumTC = 0.0;          // Sum of the target and convolution
        for (int y = - footprint; y <= footprint; y++) {
            for (int x = - footprint; x <= footprint; x++) {
                sumTC += target->kernel[y][x] * conv->kernel[y][x] / 1.0; // weight->kernel[y][x];
            }
        }
        if (!isfinite(sumTC)) {
            psTrace("psModules.imcombine", 2, "Bad sumTC at %d", i);
            return false;
        }
        for (int yOrder = 0, index = i; yOrder <= spatialOrder; yOrder++) {
            for (int xOrder = 0; xOrder <= spatialOrder - yOrder; xOrder++, index += numKernels) {
                vector->data.F64[index] = sumTC * polyValues->data.F64[yOrder][xOrder];
            }
        }
    }

    if (normAndBG) {
        // Background terms
        double sumT = 0.0;              // Sum of the target
        double sumIT = 0.0;             // Sum of the input-target product
        for (int y = - footprint; y <= footprint; y++) {
            for (int x = - footprint; x <= footprint; x++) {
                float value = target->kernel[y][x] / 1.0; // weight->kernel[y][x];
                sumIT += value * input->kernel[y][x];
                sumT += value;
            }
        }
        if (!isfinite(sumT)) {
            psTrace("psModules.imcombine", 2, "Bad sumI detected");
            return false;
        }
        if (!isfinite(sumIT)) {
            psTrace("psModules.imcombine", 2, "Bad sumIT detected");
            return false;
        }

        int normIndex = PM_SUBTRACTION_INDEX_NORM(kernels); // Index for normalisation term
        vector->data.F64[normIndex] = sumIT;
        int bgIndex = PM_SUBTRACTION_INDEX_BG(kernels); // Index for background term
        vector->data.F64[bgIndex] = sumT;
    }

    return true;
}



// Calculate the cross-matrix, composed of convolutions of each image
// Note that the cross-matrix is NOT square
static bool calculateMatrixCross(psImage *matrix, // Matrix to calculate
                                 const pmSubtractionKernels *kernels, // Kernel components
                                 const psArray *convolutions1, // Convolutions of image 1
                                 const psArray *convolutions2, // Convolutions of image 2
                                 const psKernel *image1, // Image 1 stamp
                                 const psKernel *weight, // Weight stamp
                                 const psImage *polyValues, // Spatial polynomial values
                                 int footprint // (Half-)Size of stamp
                                 )
{
    assert(matrix);
    int numKernels = kernels->num;      // Number of kernel components
    int spatialOrder = kernels->spatialOrder; // Maximum order of spatial variation
    int numSpatial = PM_SUBTRACTION_POLYTERMS(spatialOrder); // Number of spatial polynomial terms
    int numBackground = PM_SUBTRACTION_POLYTERMS(kernels->bgOrder); // Number of background terms
    int numCols = numKernels * numSpatial + 1 + numBackground; // Number of columns
    int numRows = numKernels * numSpatial; // Number of rows
    assert(matrix->numCols == numCols && matrix->numRows == numRows);
    assert(convolutions1 && convolutions1->n == numKernels);
    assert(convolutions2 && convolutions2->n == numKernels);

    int normIndex, bgIndex;             // Indices in matrix for normalisation and background terms
    PM_SUBTRACTION_INDICES(normIndex, bgIndex, kernels);

    if (!calculateMatrixSquare(matrix, convolutions1, convolutions2, weight, polyValues, numKernels,
                               spatialOrder, footprint)) {
        return false;
    }

    for (int i = 0; i < numKernels; i++) {
        // Normalisation
        psKernel *conv = convolutions2->data[i]; // Convolution
        if (!calculateMatrixElement1(matrix, i, normIndex, conv, image1, weight, polyValues, numKernels,
                                     footprint, spatialOrder, false)) {
            psTrace("psModules.imcombine", 2, "Bad sumIC at %d", i);
            return false;
        }

        // Background
        double sumC = 0.0;              // Sum of the weighting
        for (int y = - footprint; y <= footprint; y++) {
            for (int x = - footprint; x <= footprint; x++) {
                sumC += conv->kernel[y][x] / 1.0; // weight->kernel[y][x];
            }
        }
        if (!isfinite(sumC)) {
            psTrace("psModules.imcombine", 2, "Bad sumC detected at %d", i);
            return false;
        }
        for (int yOrder = 0, index = i; yOrder <= spatialOrder; yOrder++) {
            for (int xOrder = 0; xOrder <= spatialOrder - yOrder; xOrder++, index += numKernels) {
                matrix->data.F64[index][bgIndex] = sumC * polyValues->data.F64[yOrder][xOrder];
            }
        }
    }

    return true;
}


// Add in penalty term to least-squares vector
static bool calculatePenalty(psVector *vector, // Vector to which to add in penalty term
                             const pmSubtractionKernels *kernels // Kernel parameters
    )
{
    if (kernels->penalty == 0.0) {
        return true;
    }

    psVector *penalties = kernels->penalties; // Penalties for each kernel component
    int spatialOrder = kernels->spatialOrder; // Order of spatial variations
    int numKernels = kernels->num; // Number of kernel components
    for (int i = 0; i < numKernels; i++) {
        for (int yOrder = 0, index = i; yOrder <= spatialOrder; yOrder++) {
            for (int xOrder = 0; xOrder <= spatialOrder - yOrder; xOrder++, index += numKernels) {
                vector->data.F64[index] -= penalties->data.F32[i];
            }
        }
    }

    return true;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Semi-public functions
// XXX We might like to define these functions as "extern inline" but gcc currently doesn't handle this in c99
// mode.  See http://gcc.gnu.org/ml/gcc/2006-11/msg00006.html
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

// Calculate the value of a polynomial, specified by coefficients and polynomial values
double p_pmSubtractionCalculatePolynomial(const psVector *coeff, // Coefficients
                                                 const psImage *polyValues, // Polynomial values
                                                 int order, // Order of polynomials
                                                 int index, // Index at which to begin
                                                 int step // Step between subsequent indices
                                                 )
{
    double sum = 0.0;                   // Value of the polynomial sum
    for (int yOrder = 0; yOrder <= order; yOrder++) {
        for (int xOrder = 0; xOrder <= order - yOrder; xOrder++, index += step) {

            assert(index < coeff->n);

            sum += coeff->data.F64[index] * polyValues->data.F64[yOrder][xOrder];
        }
    }
    return sum;
}

double p_pmSubtractionSolutionCoeff(const pmSubtractionKernels *kernels, const psImage *polyValues,
                                           int index, bool wantDual)
{
#if 0
    // This is probably in a tight loop, so don't check inputs
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, NAN);
    PM_ASSERT_SUBTRACTION_KERNELS_SOLUTION(kernels, NAN);
    PS_ASSERT_IMAGE_NON_NULL(polyValues, NAN);
    PS_ASSERT_INT_POSITIVE(index, NAN);
#endif

    psVector *solution = wantDual ? kernels->solution2 : kernels->solution1; // Solution vector
    return p_pmSubtractionCalculatePolynomial(solution, polyValues, kernels->spatialOrder, index,
                                              kernels->num);
}

double p_pmSubtractionSolutionNorm(const pmSubtractionKernels *kernels)
{
#if 0
    // This is probably in a tight loop, so don't check inputs
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, NAN);
    PM_ASSERT_SUBTRACTION_KERNELS_SOLUTION(kernels, NAN);
    PS_ASSERT_IMAGE_NON_NULL(polyValues, NAN);
#endif

    int normIndex = PM_SUBTRACTION_INDEX_NORM(kernels); // Index for normalisation
    return kernels->solution1->data.F64[normIndex];
}

double p_pmSubtractionSolutionBackground(const pmSubtractionKernels *kernels,
                                                const psImage *polyValues)
{
#if 0
    // This is probably in a tight loop, so don't check inputs
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, NAN);
    PM_ASSERT_SUBTRACTION_KERNELS_SOLUTION(kernels, NAN);
    PS_ASSERT_IMAGE_NON_NULL(polyValues, NAN);
#endif

    int bgIndex = PM_SUBTRACTION_INDEX_BG(kernels); // Index for background
    return p_pmSubtractionCalculatePolynomial(kernels->solution1, polyValues, kernels->bgOrder, bgIndex, 1);
}

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

bool pmSubtractionCalculateEquation(pmSubtractionStampList *stamps, const pmSubtractionKernels *kernels)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, false);

    int footprint = stamps->footprint;  // Half-size of stamps
    int spatialOrder = kernels->spatialOrder; // Maximum order of spatial variation
    int numKernels = kernels->num;      // Number of kernel basis functions
    int numSpatial = PM_SUBTRACTION_POLYTERMS(spatialOrder); // Number of spatial variations
    int numBackground = PM_SUBTRACTION_POLYTERMS(kernels->bgOrder); // Number of background terms

    // Total number of parameters to solve for: coefficient of each kernel basis function, multipled by the
    // number of coefficients for the spatial polynomial, normalisation and a constant background offset.
    int numParams = numKernels * numSpatial + 1 + numBackground;

    psImage *polyValues = NULL;         // Polynomial terms

    // We iterate over each stamp, allocate the matrix and vectors if
    // necessary, and then calculate those matrix/vectors.
    for (int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
        if (stamp->status != PM_SUBTRACTION_STAMP_CALCULATE) {
            continue;
        }

        // Generate convolutions
        if (!pmSubtractionConvolveStamp(stamp, kernels, footprint)) {
            psError(PS_ERR_UNKNOWN, false, "Unable to convolve stamp %d.", i);
            psFree(polyValues);
            return NULL;
        }

#ifdef TESTING
        for (int j = 0; j < numKernels; j++) {
            if (stamp->convolutions1) {
                psString convName = NULL;
                psStringAppend(&convName, "conv1_%03d_%03d.fits", i, j);
                psFits *fits = psFitsOpen(convName, "w");
                psFree(convName);
                psKernel *conv = stamp->convolutions1->data[j];
                psFitsWriteImage(fits, NULL, conv->image, 0, NULL);
                psFitsClose(fits);
            }

            if (stamp->convolutions2) {
                psString convName = NULL;
                psStringAppend(&convName, "conv2_%03d_%03d.fits", i, j);
                psFits *fits = psFitsOpen(convName, "w");
                psFree(convName);
                psKernel *conv = stamp->convolutions2->data[j];
                psFitsWriteImage(fits, NULL, conv->image, 0, NULL);
                psFitsClose(fits);
            }
        }
#endif

        polyValues = p_pmSubtractionPolynomial(polyValues, spatialOrder, stamp->xNorm, stamp->yNorm);


            stamp->matrix1 = psImageAlloc(numParams, numParams, PS_TYPE_F64);
            stamp->vector1 = psVectorAlloc(numParams, PS_TYPE_F64);
#ifdef TESTING
            psImageInit(stamp->matrix1, NAN);
            psVectorInit(stamp->vector1, NAN);
#endif

        bool status;                    // Status of least-squares matrix/vector calculation
        switch (kernels->mode) {
          case PM_SUBTRACTION_MODE_1:
            status = calculateMatrix(stamp->matrix1, kernels, stamp->convolutions1, stamp->image1,
                                     stamp->weight, polyValues, footprint, true);
            status &= calculateVector(stamp->vector1, kernels, stamp->convolutions1, stamp->image1,
                                      stamp->image2, stamp->weight, polyValues, footprint, true);
            break;
          case PM_SUBTRACTION_MODE_2:
            status = calculateMatrix(stamp->matrix1, kernels, stamp->convolutions2, stamp->image2,
                                     stamp->weight, polyValues, footprint, true);
            status &= calculateVector(stamp->vector1, kernels, stamp->convolutions2, stamp->image2,
                                      stamp->image1, stamp->weight, polyValues, footprint, true);
            break;
          case PM_SUBTRACTION_MODE_DUAL:
            stamp->matrix2 = psImageAlloc(numKernels * numSpatial, numKernels * numSpatial, PS_TYPE_F64);
            stamp->matrixX = psImageAlloc(numParams, numKernels * numSpatial, PS_TYPE_F64);
            stamp->vector2 = psVectorAlloc(numKernels * numSpatial, PS_TYPE_F64);
#ifdef TESTING
            psImageInit(stamp->matrix2, NAN);
            psImageInit(stamp->matrixX, NAN);
            psVectorInit(stamp->vector2, NAN);
#endif
            status  = calculateMatrix(stamp->matrix1, kernels, stamp->convolutions1, stamp->image1,
                                      stamp->weight, polyValues, footprint, true);
            status &= calculateMatrix(stamp->matrix2, kernels, stamp->convolutions2, NULL,
                                      stamp->weight, polyValues, footprint, false);
            status &= calculateMatrixCross(stamp->matrixX, kernels, stamp->convolutions1,
                                           stamp->convolutions2, stamp->image1, stamp->weight, polyValues,
                                           footprint);
            status &= calculateVector(stamp->vector1, kernels, stamp->convolutions1, stamp->image1,
                                      stamp->image2, stamp->weight, polyValues, footprint, true);
            status &= calculateVector(stamp->vector2, kernels, stamp->convolutions2, NULL,
                                      stamp->image2, stamp->weight, polyValues, footprint, false);
            break;
          default:
            psAbort("Unsupported subtraction mode: %x", kernels->mode);
        }

        if (!status) {
            stamp->status = PM_SUBTRACTION_STAMP_REJECTED;
            psWarning("Rejecting stamp %d (%d,%d) because of bad equation",
                      i, (int)(stamp->x + 0.5), (int)(stamp->y + 0.5));
        } else {
            stamp->status = PM_SUBTRACTION_STAMP_USED;
        }

#ifdef TESTING
        if (psTraceGetLevel("psModules.imcombine.equation") >= 10) {
            psString matrixName = NULL;
            psStringAppend(&matrixName, "matrix1_%d.fits", i);
            psFits *matrixFile = psFitsOpen(matrixName, "w");
            psFree(matrixName);
            psFitsWriteImage(matrixFile, NULL, stamp->matrix1, 0, NULL);
            psFitsClose(matrixFile);

            matrixName = NULL;
            psStringAppend(&matrixName, "vector1_%d.fits", i);
            psImage *dummy = psImageAlloc(stamp->vector1->n, 1, PS_TYPE_F64);
            memcpy(dummy->data.F64[0], stamp->vector1->data.F64,
                   PSELEMTYPE_SIZEOF(PS_TYPE_F64) * stamp->vector1->n);
            matrixFile = psFitsOpen(matrixName, "w");
            psFree(matrixName);
            psFitsWriteImage(matrixFile, NULL, dummy, 0, NULL);
            psFree(dummy);
            psFitsClose(matrixFile);

            if (stamp->vector2) {
                matrixName = NULL;
                psStringAppend(&matrixName, "vector2_%d.fits", i);
                dummy = psImageAlloc(stamp->vector2->n, 1, PS_TYPE_F64);
                memcpy(dummy->data.F64[0], stamp->vector2->data.F64,
                       PSELEMTYPE_SIZEOF(PS_TYPE_F64) * stamp->vector2->n);
                matrixFile = psFitsOpen(matrixName, "w");
                psFree(matrixName);
                psFitsWriteImage(matrixFile, NULL, dummy, 0, NULL);
                psFree(dummy);
                psFitsClose(matrixFile);
            }

            if (kernels->mode == PM_SUBTRACTION_MODE_DUAL) {
                matrixName = NULL;
                psStringAppend(&matrixName, "matrix2_%d.fits", i);
                matrixFile = psFitsOpen(matrixName, "w");
                psFree(matrixName);
                psFitsWriteImage(matrixFile, NULL, stamp->matrix2, 0, NULL);
                psFitsClose(matrixFile);

                matrixName = NULL;
                psStringAppend(&matrixName, "matrixX_%d.fits", i);
                matrixFile = psFitsOpen(matrixName, "w");
                psFree(matrixName);
                psFitsWriteImage(matrixFile, NULL, stamp->matrixX, 0, NULL);
                psFitsClose(matrixFile);
            }
        }
#endif

    }
    psFree(polyValues);

    return true;
}

bool pmSubtractionSolveEquation(pmSubtractionKernels *kernels, const pmSubtractionStampList *stamps)
{
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, false);
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, false);

    // Check inputs
    int numParams = -1;                // Number of parameters
    int numParams2 = 0;                // Number of parameters for part solution (DUAL mode)
    for (int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
        PS_ASSERT_PTR_NON_NULL(stamp, false);
        if (stamp->status != PM_SUBTRACTION_STAMP_USED) {
            continue;
        }

        PS_ASSERT_VECTOR_NON_NULL(stamp->vector1, false);
        if (numParams == -1) {
            numParams = stamp->vector1->n;
        }
        PS_ASSERT_VECTOR_SIZE(stamp->vector1, (long)numParams, false);
        PS_ASSERT_VECTOR_TYPE(stamp->vector1, PS_TYPE_F64, false);
        PS_ASSERT_IMAGE_NON_NULL(stamp->matrix1, false);
        PS_ASSERT_IMAGE_SIZE(stamp->matrix1, numParams, numParams, false);
        PS_ASSERT_IMAGE_TYPE(stamp->matrix1, PS_TYPE_F64, false);

        if (kernels->mode == PM_SUBTRACTION_MODE_DUAL) {
            PS_ASSERT_IMAGE_NON_NULL(stamp->matrix2, false);
            PS_ASSERT_IMAGE_NON_NULL(stamp->matrixX, false);
            if (numParams2 == 0) {
                numParams2 = stamp->matrix2->numCols;
            }
            PS_ASSERT_IMAGE_SIZE(stamp->matrix2, numParams2, numParams2, false);
            PS_ASSERT_IMAGE_SIZE(stamp->matrixX, numParams, numParams2, false);
            PS_ASSERT_IMAGE_TYPE(stamp->matrix2, PS_TYPE_F64, false);
            PS_ASSERT_IMAGE_TYPE(stamp->matrixX, PS_TYPE_F64, false);
            PS_ASSERT_VECTOR_NON_NULL(stamp->vector2, false);
            PS_ASSERT_VECTOR_SIZE(stamp->vector2, (long)numParams2, false);
            PS_ASSERT_VECTOR_TYPE(stamp->vector2, PS_TYPE_F64, false);
        }
    }
    if (numParams == -1) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "No suitable stamps found.");
        return NULL;
    }

    if (kernels->mode != PM_SUBTRACTION_MODE_DUAL) {
        // Accumulate the least-squares matricies and vectors
        psImage *sumMatrix = psImageAlloc(numParams, numParams, PS_TYPE_F64); // Combined matrix
        psVector *sumVector = psVectorAlloc(numParams, PS_TYPE_F64); // Combined vector
        psVectorInit(sumVector, 0.0);
        psImageInit(sumMatrix, 0.0);
        int numStamps = 0;              // Number of good stamps
        for (int i = 0; i < stamps->num; i++) {
            pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
            if (stamp->status == PM_SUBTRACTION_STAMP_USED) {
                (void)psBinaryOp(sumMatrix, sumMatrix, "+", stamp->matrix1);
                (void)psBinaryOp(sumVector, sumVector, "+", stamp->vector1);
                numStamps++;
            }
        }
        calculatePenalty(sumVector, kernels);

        psVector *permutation = NULL;       // Permutation vector, required for LU decomposition
        psImage *luMatrix = psMatrixLUD(NULL, &permutation, sumMatrix);
        psFree(sumMatrix);
        if (!luMatrix) {
            psError(PS_ERR_UNKNOWN, true, "LU Decomposition of least-squares matrix failed.\n");
            psFree(sumVector);
            psFree(luMatrix);
            psFree(permutation);
            return NULL;
        }
        kernels->solution1 = psMatrixLUSolve(kernels->solution1, luMatrix, sumVector, permutation);
        psFree(sumVector);
        psFree(luMatrix);
        psFree(permutation);
        if (!kernels->solution1) {
            psError(PS_ERR_UNKNOWN, true, "Failed to solve the least-squares system.\n");
            return NULL;
        }
    } else {
        // Dual convolution solution

        // Accumulation of stamp matrices/vectors
        psImage *sumMatrix1 = psImageAlloc(numParams, numParams, PS_TYPE_F64);
        psImage *sumMatrix2 = psImageAlloc(numParams2, numParams2, PS_TYPE_F64);
        psImage *sumMatrixX = psImageAlloc(numParams, numParams2, PS_TYPE_F64);
        psVector *sumVector1 = psVectorAlloc(numParams, PS_TYPE_F64);
        psVector *sumVector2 = psVectorAlloc(numParams, PS_TYPE_F64);
        psImageInit(sumMatrix1, 0.0);
        psImageInit(sumMatrix2, 0.0);
        psImageInit(sumMatrixX, 0.0);
        psVectorInit(sumVector1, 0.0);
        psVectorInit(sumVector2, 0.0);

        int numStamps = 0;              // Number of good stamps
        for (int i = 0; i < stamps->num; i++) {
            pmSubtractionStamp *stamp = stamps->stamps->data[i]; // Stamp of interest
            if (stamp->status == PM_SUBTRACTION_STAMP_USED) {
                (void)psBinaryOp(sumMatrix1, sumMatrix1, "+", stamp->matrix1);
                (void)psBinaryOp(sumMatrix2, sumMatrix2, "+", stamp->matrix2);
                (void)psBinaryOp(sumMatrixX, sumMatrixX, "+", stamp->matrixX);
                (void)psBinaryOp(sumVector1, sumVector1, "+", stamp->vector1);
                (void)psBinaryOp(sumVector2, sumVector2, "+", stamp->vector2);
                numStamps++;
            }
        }
        calculatePenalty(sumVector1, kernels);
        calculatePenalty(sumVector2, kernels);

#if 0
        // Apply weighting to maximise the difference between solution coefficients for the same functions
        double fudge = PS_SQR(2 * stamps->footprint + 1); // Fudge factor
        for (int i = 0; i < kernels->num; i++) {
            sumMatrix1->data.F64[i][i] -= fudge;
            sumMatrix2->data.F64[i][i] += fudge;
        }
#endif

        // Pure matrix operations

        // A * a = Ct * b + d
        // C * a = B  * b + e
        //
        // a = (Ct * Bi * C - A)i (Ct * Bi * e - d)
        // b = Bi * (C * a - e)
        psVector *a = psVectorRecycle(kernels->solution1, numParams, PS_TYPE_F64);
        psVector *b = psVectorRecycle(kernels->solution2, numParams2, PS_TYPE_F64);
#ifdef TESTING
        psVectorInit(a, NAN);
        psVectorInit(b, NAN);
#endif
        psImage *A = sumMatrix1;
        psImage *B = sumMatrix2;
        psImage *C = sumMatrixX;
        psVector *d = sumVector1;
        psVector *e = sumVector2;

        assert(a->n == numParams);
        assert(b->n == numParams2);
        assert(A->numRows == numParams && A->numCols == numParams);
        assert(B->numRows == numParams2 && B->numCols == numParams2);
        assert(C->numRows == numParams2 && C->numCols == numParams);
        assert(d->n == numParams);
        assert(e->n == numParams2);

        psImage *Bi = psMatrixInvert(NULL, B, NULL);
        assert(Bi->numRows == numParams2 && Bi->numCols == numParams2);
        psImage *Ct = psMatrixTranspose(NULL, C);
        assert(Ct->numRows == numParams && Ct->numCols == numParams2);

        psImage *BiC = psMatrixMultiply(NULL, Bi, C);
        assert(BiC->numRows == numParams2 && BiC->numCols == numParams);
        psImage *CtBi = psMatrixMultiply(NULL, Ct, Bi);
        assert(CtBi->numRows == numParams && CtBi->numCols == numParams2);

        psImage *CtBiC = psMatrixMultiply(NULL, Ct, BiC);
        assert(CtBiC->numRows == numParams && CtBiC->numCols == numParams);

        psImage *F = (psImage*)psBinaryOp(NULL, CtBiC, "-", A);
        assert(F->numRows == numParams && F->numCols == numParams);
        float det = 0.0;
        psImage *Fi = psMatrixInvert(NULL, F, &det);
        assert(Fi->numRows == numParams && Fi->numCols == numParams);
        psTrace("psModules.imcombine", 4, "Determinant of F: %f\n", det);

        psVector *g = psVectorAlloc(numParams, PS_TYPE_F64);
#ifdef TESTING
        psVectorInit(g, NAN);
#endif
        assert(CtBi->numRows == numParams && CtBi->numCols == numParams2);
        assert(e->n == numParams2);
        assert(d->n == numParams);
        for (int i = 0; i < numParams; i++) {
            double value = 0.0;
            for (int j = 0; j < numParams2; j++) {
                value += CtBi->data.F64[i][j] * e->data.F64[j];
            }
            g->data.F64[i] = value - d->data.F64[i];
        }

        assert(Fi->numRows == numParams && Fi->numCols == numParams);
        assert(g->n == numParams);
        for (int i = 0; i < numParams; i++) {
            double value = 0.0;
            for (int j = 0; j < numParams; j++) {
                value += Fi->data.F64[i][j] * g->data.F64[j];
            }
            a->data.F64[i] = value;
        }

        psVector *h = psVectorAlloc(numParams2, PS_TYPE_F64);
#ifdef TESTING
        psVectorInit(h, NAN);
#endif
        assert(C->numRows == numParams2 && C->numCols == numParams);
        assert(a->n == numParams);
        assert(e->n == numParams2);
        for (int i = 0; i < numParams2; i++) {
            double value = 0.0;
            for (int j = 0; j < numParams; j++) {
                value += C->data.F64[i][j] * a->data.F64[j];
            }
            h->data.F64[i] = value - e->data.F64[i];
        }

        assert(Bi->numRows == numParams2 && Bi->numCols == numParams2);
        assert(h->n == numParams2);
        for (int i = 0; i < numParams2; i++) {
            double value = 0.0;
            for (int j = 0; j < numParams2; j++) {
                value += Bi->data.F64[i][j] * h->data.F64[j];
            }
            b->data.F64[i] = value;
        }


#if 0
        for (int i = 0; i < numParams; i++) {
            double aVal1 = 0.0, bVal1 = 0.0;
            for (int j = 0; j < numParams2; j++) {
                aVal1 += A->data.F64[i][j] * a->data.F64[j];
                bVal1 += Ct->data.F64[i][j] * b->data.F64[j];
            }
            bVal1 += d->data.F64[i];
            for (int j = numParams2; j < numParams; j++) {
                aVal1 += A->data.F64[i][j] * a->data.F64[j];
            }
            printf("%d: %lf\n", i, aVal1 - bVal1);
        }

        for (int i = 0; i < numParams2; i++) {
            double aVal2 = 0.0, bVal2 = 0.0;
            for (int j = 0; j < numParams2; j++) {
                aVal2 += C->data.F64[i][j] * a->data.F64[j];
                bVal2 += B->data.F64[i][j] * b->data.F64[j];
            }
            bVal2 += e->data.F64[i];
            for (int j = numParams2; j < numParams; j++) {
                aVal2 += C->data.F64[i][j] * a->data.F64[j];
            }
            printf("%d: %lf\n", i, aVal2 - bVal2);
        }
#endif

        {
            psFits *fits = psFitsOpen("sumMatrix1.fits", "w");
            psFitsWriteImage(fits, NULL, sumMatrix1, 0, NULL);
            psFitsClose(fits);
        }
        {
            psFits *fits = psFitsOpen("sumMatrix2.fits", "w");
            psFitsWriteImage(fits, NULL, sumMatrix2, 0, NULL);
            psFitsClose(fits);
        }
        {
            psFits *fits = psFitsOpen("sumMatrixX.fits", "w");
            psFitsWriteImage(fits, NULL, sumMatrixX, 0, NULL);
            psFitsClose(fits);
        }
        {
            psFits *fits = psFitsOpen("sumFinverse.fits", "w");
            psFitsWriteImage(fits, NULL, Fi, 0, NULL);
            psFitsClose(fits);
        }


        kernels->solution1 = a;
        kernels->solution2 = b;

        // XXXXX Free temporary matrices and vectors

    }

    if (psTraceGetLevel("psModules.imcombine") >= 7) {
        for (int i = 0; i < kernels->solution1->n; i++) {
            psTrace("psModules.imcombine", 7, "Solution 1 %d: %f\n", i, kernels->solution1->data.F64[i]);
        }
        if (kernels->mode == PM_SUBTRACTION_MODE_DUAL) {
            for (int i = 0; i < kernels->solution2->n; i++) {
                psTrace("psModules.imcombine", 7, "Solution 2 %d: %f\n", i, kernels->solution2->data.F64[i]);
            }
        }
     }

    return true;
}

psVector *pmSubtractionCalculateDeviations(pmSubtractionStampList *stamps,
                                           const pmSubtractionKernels *kernels)
{
    PM_ASSERT_SUBTRACTION_STAMP_LIST_NON_NULL(stamps, NULL);
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, NULL);
    PM_ASSERT_SUBTRACTION_KERNELS_SOLUTION(kernels, NULL);

    psVector *deviations = psVectorAlloc(stamps->num, PS_TYPE_F32); // Mean deviation for stamps
    int footprint = stamps->footprint; // Half-size of stamps
    long numPixels = PS_SQR(2 * footprint + 1); // Number of pixels in footprint
    double devNorm = 1.0 / (double)numPixels; // Normalisation for deviations
    int numKernels = kernels->num;      // Number of kernels

    psImage *polyValues = NULL;         // Polynomial values
    psKernel *residual = psKernelAlloc(-footprint, footprint, -footprint, footprint); // Residual image

    for (int i = 0; i < stamps->num; i++) {
        pmSubtractionStamp *stamp = stamps->stamps->data[i]; // The stamp of interest
        if (stamp->status != PM_SUBTRACTION_STAMP_USED) {
            deviations->data.F32[i] = NAN;
            continue;
        }

        // Calculate coefficients of the kernel basis functions
        polyValues = p_pmSubtractionPolynomial(polyValues, kernels->spatialOrder, stamp->xNorm, stamp->yNorm);
        double norm = p_pmSubtractionSolutionNorm(kernels); // Normalisation
        double background = p_pmSubtractionSolutionBackground(kernels, polyValues);// Difference in background

        // Calculate residuals
        psKernel *weight = stamp->weight; // Weight postage stamp
        psImageInit(residual->image, 0.0);
        if (kernels->mode != PM_SUBTRACTION_MODE_DUAL) {
            psKernel *target;           // Target postage stamp
            psKernel *source;           // Source postage stamp
            psArray *convolutions;      // Convolution postage stamps for each kernel basis function
            switch (kernels->mode) {
              case PM_SUBTRACTION_MODE_1:
                target = stamp->image2;
                source = stamp->image1;
                convolutions = stamp->convolutions1;
                break;
              case PM_SUBTRACTION_MODE_2:
                target = stamp->image1;
                source = stamp->image2;
                convolutions = stamp->convolutions2;
                break;
              default:
                psAbort("Unsupported subtraction mode: %x", kernels->mode);
            }

            for (int j = 0; j < numKernels; j++) {
                psKernel *convolution = convolutions->data[j]; // Convolution
                double coefficient = p_pmSubtractionSolutionCoeff(kernels, polyValues, j,
                                                                  false); // Coefficient
                for (int y = - footprint; y <= footprint; y++) {
                    for (int x = - footprint; x <= footprint; x++) {
                        residual->kernel[y][x] -= convolution->kernel[y][x] * coefficient;
                    }
                }
            }
            for (int y = - footprint; y <= footprint; y++) {
                for (int x = - footprint; x <= footprint; x++) {
                    residual->kernel[y][x] += target->kernel[y][x] - background - source->kernel[y][x] * norm;
                }
            }
        } else {
            // Dual convolution
            psArray *convolutions1 = stamp->convolutions1; // Convolutions of the first image
            psArray *convolutions2 = stamp->convolutions2; // Convolutions of the second image
            psKernel *image1 = stamp->image1; // The first image
            psKernel *image2 = stamp->image2; // The second image

            for (int j = 0; j < numKernels; j++) {
                psKernel *conv1 = convolutions1->data[j]; // Convolution of first image
                psKernel *conv2 = convolutions2->data[j]; // Convolution of second image
                double coeff1 = p_pmSubtractionSolutionCoeff(kernels, polyValues, j, false); // Coefficient 1
                double coeff2 = p_pmSubtractionSolutionCoeff(kernels, polyValues, j, true); // Coefficient 2

                for (int y = - footprint; y <= footprint; y++) {
                    for (int x = - footprint; x <= footprint; x++) {
                        residual->kernel[y][x] += conv2->kernel[y][x] * coeff2 - conv1->kernel[y][x] * coeff1;
                    }
                }
            }
            for (int y = - footprint; y <= footprint; y++) {
                for (int x = - footprint; x <= footprint; x++) {
                    residual->kernel[y][x] += image2->kernel[y][x] - background - image1->kernel[y][x] * norm;
                }
            }
        }

        double deviation = 0.0;         // Sum of differences
        for (int y = - footprint; y <= footprint; y++) {
            for (int x = - footprint; x <= footprint; x++) {
                double dev = PS_SQR(residual->kernel[y][x]) / weight->kernel[y][x];
                deviation += dev;
#ifdef TESTING
                residual->kernel[y][x] = dev;
#endif
            }
        }
        deviations->data.F32[i] = sqrtf(devNorm * deviation);
        psTrace("psModules.imcombine", 5, "Deviation for stamp %d (%d,%d): %f\n",
                i, (int)(stamp->x + 0.5), (int)(stamp->y + 0.5), deviations->data.F32[i]);
        if (!isfinite(deviations->data.F32[i])) {
            stamp->status = PM_SUBTRACTION_STAMP_REJECTED;
            psTrace("psModules.imcombine", 5,
                    "Rejecting stamp %d (%d,%d) because of non-finite deviation\n",
                    i, (int)(stamp->x + 0.5), (int)(stamp->y + 0.5));
            continue;
        }

#ifdef TESTING
        {
            psString filename = NULL;
            psStringAppend(&filename, "resid_%03d.fits", i);
            psFits *fits = psFitsOpen(filename, "w");
            psFree(filename);
            psFitsWriteImage(fits, NULL, residual->image, 0, NULL);
            psFitsClose(fits);
        }
        {
            psString filename = NULL;
            psStringAppend(&filename, "stamp_image_%03d.fits", i);
            psFits *fits = psFitsOpen(filename, "w");
            psFree(filename);
            psFitsWriteImage(fits, NULL, stamp->image1->image, 0, NULL);
            psFitsClose(fits);
        }
        {
            psString filename = NULL;
            psStringAppend(&filename, "stamp_weight_%03d.fits", i);
            psFits *fits = psFitsOpen(filename, "w");
            psFree(filename);
            psFitsWriteImage(fits, NULL, stamp->weight->image, 0, NULL);
            psFitsClose(fits);
        }
#endif

    }
    psFree(residual);
    psFree(polyValues);

    return deviations;
}
