#include "JumpFlooding.h"
#include "CudaGLImage.h"

#include <cmath>
#include <algorithm>
#include <iostream>

__device__
int idx2D(int x, int y, int width) {
    return y * width + x;
}

__device__
float dist2(int x, int y, JumpFlooding::Seed s) {
    float dx = float(x - s.x);
    float dy = float(y - s.y);
    return dx * dx + dy * dy;
}

/**
 * Just copies the input data into the internal used Seed array.
 */
__global__
void initSeedsKernel(cudaSurfaceObject_t inputSurface, JumpFlooding::Seed* seeds, int width, int height) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height) return;

    uchar4 pixel;
    surf2Dread(&pixel, inputSurface, x * sizeof(uchar4), y);

    bool isSeed = pixel.x > 0 || pixel.y > 0 || pixel.z > 0;

    JumpFlooding::Seed seed;
    seed.x = x;
    seed.y = y;
    seed.color = pixel;
    seed.valid = isSeed ? 1 : 0;

    seeds[idx2D(x, y, width)] = seed;
}

/**
 * Performs a jump flooding iteration.
 *
 * @param input     Input image
 * @param output    Output image
 * @param width     Width of the image
 * @param height    Height of the image
 * @param stride    Stride of this iteration
 *
 * Hint: You can call `idx2D(nx, ny, width)` to get the correct index
 * in the 1D input and output array.
 */
__global__
void jumpFloodKernel(const JumpFlooding::Seed* input, JumpFlooding::Seed* output, int width,  int height, int stride){
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height)
        return;

    // TODO (1a): Implement the jump flooding pass here.

    output[idx2D(x, y, width)] = JumpFlooding::Seed();
}


/**
 * Kernel that converts the final jump flooding result into the desired representation
 * (e.g. voronoi diagram or distance transform):
 */
__global__
void renderKernel(const JumpFlooding::Seed* seeds, cudaSurfaceObject_t outputSurface, JFARenderMode mode, int width, int height, float renderModeParam){
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x >= width || y >= height)
        return;

    // Result of the last jump flooding pass:
    JumpFlooding::Seed seed = seeds[idx2D(x, y, width)];

    uchar4 color = make_uchar4(0u, 0u, 0u, 255u);

    // TODO (1b): You might want to place helper functions here.


    if(mode == JFARenderMode::Voronoi){
        color = make_uchar4(seed.color.x, seed.color.y, seed.color.z, 255);
    } else if(mode == JFARenderMode::Distance){
        // TODO (1b)

    } else if(mode == JFARenderMode::Dilation){
        // TODO (1b)

    } else if(mode == JFARenderMode::VoronoiBF || mode == JFARenderMode::VoronoiDiff){
        // TODO (1d)

    }

    // Functions that write onto the textures rendered via OpenGL (OpenGL interoperability):
    surf2Dwrite(color, outputSurface, x * sizeof(uchar4), y);
}

// Just a helper function so we can access memory between OpenGL and CUDA:
static cudaSurfaceObject_t makeSurface(cudaGraphicsResource* resource, cudaArray_t& array){
    cudaGraphicsSubResourceGetMappedArray(&array, resource, 0, 0);

    cudaResourceDesc desc = {};
    desc.resType = cudaResourceTypeArray;
    desc.res.array.array = array;

    cudaSurfaceObject_t surface = 0;
    cudaCreateSurfaceObject(&surface, &desc);

    return surface;
}

/**
 * Initializes the Ping Pong arrays for the jump flooding algorithm (JFA).
 */
JumpFlooding::JumpFlooding(int width, int height)
    : width(width), height(height) {
    size_t count = size_t(width) * size_t(height);
    cudaMalloc(&ping, count * sizeof(JumpFlooding::Seed));
    cudaMalloc(&pong, count * sizeof(JumpFlooding::Seed));
}

JumpFlooding::~JumpFlooding() {
    if (ping)
        cudaFree(ping);

    if (pong)
        cudaFree(pong);

    ping = nullptr;
    pong = nullptr;
}

/**
 * Overall jump flooding computation.
 */
void JumpFlooding::compute(CudaGLImage& inputImage, CudaGLImage& outputImage, JFARenderMode mode, int maxIterations, float renderModeParam){
    cudaGraphicsResource* resources[] = {
        inputImage.cudaResource,
        outputImage.cudaResource
    };

    // This is just for the OpenGL texture resources (to make them accessible in CUDA):
    cudaGraphicsMapResources(2, resources);

    // This is just for the OpenGL texture resources (to make them accessible in CUDA):
    cudaArray_t inputArray;
    cudaArray_t outputArray;

    // This is just for the OpenGL texture resources (to make them accessible in CUDA):
    cudaSurfaceObject_t inputSurface = makeSurface(inputImage.cudaResource, inputArray);
    cudaSurfaceObject_t outputSurface = makeSurface(outputImage.cudaResource, outputArray);

    // Block and Grid Size for kernel call:
    dim3 block(16, 16);
    dim3 grid((width + block.x - 1) / block.x, (height + block.y - 1) / block.y);

    // First step: Copy the
    initSeedsKernel<<<grid, block>>>(inputSurface, ping, width, height);

    // Define variable for maximum strides:
    int stride = 1;
    int maxDim = std::max(width, height);

    // Calculate the maximum strides:
    while (stride < maxDim)
        stride <<= 1;   // Stride * 2

    stride >>= 1;

    // Define source and destination images:
    JumpFlooding::Seed* src = ping;
    JumpFlooding::Seed* dst = pong;

    int iteration = 0;

    // Iterate as long as stride > 0 or maximum desired iterations are reached:
    while (stride > 0 && iteration < maxIterations) {
        // Execute an iteration of the jump flooding kernel:
        jumpFloodKernel<<<grid, block>>>(src, dst, width, height, stride);

        // Swap source and destination images (ping pong):
        std::swap(src, dst);

        // Half the stride for next pass:
        stride >>= 1;

        // Increase iteration count:
        iteration++;
    }

    // Convert the indices of the jump flooding result into the
    // wanted "type", e.g. Distance Transform or Voronoi Surface.
    renderKernel<<<grid, block>>>(src, outputSurface, mode, width, height, renderModeParam);

    cudaDeviceSynchronize();

    cudaDestroySurfaceObject(inputSurface);
    cudaDestroySurfaceObject(outputSurface);
    cudaGraphicsUnmapResources(2, resources);
}
