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

#include "AdaptiveSpotFinderGPU.h"
#include "AdaptiveThreshold.h"
#include "../../common/JFJochException.h"

namespace {

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

// One ring reduction, staging per-ring sums in shared memory (fast path). Shared layout:
//   [ sum(float) | sum2(float) | count(uint32) | sum_corr(float) | sum2_corr(float) ]  x nbins
// The corrected arrays exist only when accumulate_corrected is true (the plain first pass); on the
// sigma-clip passes only the first three are launched/used.
__global__ void reduce_rings_shared(
    const uint16_t *__restrict__ pixel_to_bin,
    const float *__restrict__ corrections,
    const int32_t *__restrict__ image,
    const float *__restrict__ mean,
    const float *__restrict__ sigma,
    float clip_k,
    bool accumulate_corrected,
    double *__restrict__ sum, double *__restrict__ sum2, uint32_t *__restrict__ count,
    float *__restrict__ sum_corr, float *__restrict__ sum2_corr,
    size_t npix, int nbins) {

    // The per-block staging stays float: a block contributes only a few dozen pixels to a given ring,
    // all of similar magnitude, so there is nothing to lose there - and float keeps the shared footprint
    // (and hence the occupancy) of the hot loop unchanged. The precision that matters is in the sum over
    // ALL blocks and in the cancelling difference sum2/n - m^2 that follows it, so those are double.
    extern __shared__ float sh[];
    float *s_sum = sh;
    float *s_sum2 = &s_sum[nbins];
    uint32_t *s_count = reinterpret_cast<uint32_t *>(&s_sum2[nbins]);
    float *s_sum_corr = reinterpret_cast<float *>(&s_count[nbins]);
    float *s_sum2_corr = &s_sum_corr[nbins];

    for (int i = threadIdx.x; i < nbins; i += blockDim.x) {
        s_sum[i] = 0.0f;
        s_sum2[i] = 0.0f;
        s_count[i] = 0;
        if (accumulate_corrected) {
            s_sum_corr[i] = 0.0f;
            s_sum2_corr[i] = 0.0f;
        }
    }
    __syncthreads();

    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) {
        const int32_t v = image[idx];
        if (v == INT32_MIN || v == INT32_MAX) continue;
        const uint16_t b = pixel_to_bin[idx];
        if (b >= nbins) continue;
        const float fv = static_cast<float>(v);
        if (clip_k > 0.0f) {
            const float lo = mean[b] - clip_k * sigma[b];
            const float hi = mean[b] + clip_k * sigma[b];
            if (fv < lo || fv > hi) continue;
        }
        atomicAdd(&s_sum[b], fv);
        atomicAdd(&s_sum2[b], fv * fv);
        atomicAdd(&s_count[b], 1u);
        if (accumulate_corrected) {
            const float cv = fv * corrections[idx];
            atomicAdd(&s_sum_corr[b], cv);
            atomicAdd(&s_sum2_corr[b], cv * cv);
        }
    }
    __syncthreads();

    for (int i = threadIdx.x; i < nbins; i += blockDim.x) {
        atomicAdd(&sum[i], static_cast<double>(s_sum[i]));
        atomicAdd(&sum2[i], static_cast<double>(s_sum2[i]));
        atomicAdd(&count[i], s_count[i]);
        if (accumulate_corrected) {
            atomicAdd(&sum_corr[i], s_sum_corr[i]);
            atomicAdd(&sum2_corr[i], s_sum2_corr[i]);
        }
    }
}

