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

#include "ImagePreprocessorBufferGPU.h"
#include "PreprocessedPixel.h"

__global__ void gather_kernel(PixelView image,
                              const uint32_t *__restrict__ npixel,
                              int32_t *__restrict__ values,
                              int count) {
    // Hands back int32 in the pipeline's convention whatever width the image is stored in, so the
    // spot finder that asks for these values sees no difference.
    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, 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),
      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;
}

const uint16_t *ImagePreprocessorBufferGPU::getGPUBufferNarrow() const {
    // The same allocation, read as 16-bit. It is sized for the wide form, so the narrow image uses
    // its first half and there is nothing to allocate when a run turns out to be 16-bit.
    return reinterpret_cast<const uint16_t *>(static_cast<const int32_t *>(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>>>(
        ViewOf(*this), 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);
}