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

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

#include "pmSubtractionKernels.h"


// Free function for pmSubtractionKernels
static void subtractionKernelsFree(pmSubtractionKernels *kernels)
{
    psFree(kernels->u);
    psFree(kernels->v);
    psFree(kernels->sigma);
    psFree(kernels->xOrder);
    psFree(kernels->yOrder);
    psFree(kernels->preCalc);
}

// Raise a number to an integer power
static inline float power(float value,  // Value
                          int exp       // Exponent
    )
{
    if (exp == 0) {
        return 1.0;
    }
    float result = value;               // Result to return
    for (int i = 2; i <= exp; i++) {
        result *= value;
    }
    return result;
}

//////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Public functions
//////////////////////////////////////////////////////////////////////////////////////////////////////////////

pmSubtractionKernels *pmSubtractionKernelsAlloc(int numBasisFunctions, pmSubtractionKernelsType type,
                                                int size, int spatialOrder)
{
    pmSubtractionKernels *kernels = psAlloc(sizeof(pmSubtractionKernels)); // Kernels, to return
    psMemSetDeallocator(kernels, (psFreeFunc)subtractionKernelsFree);

    kernels->type = type;
    kernels->u = psVectorAlloc(numBasisFunctions, PS_TYPE_S32);
    kernels->v = psVectorAlloc(numBasisFunctions, PS_TYPE_S32);
    kernels->xOrder = psVectorAlloc(numBasisFunctions, PS_TYPE_S32);
    kernels->yOrder = psVectorAlloc(numBasisFunctions, PS_TYPE_S32);
    kernels->sigma = NULL;
    kernels->subIndex = 0;
    kernels->preCalc = NULL;
    kernels->size = size;
    kernels->spatialOrder = spatialOrder;

    return kernels;
}

pmSubtractionKernels *pmSubtractionKernelsPOIS(int size, int spatialOrder)
{
    PS_ASSERT_INT_POSITIVE(size, NULL);
    PS_ASSERT_INT_NONNEGATIVE(spatialOrder, NULL);

    int num = PS_SQR(2 * size + 1) * (spatialOrder + 1) * (spatialOrder + 2) / 2; // Number of basis functions

    pmSubtractionKernels *kernels = pmSubtractionKernelsAlloc(num, PM_SUBTRACTION_KERNEL_POIS,
                                                              size, spatialOrder); // The kernels

    // Generate a set of kernels for each (u,v)
    for (int v = - size, index = 0; v <= size; v++) {
        for (int u = - size; u <= size; u++) {
            // Iterate over spatial order.  This loop creates the terms for
            // x^xOrder * y^yOrder  such that (xOrder+yOrder) <= spatialOrder.
            for (int xOrder = 0; xOrder <= spatialOrder; xOrder++) {
                for (int yOrder = 0; yOrder <= spatialOrder - xOrder; yOrder++, index++) {
                    kernels->u->data.S32[index] = u;
                    kernels->v->data.S32[index] = v;
                    kernels->xOrder->data.S32[index] = xOrder;
                    kernels->yOrder->data.S32[index] = yOrder;

                    psTrace("psModules.imcombine", 7, "Kernel %d: %d %d %d %d\n", index,
                            u, v, xOrder, yOrder);
                }
            }
        }
    }

    kernels->subIndex = (num - (spatialOrder + 1) * (spatialOrder + 2) / 2) / 2;
    assert(kernels->u->data.S32[kernels->subIndex] == 0 &&
           kernels->v->data.S32[kernels->subIndex] == 0 &&
           kernels->xOrder->data.S32[kernels->subIndex] == 0 &&
           kernels->yOrder->data.S32[kernels->subIndex] == 0);

    return kernels;
}