// Same reduction with direct global atomics (used only when nbins is too large to stage in shared
// memory - a rare, high-bin-count configuration).
__global__ void reduce_rings_global(
    const uint16_t *__restrict__ pixel_to_bin,
    const float *__restrict__ corrections,
    const int32_t *__restrict__ image,
    const float *__restrict__ mean,
    const float *__restrict__ sigma,
    float clip_k,
    bool accumulate_corrected,
    double *__restrict__ sum, double *__restrict__ sum2, uint32_t *__restrict__ count,
    float *__restrict__ sum_corr, float *__restrict__ sum2_corr,
    size_t npix, int nbins) {

    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) {
        const int32_t v = image[idx];
        if (v == INT32_MIN || v == INT32_MAX) continue;
        const uint16_t b = pixel_to_bin[idx];
        if (b >= nbins) continue;
        const float fv = static_cast<float>(v);
        if (clip_k > 0.0f) {
            const float lo = mean[b] - clip_k * sigma[b];
            const float hi = mean[b] + clip_k * sigma[b];
            if (fv < lo || fv > hi) continue;
        }
        const double dv = static_cast<double>(v);
        atomicAdd(&sum[b], dv);
        atomicAdd(&sum2[b], dv * dv);
        atomicAdd(&count[b], 1u);
        if (accumulate_corrected) {
            const float cv = fv * corrections[idx];
            atomicAdd(&sum_corr[b], cv);
            atomicAdd(&sum2_corr[b], cv * cv);
        }
    }
}

// Per-ring mean/sigma from the current raw accumulators. Rings with no pixels this pass keep their
// previous value (matches the CPU, which leaves ring_mean/ring_sigma untouched when the count is 0).
__global__ void finalize_rings(const double *__restrict__ sum, const double *__restrict__ sum2,
                               const uint32_t *__restrict__ count,
                               float *__restrict__ mean, float *__restrict__ sigma, int nbins) {
    for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < nbins; b += blockDim.x * gridDim.x) {
        if (count[b] > 0) {
            // In double, then rounded to float for the clip predicate - the same two steps, in the same
            // order and the same types, as AdaptiveSpotFinderCPU::AccumulateRings.
            const double m = sum[b] / count[b];
            const double var = fmax(0.0, sum2[b] / count[b] - m * m);
            mean[b] = static_cast<float>(m);
            sigma[b] = static_cast<float>(sqrt(var));
        }
    }
}

// Flag strong pixels (value >= ring threshold, or saturated) into the packed bit buffer. Strong
// pixels are sparse, so a plain atomicOr per strong pixel is simpler than warp aggregation and the
// contention is negligible. Mirrors AdaptiveSpotFinderCPU Stage C exactly.
__global__ void flag_strong(const int32_t *__restrict__ image,
                            const uint16_t *__restrict__ pixel_to_bin,
                            const float *__restrict__ thr,
                            uint32_t *__restrict__ strong,
                            size_t npix, int nbins) {
    for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) {
        const int32_t v = image[idx];
        bool s = false;
        if (v == INT32_MAX) {
            s = true;
        } else if (v != INT32_MIN) {
            const uint16_t b = pixel_to_bin[idx];
            if (b < nbins && static_cast<float>(v) >= thr[b])
                s = true;
        }
        if (s)
            atomicOr(&strong[idx / 32], 1u << (idx % 32));
    }
}

} // namespace

