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

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

#include "pmFPA.h"
#include "pmSubtraction.h"
#include "pmSubtractionKernels.h"
#include "pmSubtractionEquation.h"

#include "pmSubtractionAnalysis.h"
#include "pmSubtractionVisual.h"

#define KERNEL_MOSAIC 2                 // Half-number of kernel instances in the mosaic image

//#define TESTING

bool pmSubtractionAnalysis(psMetadata *analysis, psMetadata *header,
                           pmSubtractionKernels *kernels, psRegion *region,
                           int numCols, int numRows)
{
    PS_ASSERT_METADATA_NON_NULL(analysis, false);
    PS_ASSERT_METADATA_NON_NULL(header, false);
    PM_ASSERT_SUBTRACTION_KERNELS_NON_NULL(kernels, false);
    PM_ASSERT_SUBTRACTION_KERNELS_SOLUTION(kernels, false);

    // Record region for subtraction
    {
        psRegion *subRegion;    // Region over which subtraction was performed
        if (region) {
            subRegion = psMemIncrRefCounter(region);
        } else {
            subRegion = psRegionAlloc(0, numCols, 0, numRows);
        }

        psMetadataAddPtr(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_REGION,
                         PS_DATA_REGION | PS_META_DUPLICATE_OK,
                         "Region over which subtraction was performed", subRegion);

        psString string = psRegionToString(*subRegion);
        psFree(subRegion);

        psMetadataAddStr(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_REGION, PS_META_DUPLICATE_OK,
                         "Region over which subtraction was performed", string);
        psFree(string);
    }

    // Record kernel
    psMetadataAddPtr(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_KERNEL,
                     PS_DATA_UNKNOWN | PS_META_DUPLICATE_OK, "Subtraction kernels", kernels);
    psMetadataAddS32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MODE,
                     PS_META_DUPLICATE_OK, "Subtraction mode", kernels->mode);
    psMetadataAddS32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MODE,
                     PS_META_DUPLICATE_OK, "Subtraction mode", kernels->mode);

    // Realisations of kernel
    {
        psTrace("psModules.imcombine", 2, "Generating diagnostics...\n");
        // Generate image with convolution kernels
        int size = kernels->size;       // Half-size of kernel
        int fullSize = 2 * size + 1 + 1; // Full size of kernel
        int imageSize = (2 * KERNEL_MOSAIC + 1) * fullSize;
        psImage *convKernels = psImageAlloc((kernels->mode == PM_SUBTRACTION_MODE_DUAL ? 2 : 1) *
                                            imageSize - 1 +
                                            (kernels->mode == PM_SUBTRACTION_MODE_DUAL ? 4 : 0),
                                            imageSize - 1, PS_TYPE_F32);
        psImageInit(convKernels, NAN);
        for (int j = -KERNEL_MOSAIC; j <= KERNEL_MOSAIC; j++) {
            for (int i = -KERNEL_MOSAIC; i <= KERNEL_MOSAIC; i++) {
                psImage *kernel = pmSubtractionKernelImage(kernels, (float)i / (float)KERNEL_MOSAIC,
                                                           (float)j / (float)KERNEL_MOSAIC,
                                                           false); // Image of the kernel
                if (!kernel) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to generate kernel image.");
                    psFree(convKernels);
                    return false;
                }

                if (psImageOverlaySection(convKernels, kernel, (i + KERNEL_MOSAIC) * fullSize,
                                          (j + KERNEL_MOSAIC) * fullSize, "=") == 0) {
                    psError(PS_ERR_UNKNOWN, false, "Unable to overlay kernel image.");
                    psFree(kernel);
                    psFree(convKernels);
                    return false;
                }
                psFree(kernel);

                if (kernels->mode == PM_SUBTRACTION_MODE_DUAL) {
                    kernel = pmSubtractionKernelImage(kernels, (float)i / (float)KERNEL_MOSAIC,
                                                      (float)j / (float)KERNEL_MOSAIC,
                                                      true); // Image of the kernel
                    if (!kernel) {
                        psError(PS_ERR_UNKNOWN, false, "Unable to generate kernel image.");
                        psFree(convKernels);
                        return false;
                    }

                    if (psImageOverlaySection(convKernels, kernel,
                                              (2 * KERNEL_MOSAIC + 1 + i + KERNEL_MOSAIC) * fullSize + 4,
                                              (j + KERNEL_MOSAIC) * fullSize, "=") == 0) {
                        psError(PS_ERR_UNKNOWN, false, "Unable to overlay kernel image.");
                        psFree(kernel);
                        psFree(convKernels);
                        return false;
                    }
                    psFree(kernel);
                }
            }
        }

        pmSubtractionVisualPlotConvKernels(convKernels);
        psMetadataAddImage(analysis, PS_LIST_TAIL, "SUBTRACTION.KERNEL.IMAGE",
                           PS_META_DUPLICATE_OK, "Realisations of kernel", convKernels);
        psFree(convKernels);
    }

    // sample difference images
    { 
        psMetadataAddArray(analysis, PS_LIST_TAIL, "SUBTRACTION.SAMPLE.STAMP.SET", PS_META_DUPLICATE_OK, "Sample Difference Stamps", kernels->sampleStamps);
    }

