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

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

#include "pmSubtractionStamps.h"
#include "pmSubtraction.h"
#include "pmSubtractionParams.h"

//////////////////////////////////////////////////////////////////////////////////////////////////////////////

#if 0
// Convolve the reference stamp by the kernel
static psKernel *convolveStamp(const pmSubtractionStamp *stamp, // Stamp to be convolved
                               const psKernel *kernel, // Kernel by which to convolve
                               int footprint // Size of area to be convolved
    )
{
    psKernel *convolution = psKernelAlloc(-footprint, footprint, -footprint, footprint); // Result
    psKernel *reference = stamp->reference; // Reference stamp, to be convolved

    // Range of kernel
    int uMin = kernel->xMin;
    int uMax = kernel->xMax;
    int vMin = kernel->yMin;
    int vMax = kernel->yMax;

    for (int y = -footprint; y <= footprint; y++) {
        psF32 *conv = &convolution->kernel[y][-footprint]; // Dereference convolution
        for (int x = -footprint; x <= footprint; x++, conv++) {
            *conv = 0.0;

            int xStart = x + uMin;      // Start index for convolution
            for (int v = vMin; v <= vMax; v++) {
                psF32 *ref = &reference->kernel[y + v][xStart]; // Dereference reference image
                psF32 *krnl = &kernel->kernel[v][uMin]; // Dereference kernel in x
                for (int u = uMin; u <= uMax; u++, ref++, krnl++) {
                    *conv += *ref * *krnl;
                }
            }
        }
    }

    return convolution;
}
#endif

// Accumulate cross-term sums for a stamp
static void accumulateCross(double *sumI, // Sum of I(x)/sigma(x)^2
                            double *sumII, // Sum of I(x)^2/sigma(x)^2
                            double *sumIC, // Sum of I(x)conv(x)/sigma(x)^2
                            const pmSubtractionStamp *stamp, // Stamp with weight
                            const psKernel *input, // Input image, I(x)
                            int kernelIndex, // Index for kernel component
                            int footprint // Size of region of interest
    )
{
    psKernel *weight = stamp->weight;   // Weight, sigma(x)^2
    psKernel *convolution = stamp->convolutions->data[kernelIndex]; // Convolution of interest

    for (int y = -footprint; y <= footprint; y++) {
        psF32 *in = &input->kernel[y][-footprint]; // Dereference input
        psF32 *wt = &weight->kernel[y][-footprint]; // Dereference weight
        psF32 *conv = &convolution->kernel[y][-footprint]; // Dereference convolution
        for (int x = -footprint; x <= footprint; x++, in++, wt++, conv++) {
            double temp = *in / *wt; // Temporary product
            *sumI += temp;
            *sumII += *in * temp;
            *sumIC += *conv * temp;
        }
    }
    return;
}

// Accumulate convolution sums for a stamp
static void accumulateConvolutions(double *sumC, // Sum of conv(x)/sigma(x)^2
                                   double *sumCC, // Sum of conv(x)^2/sigma(x)^2
                                   const pmSubtractionStamp *stamp, // Stamp with input and weight
                                   int kernelIndex, // Index for kernel component
                                   int footprint // Size of region of interest
    )
{
    psKernel *weight = stamp->weight;   // Weight, sigma(x)^2
    psKernel *convolution = stamp->convolutions->data[kernelIndex]; // Convolution of interest

    for (int y = -footprint; y <= footprint; y++) {
        psF32 *wt = &weight->kernel[y][-footprint]; // Dereference weight
        psF32 *conv = &convolution->kernel[y][-footprint]; // Dereference convolution
        for (int x = -footprint; x <= footprint; x++, wt++, conv++) {
            double convNoise = *conv / *wt; // Temporary product
            *sumC += convNoise;
            *sumCC += *conv * convNoise;
        }
    }
    return;
}

static double accumulateChi2(psKernel *input, // Input stamp
                             pmSubtractionStamp *stamp, // Stamp with weight
                             int kernelIndex, // Index for kernel component
                             double coeff, // Coefficient of convolution
                             double bg,  // Background term
                             int footprint // Size of region of interest
    )
{
    double chi2 = 0.0;
    psKernel *weight = stamp->weight;   // Weight, sigma(x)^2
    psKernel *convolution = stamp->convolutions->data[kernelIndex]; // Convolution of interest

    for (int y = -footprint; y <= footprint; y++) {
        psF32 *in = &input->kernel[y][-footprint]; // Dereference input
        psF32 *wt = &weight->kernel[y][-footprint]; // Dereference weight
        psF32 *conv = &convolution->kernel[y][-footprint]; // Dereference convolution
        for (int x = -footprint; x <= footprint; x++, in++, wt++, conv++) {
            chi2 += PS_SQR(*in - bg - coeff * *conv) / *wt;
        }
    }

    return chi2;
}

