#include "GitWorldBuildPass.h"

#include <cuda_runtime.h>

namespace
{
    __global__ void BuildFieldKernel(const float* Input, float* Output, int Count)
    {
        const int Index = blockIdx.x * blockDim.x + threadIdx.x;
        if (Index < Count)
        {
            Output[Index] = Input[Index] * (1.0f + static_cast<float>(Index) * 0.125f);
        }
    }
}

extern "C" GITWORLD_CUDA_API int GW_CudaBuildPass(
    const float* Input,
    float* Output,
    int Count,
    int* DeviceIndex)
{
    if (Input == nullptr || Output == nullptr || Count <= 0 || DeviceIndex == nullptr)
    {
        return 0;
    }

    int DeviceCount = 0;
    if (cudaGetDeviceCount(&DeviceCount) != cudaSuccess || DeviceCount == 0)
    {
        return 0;
    }

    *DeviceIndex = 0;
    if (cudaSetDevice(*DeviceIndex) != cudaSuccess)
    {
        return 0;
    }

    float* DeviceInput = nullptr;
    float* DeviceOutput = nullptr;
    const size_t Bytes = sizeof(float) * static_cast<size_t>(Count);
    if (cudaMalloc(&DeviceInput, Bytes) != cudaSuccess
        || cudaMalloc(&DeviceOutput, Bytes) != cudaSuccess)
    {
        if (DeviceInput != nullptr) cudaFree(DeviceInput);
        if (DeviceOutput != nullptr) cudaFree(DeviceOutput);
        return 0;
    }

    const bool CopiedIn = cudaMemcpy(DeviceInput, Input, Bytes, cudaMemcpyHostToDevice) == cudaSuccess;
    if (CopiedIn)
    {
        BuildFieldKernel<<<(Count + 127) / 128, 128>>>(DeviceInput, DeviceOutput, Count);
    }
    const bool Computed = CopiedIn
        && cudaGetLastError() == cudaSuccess
        && cudaDeviceSynchronize() == cudaSuccess;
    const bool CopiedOut = Computed
        && cudaMemcpy(Output, DeviceOutput, Bytes, cudaMemcpyDeviceToHost) == cudaSuccess;

    cudaFree(DeviceInput);
    cudaFree(DeviceOutput);
    return CopiedOut ? 1 : 0;
}