#ifdef TESTING
    // Generate images of the kernel components
    {
        psMetadata *header = psMetadataAlloc(); // Header
        for (int i = 0; i < kernels->solution1->n; i++) {
            psString name = NULL;       // Header keyword
            psStringAppend(&name, "SOLN%04d", i);
            psMetadataAddF64(header, PS_LIST_TAIL, name, 0, NULL, kernels->solution1->data.F64[i]);
            psFree(name);
        }
        psArray *kernelImages = pmSubtractionKernelSolutions(kernels, 0.0, 0.0, false);
        psFits *kernelFile = psFitsOpen("kernels1.fits", "w");
        (void)psFitsWriteImageCube(kernelFile, header, kernelImages, NULL);
        psFitsClose(kernelFile);
        psFree(kernelImages);
        psFree(header);
    }
    if (kernels->solution2) {
        psMetadata *header = psMetadataAlloc(); // Header
        for (int i = 0; i < kernels->solution2->n; i++) {
            psString name = NULL;       // Header keyword
            psStringAppend(&name, "SOLN%04d", i);
            psMetadataAddF64(header, PS_LIST_TAIL, name, 0, NULL, kernels->solution2->data.F64[i]);
            psFree(name);
        }
        psArray *kernelImages = pmSubtractionKernelSolutions(kernels, 0.0, 0.0, true);
        psFits *kernelFile = psFitsOpen("kernels2.fits", "w");
        (void)psFitsWriteImageCube(kernelFile, header, kernelImages, NULL);
        psFitsClose(kernelFile);
        psFree(kernelImages);
        psFree(header);
    }