// Return the initial value of chi^2
static double initialChi2(psKernel *input, // Input stamp
                          const pmSubtractionStamp *stamp, // Stamp with weight
                          int footprint // Size of convolution
    )
{
    psKernel *weight = stamp->weight;   // Weight map
    psKernel *reference = stamp->reference; // Reference stamp

    double chi2 = 0.0;                  // Chi^2
    for (int y = -footprint; y <= footprint; y++) {
        psF32 *in = &input->kernel[y][-footprint]; // Dereference input
        psF32 *wt = &weight->kernel[y][-footprint]; // Dereference weight
        psF32 *ref = &reference->kernel[y][-footprint]; // Derference reference
        for (int x = -footprint; x <= footprint; x++, in++, wt++, ref++) {
            float diff = *in - *ref;    // Temporary value
            chi2 += PS_SQR(diff) / *wt;
        }
    }

    return chi2;
}

// Subtract a convolution from the input
static void subtractConvolution(psKernel *input, // Input stamp
                                const pmSubtractionStamp *stamp, // Stamp with weight
                                int kernelIndex, // Index for kernel component
                                float coeff, // Coefficient of subtraction
                                float bg, // Background term
                                int footprint // Size of region of interest
    )
{
    psKernel *convolution = stamp->convolutions->data[kernelIndex]; // Convolution of interest

    for (int y = -footprint; y <= footprint; y++) {
        psF32 *in = &input->kernel[y][-footprint]; // Dereference input
        psF32 *conv = &convolution->kernel[y][-footprint]; // Dereference convolution
        for (int x = -footprint; x <= footprint; x++, in++, conv++) {
            *in -= *conv * coeff + bg;
        }
    }

    return;
}


