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

#define KERNEL_SIZE 10                  // Kernel half-size
#define KERNEL_SIGMA 2.0                // Gaussian width for kernel
#define MASK_BAD 0xF0                   // Mask for bad pixels
#define MASK_POOR 0x0F                  // Mask for poor pixels
#define POOR_FRAC 0.1                   // Poor fraction


static psImage *readImage(const char *name)
{
    psFits *fits = psFitsOpen(name, "r");
    psImage *image = psFitsReadImage(fits, psRegionSet(0,0,0,0), 0);
    psFitsClose(fits);
    return image;
}

void writeImage(const psImage *image,   // Image to write
                const char *name        // Name of FITS file
    )
{
    psFits *fits = psFitsOpen(name, "w");
    psFitsWriteImage(fits, NULL, image, 0, NULL);
    psFitsClose(fits);
}


int main(int argc, char *argv[])
{
    psImage *inBright = readImage("bright.fits");
    psImage *inFaint = readImage("faint.fits");
    psImage *inWeight = readImage("weight.fits");
    psImage *inMask = readImage("mask.fits");

    psImage *subMask = pmSubtractionMask(inMask, NULL, MASK_BAD | MASK_POOR, KERNEL_SIZE, KERNEL_SIZE,
                                         NAN, true);

    pmSubtractionKernels *kernels = pmSubtractionKernelsPOIS(KERNEL_SIZE, 0, 0.0, PM_SUBTRACTION_MODE_1);
    kernels->solution1 = psVectorAlloc(kernels->num + 2, PS_TYPE_F64);

    double sum = 0.0;                   // Sum of kernel
    for (int i = 0; i < kernels->num; i++) {
        int u = kernels->u->data.S32[i];
        int v = kernels->v->data.S32[i];

        sum += kernels->solution1->data.F64[i] = 123.45 * exp(-(PS_SQR(u) + PS_SQR(v)) / 0.5 / PS_SQR(KERNEL_SIGMA));
    }

    int normIndex = PM_SUBTRACTION_INDEX_NORM(kernels); // Index for normalisation
    kernels->solution1->data.F64[normIndex] = 2.0 * sum;
    int bgIndex = PM_SUBTRACTION_INDEX_BG(kernels); // Index for background
    kernels->solution1->data.F64[bgIndex] = 0.0;

    psImage *kernelImage = pmSubtractionKernelImage(kernels, 0.0, 0.0, false);
    writeImage(kernelImage, "kernel.fits");
    psFree(kernelImage);

    pmReadout *in1 = pmReadoutAlloc(NULL);
    pmReadout *in2 = pmReadoutAlloc(NULL);
    in1->image = psMemIncrRefCounter(inBright);
    in2->image = psMemIncrRefCounter(inFaint);
    in2->weight = psMemIncrRefCounter(inWeight);

    pmReadout *out1 = pmReadoutAlloc(NULL);
    pmReadout *out2 = pmReadoutAlloc(NULL);

    pmSubtractionConvolve(out1, NULL, in1, NULL, subMask, MASK_BAD, MASK_POOR, POOR_FRAC, NULL,
                          kernels, false, true);
    pmSubtractionConvolve(out2, NULL, in2, NULL, subMask, MASK_BAD, MASK_POOR, POOR_FRAC, NULL,
                          kernels, false, true);
    pmSubtractionBorder(out2->image, out2->weight, out2->mask, KERNEL_SIZE, MASK_BAD);

    psFree(subMask);
    psFree(kernels);

    psFree(inBright);
    psFree(inFaint);
    psFree(inWeight);
    psFree(inMask);

    psFree(in1);
    psFree(in2);

    writeImage(out1->image, "bright.conv.fits");
    writeImage(out2->image, "faint.conv.fits");
    writeImage(out2->weight, "weight.conv.fits");
    writeImage(out2->mask, "mask.conv.fits");

    psFree(out1);
    psFree(out2);

    return PS_EXIT_SUCCESS;
}