AdaptiveSpotFinderGPU::AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &in_mapping,
                                             std::shared_ptr<CudaStream> in_stream)
    : ImageSpotFinder(static_cast<int32_t>(in_mapping.GetWidth()),
                      static_cast<int32_t>(in_mapping.GetHeight())),
      mapping(in_mapping),
      stream(std::move(in_stream)),
      nbins(in_mapping.GetBinNumber()),
      npix(in_mapping.GetPixelToBin().size()),
      gpu_pixel_to_bin(npix),
      gpu_corrections(npix),
      gpu_sum(nbins),
      gpu_sum2(nbins),
      gpu_count(nbins),
      gpu_mean(nbins),
      gpu_sigma(nbins),
      gpu_sum_corr(nbins),
      gpu_sum2_corr(nbins),
      gpu_thr(nbins),
      gpu_strong(OutputSize()),
      host_sum(nbins),
      host_sum2(nbins),
      host_count(nbins),
      prof_sum(nbins),
      prof_sum2(nbins),
      prof_count(nbins),
      output_buffer_reg(output_buffer),
      last_profile(in_mapping) {

    // The current device, not device 0: callers round-robin engines across GPUs, so device 0's shared
    // memory and 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));
    reduce_blocks = 4 * prop.multiProcessorCount;
    flag_blocks = 4 * prop.multiProcessorCount;

    shared_plain = static_cast<size_t>(nbins) * (4 * sizeof(float) + sizeof(uint32_t));
    shared_clip = static_cast<size_t>(nbins) * (2 * sizeof(float) + sizeof(uint32_t));
    use_shared = (shared_plain < prop.sharedMemPerBlock);

    cuda_err(cudaMemcpy(gpu_pixel_to_bin, mapping.GetPixelToBin().data(), sizeof(uint16_t) * npix,
                        cudaMemcpyHostToDevice));
    cuda_err(cudaMemcpy(gpu_corrections, mapping.Corrections().data(), sizeof(float) * npix,
                        cudaMemcpyHostToDevice));
}

void AdaptiveSpotFinderGPU::ReducePass(const ImagePreprocessorBuffer &image, float clip_k,
                                       bool accumulate_corrected) {
    if (use_shared) {
        const size_t shared = accumulate_corrected ? shared_plain : shared_clip;
        reduce_rings_shared<<<reduce_blocks, reduce_threads, shared, *stream>>>(
            gpu_pixel_to_bin, gpu_corrections, image.getGPUBuffer(), gpu_mean, gpu_sigma,
            clip_k, accumulate_corrected, gpu_sum, gpu_sum2, gpu_count, gpu_sum_corr, gpu_sum2_corr,
            npix, nbins);
    } else {
        reduce_rings_global<<<reduce_blocks, reduce_threads, 0, *stream>>>(
            gpu_pixel_to_bin, gpu_corrections, image.getGPUBuffer(), gpu_mean, gpu_sigma,
            clip_k, accumulate_corrected, gpu_sum, gpu_sum2, gpu_count, gpu_sum_corr, gpu_sum2_corr,
            npix, nbins);
    }
}

void AdaptiveSpotFinderGPU::FinalizeStats() {
    const int threads = 128;
    const int blocks = (nbins + threads - 1) / threads;
    finalize_rings<<<blocks, threads, 0, *stream>>>(gpu_sum, gpu_sum2, gpu_count, gpu_mean, gpu_sigma, nbins);
}

// Host reproduction of AdaptiveSpotFinderCPU Stage B, from the clipped raw per-ring stats.
void AdaptiveSpotFinderGPU::ComputeThresholds(const SpotFindingSettings &settings) {
    int64_t n_total = 0;
    double g_sum = 0.0, g_sum2 = 0.0;
    for (int b = 0; b < nbins; ++b) {
        n_total += host_count[b];
        g_sum += host_sum[b];
        g_sum2 += host_sum2[b];
    }
    if (n_total == 0) {
        host_thr.clear();
        return;
    }

    const double E = std::max(1.0f, settings.false_pixels_per_frame);
    double p = E / static_cast<double>(n_total);
    p = std::min(std::max(p, 1e-9), 0.1);
    const float z = static_cast<float>(adaptive_threshold::NormalQuantile(1.0 - p));

    const double g_mean = g_sum / n_total;
    const double g_sigma = std::sqrt(std::max(0.0, g_sum2 / n_total - g_mean * g_mean));
    const float g_thr = adaptive_threshold::RingThreshold(static_cast<float>(g_mean),
                                                          static_cast<float>(g_sigma), p, z);

    host_thr.assign(nbins, 0.0f);
    for (int b = 0; b < nbins; ++b) {
        if (host_count[b] < adaptive_threshold::MIN_RING_PIXELS) {
            host_thr[b] = g_thr;
        } else {
            const double m = host_sum[b] / host_count[b];
            const double var = std::max(0.0, host_sum2[b] / host_count[b] - m * m);
            host_thr[b] = adaptive_threshold::RingThreshold(static_cast<float>(m),
                                                           static_cast<float>(std::sqrt(var)), p, z);
        }
    }
}

