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

__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,
    const float *__restrict__ clip_lo,
    const float *__restrict__ clip_hi,
    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();

    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];
            // clip_lo is null on the plain pass; on a clip pass it is the previous pass's
            // mean +- n sigma for this bin.
            if (clip_lo == nullptr || (val >= clip_lo[bin] && val <= clip_hi[bin])) {
                const float val2 = val * val;
                atomicAdd(&s_sum[bin], val);
                atomicAdd(&s_sum2[bin], val2);
                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,
    const float *__restrict__ clip_lo,
    const float *__restrict__ clip_hi,
    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];
            if (clip_lo == nullptr || (val >= clip_lo[bin] && val <= clip_hi[bin])) {
                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),
      gpu_clip_lo(clip_nsigma > 0.0f ? azint_bins : 0),
      gpu_clip_hi(clip_nsigma > 0.0f ? azint_bins : 0),
      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(), *stream);
    gpu_pixel_to_bin = SharedDeviceTable(integration.GetPixelToBin().data(), npixel,
                                         integration.GetPixelToBin().data(), *stream);
}

// Per-bin accept range from the accumulators of the pass just finished; mirrors
// AzIntEngine::UpdateClipLimits, kept on the device so a clip pass costs no round trip.
__global__
void gpu_azim_clip_limits(
    const float *__restrict__ azint_sum,
    const float *__restrict__ azint_sum2,
    const uint32_t *__restrict__ azint_count,
    float *__restrict__ clip_lo,
    float *__restrict__ clip_hi,
    int azint_bins,
    float nsigma) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < azint_bins; i += blockDim.x * gridDim.x) {
        const uint32_t n = azint_count[i];
        if (n < 8) {
            clip_lo[i] = -INFINITY;
            clip_hi[i] = INFINITY;
            continue;
        }
        const double dn = n;
        const double mean = azint_sum[i] / dn;
        const double sd = sqrt(fmax(0.0, azint_sum2[i] / dn - mean * mean));
        clip_lo[i] = static_cast<float>(mean - nsigma * sd);
        clip_hi[i] = static_cast<float>(mean + nsigma * sd);
    }
}

void AzIntEngineGPU::Run(const ImagePreprocessorBuffer &image, AzimuthalIntegrationProfile &profile) {
    if (image.size() != integration.GetPixelToBin().size())
        throw std::runtime_error("ImageSpotFinder::AzimIntegration: Mismatch in size");

    // Pass 0 accumulates every valid pixel; each later pass repeats it with the accept range the pass
    // before measured, so only the last pass's accumulators reach the profile. The limits stay on the
    // device between passes - a clip pass is one more read of the same pixels and nothing else.
    const int passes = PassCount();
    for (int pass = 0; pass < passes; ++pass) {
        if (pass > 0)
            gpu_azim_clip_limits<<<blocks, threads, 0, *stream>>>(
                gpu_sum, gpu_sum2, gpu_count, gpu_clip_lo, gpu_clip_hi, azint_bins, clip_nsigma);
        const float *lo = pass > 0 ? gpu_clip_lo.get() : nullptr;
        const float *hi = pass > 0 ? gpu_clip_hi.get() : nullptr;

        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, lo, hi, npixel, azint_bins
            );
        } else {
            gpu_azim<<<blocks, threads, 0, *stream>>>(
                gpu_pixel_to_bin->get(),gpu_azint_correction->get(),image.getGPUBuffer(), gpu_sum, gpu_sum2,
                gpu_count, lo, hi, npixel, azint_bins
            );
        }
    }

    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);
}