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

#include "AzIntEngineGPU.h"

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

// Pushes one ring run's totals to the shared accumulators; nothing for an empty run.
__device__ __forceinline__ void flush_azim_run(float *s_sum, float *s_sum2, uint32_t *s_count,
                                               int b, float r_sum, float r_sum2, uint32_t r_count) {
    if (r_count == 0) return;   // also covers the initial "no ring yet"
    atomicAdd(&s_sum[b], r_sum);
    atomicAdd(&s_sum2[b], r_sum2);
    atomicAdd(&s_count[b], r_count);
}

__global__
void gpu_azim_shared(
    const uint16_t *__restrict__ pixel_to_bin,
    const float *__restrict__ corrections,
    const int32_t *__restrict__ input_buffer,
    float *__restrict__ azint_sum,
    float *__restrict__ azint_sum2,
    uint32_t *__restrict__ azint_count,
    size_t num_pixels,
    int azint_bins) {
    extern __shared__ float shared[];

    float *s_sum = shared;
    float *s_sum2 = &s_sum[azint_bins];
    uint32_t *s_count = (uint32_t *) &s_sum2[azint_bins];

    // Initialize shared memory
    for (int i = threadIdx.x; i < azint_bins; i += blockDim.x) {
        s_sum[i] = 0.0f;
        s_sum2[i] = 0.0f;
        s_count[i] = 0;
    }

    __syncthreads();

    // Four pixels per thread, read as vector loads, and a running total per ring pushed to shared
    // memory only when the ring changes: consecutive pixels along a row mostly share a ring, and an
    // atomic per pixel on the same few addresses is what this kernel was limited by. The same scheme
    // as the adaptive spot finder's ring pass (reduce_rings_shared). The buffers come straight from
    // cudaMalloc, aligned for int4/float4; the npix % 4 leftovers are done one at a time below.
    const size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;
    const size_t nquad = num_pixels / 4;
    for (size_t q = blockIdx.x * blockDim.x + threadIdx.x; q < nquad; q += stride) {
        const int4 v4 = reinterpret_cast<const int4 *>(input_buffer)[q];
        const ushort4 b4 = reinterpret_cast<const ushort4 *>(pixel_to_bin)[q];
        const float4 c4 = reinterpret_cast<const float4 *>(corrections)[q];
        const int32_t vq[4] = {v4.x, v4.y, v4.z, v4.w};
        const uint16_t bq[4] = {b4.x, b4.y, b4.z, b4.w};
        const float cq[4] = {c4.x, c4.y, c4.z, c4.w};

        int r_b = -1;
        float r_sum = 0.0f, r_sum2 = 0.0f;
        uint32_t r_count = 0;
        #pragma unroll
        for (int k = 0; k < 4; k++) {
            const int32_t v = vq[k];
            const int b = bq[k];
            if (v == INT32_MIN || v == INT32_MAX || b >= azint_bins) continue;
            if (b != r_b) {
                flush_azim_run(s_sum, s_sum2, s_count, r_b, r_sum, r_sum2, r_count);
                r_b = b;
                r_sum = 0.0f; r_sum2 = 0.0f; r_count = 0;
            }
            const float val = static_cast<float>(v) * cq[k];
            r_sum += val;
            r_sum2 += val * val;
            r_count += 1;
        }
        flush_azim_run(s_sum, s_sum2, s_count, r_b, r_sum, r_sum2, r_count);
    }

    for (size_t idx = 4 * nquad + blockIdx.x * blockDim.x + threadIdx.x; idx < num_pixels; idx += stride) {
        const uint16_t bin = pixel_to_bin[idx];
        const int32_t v = input_buffer[idx];
        if (bin < azint_bins && v != INT32_MIN && v != INT32_MAX) {
            const float val = static_cast<float>(v) * corrections[idx];
            atomicAdd(&s_sum[bin], val);
            atomicAdd(&s_sum2[bin], val * val);
            atomicAdd(&s_count[bin], 1);
        }
    }

    __syncthreads();

    // Merge to global memory
    for (unsigned int i = threadIdx.x; i < azint_bins; i += blockDim.x) {
        atomicAdd(&azint_sum[i], s_sum[i]);
        atomicAdd(&azint_sum2[i], s_sum2[i]);
        atomicAdd(&azint_count[i], s_count[i]);
    }
}