#endif


    // Kernel shape
    {
        psImage *image = pmSubtractionKernelImage(kernels, 0.5, 0.5, false); // Image of the kernel
        if (!image) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate image of kernel.");
            return false;
        }
        int size = kernels->size;       // Half-size of kernel
        int fullSize = 2 * size + 1;    // Full size of kernel
        float norm = 0.0;               // Normalisation (kernel sum)
        for (int y = 0; y < fullSize; y++) {
            for (int x = 0; x < fullSize; x++) {
                norm += image->data.F32[y][x];
            }
        }
        float max = -INFINITY;          // Maximum fraction
        for (int r = 1; r < size; r++) {
            unsigned long r2 = PS_SQR(r); // r^2
            float sum = 0.0;            // Sum within circle
            for (int y = 0, v = -size; y < fullSize; y++, v++) {
                unsigned long v2 = PS_SQR(v); // y^2
                for (int x = 0, u = -size; x < fullSize; x++, u++) {
                    unsigned long u2 = PS_SQR(u); // u^2
                    if (u2 + v2 <= r2) {
                        sum += image->data.F32[y][x];
                    }
                }
            }
            float frac = sum / norm;    // Fraction of flux moving towards centre
            psTrace("psModules.imcombine", 5, "Deconvolution fraction at %d: %f\n", r, frac);
            max = PS_MAX(max, frac);
        }
        psFree(image);

        {
            psMetadataItem *item = psMetadataLookup(analysis, PM_SUBTRACTION_ANALYSIS_DECONV_MAX); // Previous
            if (item) {
                max = item->data.F32 = PS_MAX(item->data.F32, max);
            } else {
                psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_DECONV_MAX, 0,
                                 "Maximum deconvolution fraction", max);
            }
        }

        {
            psMetadataItem *item = psMetadataLookup(header, PM_SUBTRACTION_ANALYSIS_DECONV_MAX); // Previous
            if (item) {
                item->data.F32 = max;
            } else {
                psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_DECONV_MAX, 0,
                                 "Maximum deconvolution fraction", max);
            }
        }
    }

    // Kernel moments
    {
        psImage *image = pmSubtractionKernelImage(kernels, 0.5, 0.5, false); // Image of the kernel
        if (!image) {
            psError(PS_ERR_UNKNOWN, false, "Unable to generate image of kernel.");
            return false;
        }
        double m00 = 0, m10 = 0, m01 = 0, m20 = 0, m11 = 0, m02 = 0; // Moments to calculate
        int size = kernels->size;       // Half-size of kernel
        int fullSize = 2 * size + 1;    // Full size of kernel
        for (int y = 0, v = -size; y < fullSize; y++, v++) {
            for (int x = 0, u = -size; x < fullSize; x++, u++) {
                float value = image->data.F32[y][x]; // Value of kernel
                m00 += value;
                m10 += u * value;
                m01 += v * value;
                m20 += u * u * value;
                m11 += u * v * value;
                m02 += v * v * value;
            }
        }
        psFree(image);

        // Convert first moments to centroids
        m10 /= m00;
        m01 /= m00;

        // Convert second moments to covariance
        m20 = m20 / m00 - PS_SQR(m10);
        m02 = m02 / m00 - PS_SQR(m01);
        m11 = m11 / m00 - m10 * m01;

        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_NORM,
                         PS_META_DUPLICATE_OK, "Normalisation", m00);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MX,
                         PS_META_DUPLICATE_OK, "Moment in x", m10);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MY,
                         PS_META_DUPLICATE_OK, "Moment in y", m01);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MXX,
                         PS_META_DUPLICATE_OK, "Moment in xx", m20);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MXY,
                         PS_META_DUPLICATE_OK, "Moment in xy", m11);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MYY,
                         PS_META_DUPLICATE_OK, "Moment in yy", m02);

        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_NORM,
                         PS_META_DUPLICATE_OK, "Normalisation", m00);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MX,
                         PS_META_DUPLICATE_OK, "Moment in x", m10);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MY,
                         PS_META_DUPLICATE_OK, "Moment in y", m01);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MXX,
                         PS_META_DUPLICATE_OK, "Moment in xx", m20);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MXY,
                         PS_META_DUPLICATE_OK, "Moment in xy", m11);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_MYY,
                         PS_META_DUPLICATE_OK, "Moment in yy", m02);
    }

    // Difference in background
    {
        psImage *polyValues = p_pmSubtractionPolynomialFromCoords(NULL, kernels, numCols, numRows,
                                                                  numCols / 2.0, numRows / 2.0); // Polynomial
        float bg = p_pmSubtractionSolutionBackground(kernels, polyValues); // Background difference

        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_BGDIFF,
                         PS_META_DUPLICATE_OK, "Background difference", bg);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_BGDIFF,
                         PS_META_DUPLICATE_OK, "Background difference", bg);
        psFree(polyValues);
    }

    // Quality of fit
    {
        psMetadataAddS32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_STAMPS, 0, "Number of stamps",
                         kernels->numStamps);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_DEV_MEAN, 0, "Mean stamp deviation",
                         kernels->mean);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_DEV_RMS, 0, "RMS stamp deviation",
                         kernels->rms);

        psMetadataAddS32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_STAMPS, 0, "Number of stamps",
                         kernels->numStamps);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_DEV_MEAN, 0, "Mean stamp deviation",
                         kernels->mean);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_DEV_RMS, 0, "RMS stamp deviation",
                         kernels->rms);

        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FSIGMA_RES_MEAN,  0, "Mean Fractional Sigma of Residuals", kernels->fSigResMean);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FSIGMA_RES_STDEV, 0, "Mean Fractional Sigma of Residuals", kernels->fSigResStdev);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMIN_RES_MEAN,  0, "Mean Fractional Sigma of Residuals", kernels->fSigResMean);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMIN_RES_STDEV, 0, "Mean Fractional Sigma of Residuals", kernels->fSigResStdev);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMAX_RES_MEAN,  0, "Mean Fractional Sigma of Residuals", kernels->fSigResMean);
        psMetadataAddF32(analysis, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMAX_RES_STDEV, 0, "Mean Fractional Sigma of Residuals", kernels->fSigResStdev);

        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FSIGMA_RES_MEAN,  0, "Mean Fractional Sigma of Residuals", kernels->fSigResMean);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FSIGMA_RES_STDEV, 0, "Mean Fractional Sigma of Residuals", kernels->fSigResStdev);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMIN_RES_MEAN,  0, "Mean Fractional Sigma of Residuals", kernels->fSigResMean);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMIN_RES_STDEV, 0, "Mean Fractional Sigma of Residuals", kernels->fSigResStdev);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMAX_RES_MEAN,  0, "Mean Fractional Sigma of Residuals", kernels->fSigResMean);
        psMetadataAddF32(header, PS_LIST_TAIL, PM_SUBTRACTION_ANALYSIS_FMAX_RES_STDEV, 0, "Mean Fractional Sigma of Residuals", kernels->fSigResStdev);
    }

    return true;
}
