// Copyright (C) 2006  Paul A. Price (price@ifa.hawaii.edu)
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

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

#include <stdio.h>
#include <math.h>
#include <assert.h>
#include "pslib.h"
#include "pois.h"

/* Calculate the matrix and vector for the matrix equation */
int poisCalculateEquation(psArray *stamps, // The stamps
                          const psImage *refImage, // The reference image
                          const psImage *inImage, // The input image
                          const psArray *kernelParams, // Kernel parameters
                          const poisConfig *config // Configuration
    )
{
    // Check inputs
    assert(stamps);
    assert(refImage);
    assert(inImage);
    assert(kernelParams);
    assert(config);

    // In addition to the kernel parameters, need one extra parameter to fit the (assumed constant) difference
    // in the background levels.
    int numSolveParams = kernelParams->n + 1;
    int bgIndex = kernelParams->n;      // Index in matrix for the background (it's the last one, after the
                                        // kernel parameters)
    float **reference = refImage->data.F32; // The reference pixels
    float **input = inImage->data.F32;  // The input pixels
    int nPix = 0;                       // Number of pixels used in calculation (for interest)
    float nx = 0.5 * (float)refImage->numCols; // Half-size of input image in x
    float ny = 0.5 * (float)refImage->numRows; // Half-size of input image in y
    int numKernelParams = kernelParams->n; // Number of kernel parameters

    psPolynomial2D *poly = psPolynomial2DAlloc(PS_POLYNOMIAL_ORD, config->spatialOrder + 1,
                                               config->spatialOrder + 1); // Polynomial for evaluation
    for (int j = 0; j < config->spatialOrder + 1; j++) {
        for (int i = 0; i < config->spatialOrder + 1; i++) {
            poly->coeff[j][i] = 1.0;
            poly->mask[j][i] = 1;       // Mask all coefficients; unmask to evaluate
        }
    }

    psImage *polyValues = psImageAlloc(config->spatialOrder + 1, config->spatialOrder + 1,
                                       PS_TYPE_F64); // Evaluations for each of the polynomial orders

    /* Iterate over the stamps */
    for (int s = 0; s < stamps->n; s++) {
        poisStamp *stamp = stamps->data[s]; // The stamp

        if (stamp->status == POIS_STAMP_RECALC) {
            psTrace("pois.calculateEquation", 5, "Calculating for stamp %d: %d,%d\n", s, stamp->x, stamp->y);

            // Initialise the matrix and vector if required
            if (! stamp->matrix) {
                psTrace("pois.calculateEquation", 9, "Allocating matrix: %dx%d\n", numSolveParams,
                        numSolveParams);
                stamp->matrix = psImageAlloc(numSolveParams, numSolveParams, PS_TYPE_F64);
            }
            if (! stamp->vector) {
                psTrace("pois.calculateEquation", 3, "Allocating vector: %d\n", numSolveParams);
                stamp->vector = psVectorAlloc(numSolveParams, PS_TYPE_F64);
                stamp->vector->n = numSolveParams;
            }

            // Dereference, for convenience
            psImage *matrix = stamp->matrix;
            psVector *vector = stamp->vector;

            // Will assert on the type, since can't do much about it if it's wrong.
            assert(matrix->type.type == PS_TYPE_F64);
            assert(vector->type.type == PS_TYPE_F64);

            // Initialise the matrix and vector
            for (int j = 0; j < numSolveParams; j++) {
                for (int i = 0; i < numSolveParams; i++) {
                    matrix->data.F64[j][i] = 0.0;
                }
                vector->data.F64[j] = 0.0;
            }

            // Evaluate the polynomial for each order
            for (int j = 0; j < config->spatialOrder + 1; j++) {
                for (int i = 0; i < config->spatialOrder + 1 - j; i++) {
                    poly->mask[j][i] = 0;
                    polyValues->data.F64[j][i] = psPolynomial2DEval(poly, ((double)stamp->x - nx)/nx,
                                                                    ((double)stamp->y - ny)/ny);
                    poly->mask[j][i] = 1;
                }
            }

            // Iterate over the pixels
            for (int y = stamp->y - config->footprint; y < stamp->y + config->footprint; y++) {
                for (int x = stamp->x - config->footprint; x < stamp->x + config->footprint; x++) {
                    float invNoise2 = 1.0/reference[y][x]; // The inverse of the noise, squared.
                    nPix++;

                    /* Iterate over the first convolution */
                    for (int k1 = 0; k1 < numKernelParams; k1++) {
                        poisKernelBasis *kernel1 = kernelParams->data[k1]; // Kernel basis #1

                        int u1 = kernel1->u; // Offset in x
                        int v1 = kernel1->v; // Offset in y
                        int i1 = kernel1->xOrder; // Polynomial order in x
                        int j1 = kernel1->yOrder; // Polynomial order in y
                        // First convolution
                        float conv1 = polyValues->data.F64[j1][i1] * reference[y - v1][x - u1];

                        // Assuming that the first kernel component is 0 order in x and y, and 0 offset
                        if (k1 != 0) {
                            conv1 -= reference[y][x];
                        }

                        /* Iterate over the second convolution */
                        for (int k2 = k1; k2 < numKernelParams; k2++) {
                            poisKernelBasis *kernel2 = kernelParams->data[k2]; // Kernel basis #2

                            int u2 = kernel2->u; // Offset in x
                            int v2 = kernel2->v; // Offset in y
                            int i2 = kernel2->xOrder; // Polynomial order in x
                            int j2 = kernel2->yOrder; // Polynomial order in y
                            // Second convolution
                            float conv2 = polyValues->data.F64[j2][i2] * reference[y - v2][x - u2];

                            // Assuming that the first kernel component is 0 order in x and y, and 0 offset
                            if (k2 != 0) {
                                conv2 -= reference[y][x];
                            }

                            // Add into the matrix element
                            matrix->data.F64[k1][k2] += conv1 * // First convolution
                                conv2 * // Second convolution
                                invNoise2; // Divide by sigma^2, ~ poisson.
                        } // Second convolution

                        // Add into the vector element
                        vector->data.F64[k1] += input[y][x] * // Input image, no convolution
                            conv1 *     // Convolved reference image
                            invNoise2; // Divide by sigma^2, ~ poisson.

                        /* Background term */
                        matrix->data.F64[k1][bgIndex] += conv1 * // Convolved reference image
                            invNoise2; // Divide by sigma^2, ~ poisson.

                    } // First convolution

                    /* Background only terms */
                    matrix->data.F64[bgIndex][bgIndex] += invNoise2;
                    vector->data.F64[bgIndex] += input[y][x] * invNoise2;
                }
            } // Iterating over pixels in stamp

            // Fill in other side of symmetric matrix
            for (int k1 = 0; k1 < kernelParams->n; k1++) {
                for (int k2 = 0; k2 < k1; k2++) {
                    matrix->data.F64[k1][k2] = matrix->data.F64[k2][k1];
                }
                matrix->data.F64[bgIndex][k1] = matrix->data.F64[k1][bgIndex];
            }

            for (int k = 0; k < kernelParams->n; k++) {
                poisKernelBasis *kernel = kernelParams->data[k]; // Kernel basis #2
                int u = kernel->u; // Offset in x
                int v = kernel->v; // Offset in y
                matrix->data.F64[k][k] += config->penalty * (float)(u*u + v*v);
            }

            stamp->status = POIS_STAMP_USED; // We've calculated it now, so don't need to do so again.

        } // Legitimate stamps

    } // Iterating over stamps

    psFree(polyValues);
    psFree(poly);

    // Return number of pixels over which matrix has been calculated
    return nPix;
}
