// 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.
// The weighted sums go per ROI and panel (roi * slots + panel; slots = 1 and panel 0 without a panel
// table); panel_box holds x0, y0, width, height of each panel, and a pixel on none is left out.
__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,
    const int *__restrict__ panel_box,
    int npanels,
    int slots,
    unsigned long long *__restrict__ roi_sum,
    unsigned long long *__restrict__ roi_sum2,
    unsigned long long *__restrict__ roi_pixels,
    unsigned long long *__restrict__ roi_panel_v,
    unsigned long long *__restrict__ roi_panel_x,
    unsigned long long *__restrict__ roi_panel_y,
    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_pv     = &s_pixels[roi_count];
    unsigned long long *s_xw     = &s_pv[roi_count * slots];
    unsigned long long *s_yw     = &s_xw[roi_count * slots];
    int *s_max = (int *) &s_yw[roi_count * slots];

    for (int r = threadIdx.x; r < roi_count; r += blockDim.x) {
        s_sum[r] = 0;
        s_sum2[r] = 0;
        s_pixels[r] = 0;
        s_max[r] = INT_MIN;
    }
    for (int k = threadIdx.x; k < roi_count * slots; k += blockDim.x) {
        s_pv[k] = 0;
        s_xw[k] = 0;
        s_yw[k] = 0;
    }
    __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;
        int panel = 0;
        if (npanels > 0) {
            panel = -1;
            for (int p = 0; p < npanels && panel < 0; p++)
                if (x >= panel_box[4 * p] && x < panel_box[4 * p] + panel_box[4 * p + 2]
                    && y >= panel_box[4 * p + 1] && y < panel_box[4 * p + 1] + panel_box[4 * p + 3])
                    panel = p;
            if (panel < 0)
                continue;
        }
        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_pv[r * slots + panel], val_u);
                atomicAdd(&s_xw[r * slots + panel], vx_u);
                atomicAdd(&s_yw[r * slots + panel], 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]);
        atomicMax(&roi_max[r], s_max[r]);
    }
    for (int k = threadIdx.x; k < roi_count * slots; k += blockDim.x) {
        atomicAdd(&roi_panel_v[k], s_pv[k]);
        atomicAdd(&roi_panel_x[k], s_xw[k]);
        atomicAdd(&roi_panel_y[k], s_yw[k]);
    }
}

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_panel_v(roi_count * panel_slots),
      gpu_panel_x(roi_count * panel_slots),
      gpu_panel_y(roi_count * panel_slots),
      gpu_max(roi_count),
      host_sum(roi_count),
      host_sum2(roi_count),
      host_pixels(roi_count),
      host_panel_v(roi_count * panel_slots),
      host_panel_x(roi_count * panel_slots),
      host_panel_y(roi_count * panel_slots),
      host_max(roi_count),
      max_init(roi_count, INT_MIN) {

    // The current device, not device 0: workers are pinned round-robin across the GPUs, so device 0's
    // SM count can belong to a different card than the one these kernels launch on.
    int device = 0;
    cuda_err(cudaGetDevice(&device));
    cudaDeviceProp prop{};
    cuda_err(cudaGetDeviceProperties(&prop, device));
    threads = 128;
    blocks = 4 * prop.multiProcessorCount;
    shared_needed = roi_count * (3 * sizeof(unsigned long long) + sizeof(int))
                    + roi_count * panel_slots * 3 * sizeof(unsigned long long);

    std::vector<int> box;
    for (const auto &m: modules.Get())
        box.insert(box.end(), {static_cast<int>(m.x0), static_cast<int>(m.y0), static_cast<int>(m.width),
                               static_cast<int>(m.height)});
    if (!box.empty()) {
        gpu_panel_box = CudaDevicePtr<int>(box.size());
        cuda_err(cudaMemcpyAsync(gpu_panel_box, box.data(), sizeof(int) * box.size(),
                                 cudaMemcpyHostToDevice, *stream));
    }

    // 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));
    const size_t slots = roi_count * panel_slots;
    cuda_err(cudaMemsetAsync(gpu_panel_v, 0, sizeof(unsigned long long) * slots, *stream));
    cuda_err(cudaMemsetAsync(gpu_panel_x, 0, sizeof(unsigned long long) * slots, *stream));
    cuda_err(cudaMemsetAsync(gpu_panel_y, 0, sizeof(unsigned long long) * slots, *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_panel_box, static_cast<int>(modules.size()), static_cast<int>(panel_slots),
        gpu_sum, gpu_sum2, gpu_pixels, gpu_panel_v, gpu_panel_x, gpu_panel_y, gpu_max);
    cuda_err(cudaGetLastError());

    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_panel_v.data(), gpu_panel_v, sizeof(unsigned long long) * slots, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_panel_x.data(), gpu_panel_x, sizeof(unsigned long long) * slots, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(host_panel_y.data(), gpu_panel_y, sizeof(unsigned long long) * slots, 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_max[r] = (host_max[r] == INT_MIN) ? INT64_MIN : static_cast<int64_t>(host_max[r]);
    }
    for (size_t k = 0; k < slots; k++) {
        roi_panel_v[k] = static_cast<int64_t>(host_panel_v[k]);
        roi_panel_x[k] = static_cast<int64_t>(host_panel_x[k]);
        roi_panel_y[k] = static_cast<int64_t>(host_panel_y[k]);
    }
    PlaceWeighted();
    Export(out);
}