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

#include <climits>

#include "ROIIntegrationGPU.h"
#include "../../common/DiffractionExperiment.h"

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

// One pixel carries a 16-bit mask, so it can feed any subset of the ROIs.
// Each block reduces into shared memory first to keep global atomics low.
__global__
void gpu_roi(
    const uint16_t *__restrict__ roi_map,
    const int32_t  *__restrict__ input_buffer,
    size_t num_pixels,
    size_t width,
    int roi_count,
    unsigned long long *__restrict__ roi_sum,
    unsigned long long *__restrict__ roi_sum2,
    unsigned long long *__restrict__ roi_pixels,
    unsigned long long *__restrict__ roi_x_weighted,
    unsigned long long *__restrict__ roi_y_weighted,
    int *__restrict__ roi_max) {
    extern __shared__ unsigned long long shared[];
    unsigned long long *s_sum    = shared;
    unsigned long long *s_sum2   = &s_sum[roi_count];
    unsigned long long *s_pixels = &s_sum2[roi_count];
    unsigned long long *s_xw     = &s_pixels[roi_count];
    unsigned long long *s_yw     = &s_xw[roi_count];
    int *s_max = (int *) &s_yw[roi_count];

    for (int r = threadIdx.x; r < roi_count; r += blockDim.x) {
        s_sum[r] = 0;
        s_sum2[r] = 0;
        s_pixels[r] = 0;
        s_xw[r] = 0;
        s_yw[r] = 0;
        s_max[r] = INT_MIN;
    }
    __syncthreads();

    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
         idx < num_pixels;
         idx += blockDim.x * gridDim.x) {
        const uint16_t mask = roi_map[idx];
        if (mask == 0)
            continue;

        const int32_t v = input_buffer[idx];
        if (v == INT32_MIN)        // masked/bad pixel
            continue;
        const bool saturated = (v == INT32_MAX);

        const long long val = v;
        const long long x = idx % width;
        const long long y = idx / width;
        const unsigned long long val_u  = (unsigned long long) val;
        const unsigned long long val2_u = (unsigned long long) (val * val);
        const unsigned long long vx_u   = (unsigned long long) (val * x);
        const unsigned long long vy_u   = (unsigned long long) (val * y);

        for (int r = 0; r < roi_count; r++) {
            if (!(mask & (1u << r)))
                continue;
            if (!saturated) {
                atomicAdd(&s_sum[r], val_u);
                atomicAdd(&s_sum2[r], val2_u);
                atomicAdd(&s_pixels[r], 1ULL);
                atomicAdd(&s_xw[r], vx_u);
                atomicAdd(&s_yw[r], vy_u);
            }
            atomicMax(&s_max[r], v);
        }
    }
    __syncthreads();

    for (int r = threadIdx.x; r < roi_count; r += blockDim.x) {
        atomicAdd(&roi_sum[r], s_sum[r]);
        atomicAdd(&roi_sum2[r], s_sum2[r]);
        atomicAdd(&roi_pixels[r], s_pixels[r]);
        atomicAdd(&roi_x_weighted[r], s_xw[r]);
        atomicAdd(&roi_y_weighted[r], s_yw[r]);
        atomicMax(&roi_max[r], s_max[r]);
    }
}

ROIIntegrationGPU::ROIIntegrationGPU(const DiffractionExperiment &experiment, std::shared_ptr<CudaStream> stream)
    : ROIIntegration(experiment),
      stream(stream),
      gpu_roi_map(npixel),
      gpu_sum(roi_count),
      gpu_sum2(roi_count),
      gpu_pixels(roi_count),
      gpu_x_weighted(roi_count),
      gpu_y_weighted(roi_count),
      gpu_max(roi_count),
      host_sum(roi_count),
      host_sum2(roi_count),
      host_pixels(roi_count),
      host_x_weighted(roi_count),
      host_y_weighted(roi_count),
      host_max(roi_count),
      max_init(roi_count, INT_MIN) {

    cudaDeviceProp prop{};
    cuda_err(cudaGetDeviceProperties(&prop, 0));
    threads = 128;
    blocks = 4 * prop.multiProcessorCount;
    shared_needed = roi_count * (5 * sizeof(unsigned long long) + sizeof(int));

    // On this engine's stream, like every other operation it issues: the streams are non-blocking, so a
    // NULL-stream copy is no longer ordered against the kernels that read the map. The one-time
    // synchronise leaves the constructor with the upload settled rather than in flight.
    cuda_err(cudaMemcpyAsync(gpu_roi_map, roi_map.data(), sizeof(uint16_t) * npixel,
                             cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaStreamSynchronize(*stream));
}

void ROIIntegrationGPU::Run(const ImagePreprocessorBuffer &image, std::map<std::string, ROIMessage> &out) {
    if (image.size() != npixel)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "ROIIntegration: mismatch in image size");

    cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(unsigned long long) * roi_count, *stream));
    cuda_err(cudaMemsetAsync(gpu_sum2, 0, sizeof(unsigned long long) * roi_count, *stream));
    cuda_err(cudaMemsetAsync(gpu_pixels, 0, sizeof(unsigned long long) * roi_count, *stream));
    cuda_err(cudaMemsetAsync(gpu_x_weighted, 0, sizeof(unsigned long long) * roi_count, *stream));
    cuda_err(cudaMemsetAsync(gpu_y_weighted, 0, sizeof(unsigned long long) * roi_count, *stream));
    cuda_err(cudaMemcpyAsync(gpu_max, max_init.data(), sizeof(int) * roi_count, cudaMemcpyHostToDevice, *stream));

    gpu_roi<<<blocks, threads, shared_needed, *stream>>>(
        gpu_roi_map, image.getGPUBuffer(), npixel, width, roi_count,
        gpu_sum, gpu_sum2, gpu_pixels, gpu_x_weighted, gpu_y_weighted, gpu_max);

    cudaMemcpyAsync(host_sum.data(), gpu_sum, sizeof(unsigned long long) * roi_count, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_sum2.data(), gpu_sum2, sizeof(unsigned long long) * roi_count, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_pixels.data(), gpu_pixels, sizeof(unsigned long long) * roi_count, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_x_weighted.data(), gpu_x_weighted, sizeof(unsigned long long) * roi_count, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_y_weighted.data(), gpu_y_weighted, sizeof(unsigned long long) * roi_count, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_max.data(), gpu_max, sizeof(int) * roi_count, cudaMemcpyDeviceToHost, *stream);
    cuda_err(cudaStreamSynchronize(*stream));

    for (uint16_t r = 0; r < roi_count; r++) {
        roi_sum[r] = static_cast<int64_t>(host_sum[r]);
        roi_sum2[r] = host_sum2[r];
        roi_pixels[r] = host_pixels[r];
        roi_x_weighted[r] = static_cast<int64_t>(host_x_weighted[r]);
        roi_y_weighted[r] = static_cast<int64_t>(host_y_weighted[r]);
        roi_max[r] = (host_max[r] == INT_MIN) ? INT64_MIN : static_cast<int64_t>(host_max[r]);
    }
    Export(out);
}