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

#include "tap.h"
#include "pstap.h"

#define IMAGE_SIZE 21
#define KERNEL_SIZE 3
#define TOL 1.0e-6


// Generate image with single high pixel
static psImage *generateImage(void)
{
    psImage *image = psImageAlloc(IMAGE_SIZE, IMAGE_SIZE, PS_TYPE_F32);
    psImageInit(image, 0.0);
    image->data.F32[IMAGE_SIZE/2][IMAGE_SIZE/2] = 1.0;
    return image;
}

// Generate Gaussian kernel
static psKernel *generateKernel(void)
{
    psKernel *kernel = psKernelAlloc(- KERNEL_SIZE, KERNEL_SIZE, - KERNEL_SIZE, KERNEL_SIZE);
    for (int y = - KERNEL_SIZE; y <= KERNEL_SIZE; y++) {
        for (int x = - KERNEL_SIZE; x <= KERNEL_SIZE; x++) {
            kernel->kernel[y][x] = exp(- x*x - y*y);
        }
    }
    return kernel;
}

// Check the convolved image matches the kernel
static bool checkConvolved(const psImage *image, const psKernel *kernel)
{
    bool correct = true;
    for (int y = 0; y < IMAGE_SIZE; y++) {
        for (int x = 0; x < IMAGE_SIZE; x++) {
            if (x < IMAGE_SIZE/2 - KERNEL_SIZE || x > IMAGE_SIZE/2 + KERNEL_SIZE ||
                y < IMAGE_SIZE/2 - KERNEL_SIZE || y > IMAGE_SIZE/2 + KERNEL_SIZE) {
                if (fabs(image->data.F32[y][x]) > TOL) {
                    diag("%d,%d --> %f", x, y, fabs(image->data.F32[y][x]));
                    correct = false;
                }
            } else {
                // Position relative to the centre
                int kx = x - IMAGE_SIZE/2;
                int ky = y - IMAGE_SIZE/2;
                if (fabs(image->data.F32[y][x] - kernel->kernel[ky][kx]) > TOL) {
                    diag("%d,%d --> %f", x, y, fabs(image->data.F32[y][x] - kernel->kernel[ky][kx]));
                    correct = false;
                }
            }
        }
    }
    return correct;
}


int main(int argc, char *argv[])
{
    plan_tests(10);

    diag("psImageConvolve tests");

    {
        // Direct convolution: 5 tests
        psMemId id = psMemGetId();

        psImage *image = generateImage();
        psKernel *kernel = generateKernel();

        psImage *convolved = psImageConvolveDirect(image, kernel);
        ok(convolved, "convolution result");
        skip_start(!convolved, 3, "convolution failed");
        ok(convolved->type.type == PS_TYPE_F32, "output type");
        ok(convolved->numCols == IMAGE_SIZE && convolved->numRows == IMAGE_SIZE, "output size %dx%d",
           convolved->numCols, convolved->numRows);
        ok(checkConvolved(convolved, kernel), "convolved image correct");
        psFree(convolved);
        skip_end();

        psFree(kernel);
        psFree(image);

        ok(!psMemCheckLeaks(id, NULL, NULL, false), "no memory leaks");
    }

    {
        // FFT convolution: 5 tests
        psMemId id = psMemGetId();

        psImage *image = generateImage();
        psKernel *kernel = generateKernel();

        psImage *convolved = psImageConvolveFFT(image, kernel, 0.0);
        ok(convolved, "convolution result");
        skip_start(!convolved, 3, "convolution failed");
        ok(convolved->type.type == PS_TYPE_F32, "output type");
        ok(convolved->numCols == IMAGE_SIZE && convolved->numRows == IMAGE_SIZE, "output size %dx%d",
           convolved->numCols, convolved->numRows);
        ok(checkConvolved(convolved, kernel), "convolved image correct");
        psFree(convolved);
        skip_end();

        psFree(kernel);
        psFree(image);

        ok(!psMemCheckLeaks(id, NULL, NULL, false), "no memory leaks");
    }


    return exit_status();
}