pmSubtractionKernels *pmSubtractionKernelsOptimumISIS(pmSubtractionKernelsType type, int size, int inner,
                                                     int spatialOrder, psVector *fwhms, int maxOrder,
                                                     const psArray *stamps, int footprint, float tolerance)
{
    if (type != PM_SUBTRACTION_KERNEL_ISIS && type != PM_SUBTRACTION_KERNEL_GUNK) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Invalid kernel type: %x\n", type);
        return NULL;
    }
    PS_ASSERT_INT_NONNEGATIVE(size, NULL);
    PS_ASSERT_INT_NONNEGATIVE(inner, NULL);
    PS_ASSERT_INT_NONNEGATIVE(spatialOrder, NULL);
    PS_ASSERT_VECTOR_NON_NULL(fwhms, NULL);
    PS_ASSERT_VECTOR_TYPE(fwhms, PS_TYPE_F32, NULL);
    PS_ASSERT_INT_NONNEGATIVE(maxOrder, NULL);
    PS_ASSERT_ARRAY_NON_NULL(stamps, NULL);
    PS_ASSERT_INT_NONNEGATIVE(footprint, NULL);
    PS_ASSERT_FLOAT_LARGER_THAN(tolerance, 0.0, NULL);

    // Generate the kernels to test
    int numGaussians = fwhms->n;       // Number of Gaussians
    int numKernels = numGaussians * (maxOrder + 1) * (maxOrder + 2) / 2; // Number of kernel components
    psString params = NULL;             // Parameter, for description
    for (int i = 0; i < numGaussians; i++) {
        psStringAppend(&params, "%.2f,", fwhms->data.F32[i]);
    }
    params[strlen(params) - 1] = '\0';

    psVector *orders = psVectorAlloc(numGaussians, PS_TYPE_S32); // Polynomial orders
    psVectorInit(orders, maxOrder);
    pmSubtractionKernels *kernels = p_pmSubtractionKernelsRawISIS(size, spatialOrder,
                                                                  fwhms, orders); // Kernels
    psFree(orders);
    psFree(kernels->description);
    kernels->description = NULL;
    psStringAppend(&kernels->description, "OptISIS(%d,(%s),%d)", size, params, spatialOrder);
    psFree(params);

    // Need to save the stamp inputs --- we're changing the values!
    int numStamps = stamps->n;          // Number of stamps
    psArray *inputs = psArrayAlloc(numStamps); // Deep copies of the inputs
    for (int i = 0; i < numStamps; i++) {
        pmSubtractionStamp *stamp = stamps->data[i]; // Stamp of interest
        psKernel *input = stamp->input; // Input image of interest
        psImage *copy = psImageCopy(NULL, input->image, PS_TYPE_F32); // Copy of the image
        inputs->data[i] = psKernelAllocFromImage(copy, size + footprint, size + footprint);
        psFree(copy);                   // Drop reference
    }

    // Generate the convolutions
    for (int i = 0; i < numStamps; i++) {
        pmSubtractionStamp *stamp = stamps->data[i]; // Stamp of interest
        if (!stamp->convolutions) {
            stamp->convolutions = psArrayAlloc(numKernels);
        }

        for (int j = 0; j < numKernels; j++) {
            psKernel *kernel = kernels->preCalc->data[j]; // Precalculated kernel
            stamp->convolutions->data[j] = p_pmSubtractionConvolveStampPrecalc(stamp->reference, kernel);
        }
    }

    // Calculate sums invariant with the input image
    double sum1 = 0.0;                  // sum of 1/sigma(x,y)^2
    for (int i = 0; i < numStamps; i++) {
        pmSubtractionStamp *stamp = stamps->data[i]; // Stamp of interest
        psKernel *weight = stamp->weight; // Weight map for stamp

        for (int v = -footprint; v <= footprint; v++) {
            psF32 *wt = weight->kernel[v]; // Dereference weight map
            for (int u = -footprint; u <= footprint; u++, wt++) {
                sum1 += 1.0 / *wt;
            }
        }
    }
    psVector *sumC = psVectorAlloc(numKernels, PS_TYPE_F64); // sum of R(x)*k(u)/sigma(x)^2
    psVector *sumCC = psVectorAlloc(numKernels, PS_TYPE_F64); // sum of [R(x)*k(u)]^2/sigma(x)^2
    psVectorInit(sumC, 0.0);
    psVectorInit(sumCC, 0.0);
    for (int i = 0; i < numKernels; i++) {
        for (int j = 0; j < numStamps; j++) {
            accumulateConvolutions(&sumC->data.F64[i], &sumCC->data.F64[i], stamps->data[j], i, footprint);
        }
    }

    // Initial chi^2
    double lastChi2 = 0.0;              // Chi^2 from last iteration
    int numPixels = numStamps * PS_SQR(2 * footprint + 1); // Number of pixels contributing to chi^2
    for (int i = 0; i < numStamps; i++) {
        lastChi2 += initialChi2(inputs->data[i], stamps->data[i], footprint);
    }
    lastChi2 /= numPixels;

    // Rank the kernel components
    psVector *ranking = psVectorAlloc(numKernels, PS_TYPE_S32); // Ranking of the kernel components
    psVectorInit(ranking, -1);
    int cutIndex = -1;                  // Index at which to cut off kernels
    for (int iter = 0; iter < numKernels; iter++) {
        int bestIndex = -1;             // Index of best kernel component
        double bestChi2 = INFINITY;     // Value of best chi^2
        double bestCoeff = 0;           // Value of best coefficient
        double bestBG = 0;              // Value of best background

        for (int i = 0; i < numKernels; i++) {
            if (ranking->data.S32[i] >= 0) {
                continue;
            }

            double sumI = 0.0;          // sum of I(x)/sigma(x)^2
            double sumII = 0.0;         // sum of I(x)^2/sigma(x)^2
            double sumIC = 0.0;         // sum of I(x)C(x)/sigma(x)^2

            for (int j = 0; j < numStamps; j++) {
                accumulateCross(&sumI, &sumII, &sumIC, stamps->data[j], inputs->data[j], i, footprint);
            }

            double invDet = 1.0 / (sum1 * sumCC->data.F64[i] - PS_SQR(sumC->data.F64[i])); // Determinant^-1
            double coeff = invDet * (sum1 * sumIC - sumC->data.F64[i] * sumI); // Coefficient for kernel
            double bg = invDet * (sumCC->data.F64[i] * sumI - sumC->data.F64[i] * sumIC); // Background

            double chi2 = 0.0;          // Chi^2
            for (int j = 0; j < numStamps; j++) {
                chi2 += accumulateChi2(inputs->data[j], stamps->data[j], i, coeff, bg, footprint);
            }

            if (chi2 < bestChi2) {
                bestIndex = i;
                bestCoeff = coeff;
                bestChi2 = chi2;
                bestBG = bg;
            }

            psTrace("psModules.imcombine", 8, "%d: %lf %lf %lf %lf %lf %lf\n", i, sum1, sumI, sumII, sumIC,
                    sumC->data.F64[i], sumCC->data.F64[i]);
            psTrace("psModules.imcombine", 6, "%d: %lf %lf %lf\n", i, coeff, bg, chi2);
        }
        bestChi2 /= numPixels;

        if (bestIndex == -1) {
            psError(PS_ERR_UNKNOWN, false, "Unable to find best kernel component in round %d.", iter);
            psFree(inputs);
            psFree(sumC);
            psFree(sumCC);
            psFree(ranking);
            psFree(kernels);
            return NULL;
        }

        // And the winner is....
        ranking->data.S32[bestIndex] = iter;
        // Remove its contribution, and don't include it in the future.
        for (int j = 0; j < numStamps; j++) {
            subtractConvolution(inputs->data[j], stamps->data[j], bestIndex, bestCoeff, bestBG, footprint);
        }

        double diff = lastChi2 - bestChi2; // Difference in chi^2 between iterations

        psTrace("psModules.imcombine", 3, "The winner of round %d is %d (%f,%d,%d): %lf (%lf) --> %lf\n",
                iter, bestIndex, kernels->widths->data.F32[bestIndex], kernels->u->data.S32[bestIndex],
                kernels->v->data.S32[bestIndex], bestCoeff, diff, bestChi2);

        if (fabsf(diff) < tolerance) {
            cutIndex = iter;
            break;
        }

        lastChi2 = bestChi2;
    }
    psFree(inputs);
    psFree(sumC);
    psFree(sumCC);

    if (cutIndex < 0) {
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Unable to converge to tolerance %g\n", tolerance);
        psFree(ranking);
        psFree(kernels);
        return NULL;
    }

    int newSize = cutIndex + 1;         // Size of new kernel basis set
    psTrace("psModules.imcombine", 2, "Accepting %d kernels.\n", newSize);
    psVector *uNew = psVectorAlloc(newSize, PS_TYPE_S32);
    psVector *vNew = psVectorAlloc(newSize, PS_TYPE_S32);
    psVector *widthsNew = psVectorAlloc(newSize, PS_TYPE_F32);
    psArray *preCalcNew = psArrayAlloc(newSize);

    for (int i = 0; i < numKernels; i++) {
        int rank = ranking->data.S32[i]; // This kernel component's ranking
        if (rank >= 0 && rank < newSize) {
            uNew->data.S32[rank] = kernels->u->data.S32[i];
            vNew->data.S32[rank] = kernels->v->data.S32[i];
            widthsNew->data.F32[rank] = kernels->widths->data.F32[i];
            preCalcNew->data[rank] = psMemIncrRefCounter(kernels->preCalc->data[i]);
        }
    }
    psFree(kernels->u);
    psFree(kernels->v);
    psFree(kernels->widths);
    psFree(kernels->preCalc);
    kernels->u = uNew;
    kernels->v = vNew;
    kernels->widths = widthsNew;
    kernels->preCalc = preCalcNew;
    kernels->num = newSize;

    psFree(ranking);

    // Maintain photometric scaling
    if (type == PM_SUBTRACTION_KERNEL_ISIS) {
        psKernel *subtract = kernels->preCalc->data[0]; // Kernel to subtract from the rest
        for (int i = 1; i < newSize; i++) {
            if (kernels->u->data.S32[i] % 2 == 0 && kernels->v->data.S32[i] % 2 == 0) {
                psKernel *kernel = kernels->preCalc->data[i]; // Kernel of interest
                psBinaryOp(kernel->image, kernel->image, "-", subtract->image);
            }
        }
    } else if (type == PM_SUBTRACTION_KERNEL_GUNK) {
        psStringPrepend(&kernels->description, "GUNK=");
        psStringAppend(&kernels->description, "+POIS(%d,%d)", inner, spatialOrder);

        for (int i = 0; i < newSize; i++) {
            if (kernels->u->data.S32[i] % 2 == 0 && kernels->v->data.S32[i] % 2 == 0) {
                psKernel *kernel = kernels->preCalc->data[i]; // Kernel of interest
                kernel->kernel[0][0] -= 1.0;
            }
        }

        if (!p_pmSubtractionKernelsAddGrid(kernels, numGaussians, inner)) {
            psAbort("Should never get here.");
        }

        kernels->subIndex = PS_SQR(2 * inner + 1) / 2 + numGaussians;
        assert(kernels->u->data.S32[kernels->subIndex] == 0 &&
               kernels->v->data.S32[kernels->subIndex] == 0);
    }

    return kernels;
}