__global__
void gpu_azim(
    const uint16_t *__restrict__ pixel_to_bin,
    const float *__restrict__ corrections,
    const int32_t *__restrict__ input_buffer,
    float *__restrict__ azint_sum,
    float *__restrict__ azint_sum2,
    uint32_t *__restrict__ azint_count,
    size_t num_pixels,
    int azint_bins) {
    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
         idx < num_pixels;
         idx += blockDim.x * gridDim.x) {
        uint16_t bin = pixel_to_bin[idx];

        int32_t v = input_buffer[idx];
        bool valid = (v != INT32_MIN) & (v != INT32_MAX);

        if (bin < azint_bins && valid) {
            const float val = static_cast<float>(v) * corrections[idx];
            const float val2 = val * val;
            atomicAdd(&azint_sum[bin], val);
            atomicAdd(&azint_sum2[bin], val2);
            atomicAdd(&azint_count[bin], 1);
        }
    }
}

AzIntEngineGPU::AzIntEngineGPU(const AzimuthalIntegrationMapping &integration, std::shared_ptr<CudaStream> stream)
    : AzIntEngine(integration),
      stream(stream),
      gpu_sum(azint_bins),
      gpu_sum2(azint_bins),
      gpu_count(azint_bins),
      cpu_sum_reg(azint_sum),
      cpu_sum2_reg(azint_sum2),
      cpu_count_reg(azint_count) {

    int device = 0;
    cuda_err(cudaGetDevice(&device));       // this worker's GPU, not necessarily 0
    cudaDeviceProp prop{};
    cuda_err(cudaGetDeviceProperties(&prop, device));

    threads = 128;
    blocks = 4 * prop.multiProcessorCount;
    shared_size = prop.sharedMemPerBlock;
    shared_needed = azint_bins * (2 * sizeof(float) + sizeof(uint32_t));

    // Geometry-only, so shared per GPU: the first engine on this device uploads them, the rest reuse
    // them. Keyed by the mapping's own vectors, which outlive every engine built from it.
    gpu_azint_correction = SharedDeviceTable(integration.Corrections().data(), npixel,
                                             integration.Corrections().data(),
                                             integration.GetCorrectionsChecksum(), *stream);
    gpu_pixel_to_bin = SharedDeviceTable(integration.GetPixelToBin().data(), npixel,
                                         integration.GetPixelToBin().data(),
                                         integration.GetPixelToBinChecksum(), *stream);
}

void AzIntEngineGPU::Run(const ImagePreprocessorBuffer &image, AzimuthalIntegrationProfile &profile) {
    if (image.size() != integration.GetPixelToBin().size())
        throw std::runtime_error("ImageSpotFinder::AzimIntegration: Mismatch in size");
    cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(float) * azint_bins, *stream));
    cuda_err(cudaMemsetAsync(gpu_sum2, 0, sizeof(float) * azint_bins, *stream));
    cuda_err(cudaMemsetAsync(gpu_count, 0, sizeof(uint32_t) * azint_bins, *stream));

    if (shared_needed < shared_size) {
        gpu_azim_shared<<<blocks, threads, shared_needed, *stream>>>(
        gpu_pixel_to_bin->get(),gpu_azint_correction->get(),image.getGPUBuffer(), gpu_sum, gpu_sum2,
        gpu_count, npixel, azint_bins
        );
        cuda_err(cudaGetLastError());
    } else {
        gpu_azim<<<blocks, threads, 0, *stream>>>(
            gpu_pixel_to_bin->get(),gpu_azint_correction->get(),image.getGPUBuffer(), gpu_sum, gpu_sum2,
            gpu_count, npixel, azint_bins
        );
        cuda_err(cudaGetLastError());
    }

    cudaMemcpyAsync(azint_sum.data(), gpu_sum, sizeof(float) * azint_bins, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(azint_sum2.data(), gpu_sum2, sizeof(float) * azint_bins, cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(azint_count.data(), gpu_count, sizeof(uint32_t) * azint_bins, cudaMemcpyDeviceToHost, *stream);
    cuda_err(cudaStreamSynchronize(*stream));

    profile.Clear(integration);
    profile.Add(azint_sum, azint_sum2, azint_count);
}