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

// centre holds six values per tile: its centre x, y and the box x0, y0, x1, y1 it is clipped to.
__global__ void gather_tiles_kernel(const int32_t *__restrict__ image, int width,
                                    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 int32_t *c = centre + 6 * tile;
        const int px = c[0] - half + p % side;
        const int py = c[1] - half + p / side;
        tiles[i] = (px < c[2] || py < c[3] || px >= c[4] || py >= c[5])
                   ? 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 (6 * x.size() > tile_centre.size()) {
        tile_centre.resize(std::max(6 * x.size(), 2 * tile_centre.size()));
        gpu_tile_centre = CudaDevicePtr<int32_t>(tile_centre.size());
    }
    for (size_t i = 0; i < x.size(); i++) {
        const PanelBox b = TileBox(x[i], y[i], width);
        int32_t *c = tile_centre.data() + 6 * i;
        c[0] = x[i];
        c[1] = y[i];
        c[2] = b.x0;
        c[3] = b.y0;
        c[4] = b.x1;
        c[5] = b.y1;
    }
    cuda_err(cudaMemcpyAsync(gpu_tile_centre.get(), tile_centre.data(), 6 * 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,
                                                           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();
}