// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "ImagePreprocessorBufferGPU.h"

#include <algorithm>
#include <climits>

static inline void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

__global__ void gather_kernel(const int32_t *__restrict__ image,
                              const uint32_t *__restrict__ npixel,
                              int32_t *__restrict__ values,
                              int count) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += blockDim.x * gridDim.x)
        values[i] = image[npixel[i]];
}

__global__ void gather_tiles_kernel(const int32_t *__restrict__ image, int width, int height,
                                    const int32_t *__restrict__ centre, int half,
                                    int32_t *__restrict__ tiles, size_t count) {
    const int side = 2 * half + 1;
    const size_t area = static_cast<size_t>(side) * side;
    for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += static_cast<size_t>(blockDim.x) * gridDim.x) {
        const size_t tile = i / area;
        const int p = static_cast<int>(i - tile * area);
        const int px = centre[2 * tile] - half + p % side;
        const int py = centre[2 * tile + 1] - half + p / side;
        tiles[i] = (px < 0 || py < 0 || px >= width || py >= height)
                   ? INT32_MIN : image[static_cast<size_t>(py) * width + px];
    }
}

ImagePreprocessorBufferGPU::ImagePreprocessorBufferGPU(size_t npixel, bool host_mirror)
    : ImagePreprocessorBuffer(npixel, host_mirror),
      gpu_image(npixel),
      // A no-op when the mirror was not allocated: CudaRegisteredVector skips an empty vector.
      buffer_reg(buffer),
      max_gather(StrongPixelLimit(npixel)),
      gpu_gather_index(max_gather),
      gpu_gather_value(max_gather) {
}

int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() {
    return gpu_image;
}

const int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() const {
    return gpu_image;
}

void ImagePreprocessorBufferGPU::Gather(const std::vector<uint32_t> &npixel, std::vector<int32_t> &values) const {
    values.resize(npixel.size());
    if (npixel.empty())
        return;

    const int count = static_cast<int>(npixel.size());
    cudaMemcpyAsync(gpu_gather_index.get(), npixel.data(), count * sizeof(uint32_t),
                    cudaMemcpyHostToDevice, gather_stream);
    gather_kernel<<<(count + 255) / 256, 256, 0, gather_stream>>>(
        gpu_image.get(), gpu_gather_index.get(), gpu_gather_value.get(), count);
    cuda_err(cudaGetLastError());
    cudaMemcpyAsync(values.data(), gpu_gather_value.get(), count * sizeof(int32_t),
                    cudaMemcpyDeviceToHost, gather_stream);
    cudaStreamSynchronize(gather_stream);
}

const int32_t *ImagePreprocessorBufferGPU::GatherTiles(const std::vector<int32_t> &x, const std::vector<int32_t> &y,
                                                       int width, int half) {
    const size_t side = 2 * half + 1;
    const size_t count = x.size() * side * side;
    if (count == 0)
        return nullptr;
    if (count > tile_capacity) {
        tile_capacity = std::max(count, 2 * tile_capacity);
        gpu_tiles = CudaDevicePtr<int32_t>(tile_capacity);
        host_tiles.resize(tile_capacity);
    }
    if (2 * x.size() > tile_centre.size()) {
        tile_centre.resize(std::max(2 * x.size(), 2 * tile_centre.size()));
        gpu_tile_centre = CudaDevicePtr<int32_t>(tile_centre.size());
    }
    for (size_t i = 0; i < x.size(); i++) {
        tile_centre[2 * i] = x[i];
        tile_centre[2 * i + 1] = y[i];
    }
    cuda_err(cudaMemcpyAsync(gpu_tile_centre.get(), tile_centre.data(), 2 * x.size() * sizeof(int32_t),
                             cudaMemcpyHostToDevice, gather_stream));
    const int blocks = static_cast<int>(std::min<size_t>((count + 255) / 256, 4096));
    gather_tiles_kernel<<<blocks, 256, 0, gather_stream>>>(gpu_image.get(), width, static_cast<int>(npixels / width),
                                                           gpu_tile_centre.get(), half, gpu_tiles.get(), count);
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(host_tiles.data(), gpu_tiles.get(), count * sizeof(int32_t),
                             cudaMemcpyDeviceToHost, gather_stream));
    cuda_err(cudaStreamSynchronize(gather_stream));
    return host_tiles.data();
}

const int32_t *ImagePreprocessorBufferGPU::HostImage() {
    buffer.resize(npixels);
    cuda_err(cudaMemcpyAsync(buffer.data(), gpu_image.get(), npixels * sizeof(int32_t),
                             cudaMemcpyDeviceToHost, gather_stream));
    cuda_err(cudaStreamSynchronize(gather_stream));
    return buffer.data();
}