void AdaptiveSpotFinderGPU::Detect(const ImagePreprocessorBuffer &image,
                                   const SpotFindingSettings &settings) {
    if (image.size() != npix)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "AdaptiveSpotFinderGPU::Detect: mismatch in pixel size");

    // --- Stage A: robust per-ring background (one plain pass + two sigma-clip passes) ---
    cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(double) * nbins, *stream));
    cuda_err(cudaMemsetAsync(gpu_sum2, 0, sizeof(double) * nbins, *stream));
    cuda_err(cudaMemsetAsync(gpu_count, 0, sizeof(uint32_t) * nbins, *stream));
    cuda_err(cudaMemsetAsync(gpu_mean, 0, sizeof(float) * nbins, *stream));
    cuda_err(cudaMemsetAsync(gpu_sigma, 0, sizeof(float) * nbins, *stream));
    cuda_err(cudaMemsetAsync(gpu_sum_corr, 0, sizeof(float) * nbins, *stream));
    cuda_err(cudaMemsetAsync(gpu_sum2_corr, 0, sizeof(float) * nbins, *stream));

    ReducePass(image, 0.0f, true);   // plain pass also fills the corrected profile accumulators
    FinalizeStats();

    // Snapshot the plain corrected profile (and its pixel count) before the raw accumulators are
    // re-zeroed for the sigma-clip passes.
    cuda_err(cudaMemcpyAsync(prof_sum.data(), gpu_sum_corr, sizeof(float) * nbins, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(prof_sum2.data(), gpu_sum2_corr, sizeof(float) * nbins, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(prof_count.data(), gpu_count, sizeof(uint32_t) * nbins, cudaMemcpyDeviceToHost, *stream));

    for (int pass = 0; pass < 2; ++pass) {
        cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(double) * nbins, *stream));
        cuda_err(cudaMemsetAsync(gpu_sum2, 0, sizeof(double) * nbins, *stream));
        cuda_err(cudaMemsetAsync(gpu_count, 0, sizeof(uint32_t) * nbins, *stream));
        ReducePass(image, 3.0f, false);
        FinalizeStats();
    }

    // Snapshot the clipped raw stats that drive the threshold.
    cuda_err(cudaMemcpyAsync(host_sum.data(), gpu_sum, sizeof(double) * nbins, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(host_sum2.data(), gpu_sum2, sizeof(double) * nbins, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(host_count.data(), gpu_count, sizeof(uint32_t) * nbins, cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaStreamSynchronize(*stream));

    // --- Stage B: per-ring threshold on the host (shared with the CPU finder) ---
    ComputeThresholds(settings);

    // The profile is a byproduct even when the frame has no valid pixels for detection.
    last_profile.Clear(mapping);
    last_profile.Add(prof_sum, prof_sum2, prof_count);

    if (host_thr.empty()) {
        // Nothing valid to threshold against: leave no strong pixels for ExtractSpots to build on.
        std::fill(output_buffer.begin(), output_buffer.end(), 0);
        return;
    }

    // --- Stage C: flag strong pixels into the bit buffer (value >= ring threshold) ---
    cuda_err(cudaMemcpyAsync(gpu_thr, host_thr.data(), sizeof(float) * nbins, cudaMemcpyHostToDevice, *stream));
    cuda_err(cudaMemsetAsync(gpu_strong, 0, OutputByteSize(), *stream));
    flag_strong<<<flag_blocks, flag_threads, 0, *stream>>>(
        image.getGPUBuffer(), gpu_pixel_to_bin, gpu_thr, gpu_strong, npix, nbins);
    cuda_err(cudaMemcpyAsync(output_buffer.data(), gpu_strong, OutputByteSize(), cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaStreamSynchronize(*stream));
}