#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"

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


bool pmSubtractionAnalysis(psMetadata *analysis, pmSubtractionKernels *kernels, psRegion *region,
                           int numCols, int numRows)
{
    if (analysis) {
        PS_ASSERT_METADATA_NON_NULL(analysis, 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);
        psFree(subRegion);
    }

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

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


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


    // Set the variance factors
    float vf1 = 1.0, vf2 = 1.0;         // Variance factors for each image
    switch (kernels->mode) {
      case PM_SUBTRACTION_MODE_1:
        vf1 = pmSubtractionVarianceFactor(kernels, 0.5, 0.5, false);
        break;
      case PM_SUBTRACTION_MODE_2:
        vf2 = pmSubtractionVarianceFactor(kernels, 0.5, 0.5, false);
        break;
      case PM_SUBTRACTION_MODE_DUAL:
        vf1 = pmSubtractionVarianceFactor(kernels, 0.5, 0.5, false);
        vf2 = pmSubtractionVarianceFactor(kernels, 0.5, 0.5, true);
        break;
      default:
        psAbort("Invalid subtraction mode: %x", kernels->mode);
    }

    // Weight by the area
    if (region) {
        float norm = (region->x1 - region->x0 + 1) * (region->y1 - region->y0 + 1) / (numCols * numRows);
        vf1 *= norm;
        vf2 *= norm;
    }

    // Update the variance factor
#define UPDATE_VARFACTOR(VF, ANALYSIS) { \
    psMetadataItem *vfItem = psMetadataLookup(analysis, ANALYSIS); \
    if (vfItem) { \
        psAssert(vfItem->type == PS_TYPE_F32, "Should be the type we said."); \
        vfItem->data.F32 += VF; \
    } else { \
        psMetadataAddF32(analysis, PS_LIST_TAIL, ANALYSIS, 0, "Variance factor weighted by the area", VF); \
    } \
}

    UPDATE_VARFACTOR(vf1, PM_SUBTRACTION_ANALYSIS_VARFACTOR_1);
    UPDATE_VARFACTOR(vf2, PM_SUBTRACTION_ANALYSIS_VARFACTOR_2);

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

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

    return true;
}