pmSubtractionKernels *pmSubtractionKernelsISIS(int size, int spatialOrder,
                                               const psVector *sigmas, const psVector *orders)
{
    PS_ASSERT_VECTOR_NON_NULL(sigmas, NULL);
    PS_ASSERT_VECTOR_TYPE(sigmas, PS_TYPE_F32, NULL);
    PS_ASSERT_VECTOR_NON_NULL(orders, NULL);
    PS_ASSERT_VECTOR_TYPE(orders, PS_TYPE_S32, NULL);
    PS_ASSERT_VECTORS_SIZE_EQUAL(sigmas, orders, NULL);
    PS_ASSERT_INT_POSITIVE(size, NULL);
    PS_ASSERT_INT_NONNEGATIVE(spatialOrder, NULL);

    int numGaussians = sigmas->n;       // Number of Gaussians

    int num = 0;                        // Number of basis functions
    for (int i = 0; i < numGaussians; i++) {
        int gaussOrder = orders->data.S32[i]; // Polynomial order to apply to Gaussian
        num += (gaussOrder + 1) * (gaussOrder + 2) / 2;
    }
    num *= (spatialOrder + 1) * (spatialOrder + 2) / 2;

    pmSubtractionKernels *kernels = pmSubtractionKernelsAlloc(num, PM_SUBTRACTION_KERNEL_ISIS,
                                                              size, spatialOrder); // The kernels
    kernels->sigma = psVectorAlloc(num, PS_TYPE_F32);
    kernels->preCalc = psArrayAlloc(num);

    // Set the kernel parameters
    for (int i = 0, index = 0; i < numGaussians; i++) {
        // Iterate over (u,v) order
        for (int uOrder = 0; uOrder <= orders->data.S32[i]; uOrder++) {
            for (int vOrder = 0; vOrder <= orders->data.S32[i] - uOrder; vOrder++) {

                // Set the pre-calculated kernel
                psKernel *preCalc = psKernelAlloc(-size, size, -size, size);
                for (int v = -size; v <= size; v++) {
                    for (int u = -size; u <= size; u++) {
                        preCalc->kernel[v][u] = power(u, uOrder) * power(v, vOrder) *
                            expf(-0.5 * (PS_SQR(u) + PS_SQR(v)) / PS_SQR(sigmas->data.F32[i]));
                    }
                }

#if 0
                // Subtract a particular kernel in order to preserve photometric calibration across image
                if (spatialOrder > 0 && index != kernels->subIndex) {
                    psKernel *subKernel = kernels->preCalc->data[kernels->subIndex]; // Kernel to subtract
                    (void)psBinaryOp(preCalc->image, preCalc->image, "-", subKernel->image);
                }
#endif

                // Iterate over spatial order.  This loop creates the terms for
                // x^xOrder * y^yOrder  such that (xOrder+yOrder) <= spatialOrder.
                for (int xOrder = 0; xOrder <= spatialOrder; xOrder++) {
                    for (int yOrder = 0; yOrder <= spatialOrder - xOrder; yOrder++, index++) {
                        kernels->sigma->data.F32[index] = sigmas->data.F32[i];
                        kernels->u->data.S32[index] = uOrder;
                        kernels->v->data.S32[index] = vOrder;
                        kernels->xOrder->data.S32[index] = xOrder;
                        kernels->yOrder->data.S32[index] = yOrder;
                        kernels->preCalc->data[index] = psMemIncrRefCounter(preCalc);

                        psTrace("psModules.imcombine", 7, "Kernel %d: %f %d %d %d %d\n", index,
                                sigmas->data.F32[i], uOrder, vOrder, xOrder, yOrder);
                    }
                }

                psFree(preCalc);        // Drop reference
            }
        }
    }

    kernels->subIndex = 0;
    assert(kernels->u->data.S32[kernels->subIndex] == 0 &&
           kernels->v->data.S32[kernels->subIndex] == 0 &&
           kernels->xOrder->data.S32[kernels->subIndex] == 0 &&
           kernels->yOrder->data.S32[kernels->subIndex] == 0);

    if (psTraceGetLevel("psModules.imcombine.kernel") >= 10) {
        for (int i = 0; i < num; i++) {
            psKernel *kernel = kernels->preCalc->data[i]; // Kernel of interest
            psString kernelName = NULL;
            psStringAppend(&kernelName, "kernel%03d.fits", i);
            psFits *kernelFile = psFitsOpen(kernelName, "w");
            psFree(kernelName);
            psFitsWriteImage(kernelFile, NULL, kernel->image, 0, NULL);
            psFitsClose(kernelFile);
        }
    }

    return kernels;
}

pmSubtractionKernels *pmSubtractionKernelsGenerate(pmSubtractionKernelsType type, int size, int spatialOrder,
                                                   const psVector *sigmas, const psVector *orders)
{
    switch (type) {
      case PM_SUBTRACTION_KERNEL_POIS:
        return pmSubtractionKernelsPOIS(size, spatialOrder);
      case PM_SUBTRACTION_KERNEL_ISIS:
        return pmSubtractionKernelsISIS(size, spatialOrder, sigmas, orders);
      default:
        psError(PS_ERR_BAD_PARAMETER_VALUE, true, "Unknown kernel type: %x", type);
        return NULL;
    }
}
