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

#include "ImagePreprocessorBufferGPU.h"

__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]];
}

ImagePreprocessorBufferGPU::ImagePreprocessorBufferGPU(size_t npixel)
    : ImagePreprocessorBuffer(npixel),
      gpu_image(npixel),
      buffer_reg(buffer),
      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);
    cudaMemcpyAsync(values.data(), gpu_gather_value.get(), count * sizeof(int32_t),
                    cudaMemcpyDeviceToHost, gather_stream);
    cudaStreamSynchronize(gather_stream);
}