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

#include "BeamCenterBackgroundGPU.h"

#include <cub/device/device_radix_sort.cuh>

#include "../indexing/CUDAMemHelpers.h"

namespace {

constexpr int THREADS = 256;

void check(cudaError_t err, const char *what) {
    if (err != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                              std::string("Beam centre from background: ") + what + ": " + cudaGetErrorString(err));
}

// Every pixel's cell (BackgroundBand::CELLS where it is outside the band or unusable), its index,
// and its derivatives.
// A pixel on no panel is outside the band.
__global__ void bin_kernel(BackgroundBand band, int width, size_t npixels, float beam_x, float beam_y,
                           const BackgroundPanel *__restrict__ panels, int npanels,
                           const char *__restrict__ usable, int32_t *__restrict__ key,
                           int32_t *__restrict__ index, float *__restrict__ jac_x,
                           float *__restrict__ jac_y) {
    const size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x;
    if (i >= npixels)
        return;
    int cell = BackgroundBand::CELLS;
    float jx = 0.0f, jy = 0.0f;
    if (usable[i]) {
        const int x = static_cast<int>(i % width), y = static_cast<int>(i / width);
        float px = static_cast<float>(x), py = static_cast<float>(y);
        bool on_panel = npanels == 0;
        for (int p = 0; p < npanels && !on_panel; p++) {
            const BackgroundPanel &m = panels[p];
            if (x >= m.x0 && x < m.x0 + m.width && y >= m.y0 && y < m.y0 + m.height) {
                BackgroundPlanePoint(m, x, y, px, py);
                on_panel = true;
            }
        }
        if (on_panel) {
            const int c = BackgroundBandCell(band, px, py, beam_x, beam_y, jx, jy);
            if (c >= 0)
                cell = c;
        }
    }
    key[i] = cell;
    index[i] = static_cast<int32_t>(i);
    jac_x[i] = jx;
    jac_y[i] = jy;
}

// Where each cell's pixels start in the sorted keys; offset[CELLS] is the number of binned pixels.
__global__ void offset_kernel(const int32_t *__restrict__ sorted_key, size_t n, int32_t *__restrict__ offset) {
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c > BackgroundBand::CELLS)
        return;
    size_t lo = 0, hi = n;
    while (lo < hi) {
        const size_t mid = (lo + hi) / 2;
        if (sorted_key[mid] < c) lo = mid + 1;
        else hi = mid;
    }
    offset[c] = static_cast<int32_t>(lo);
}

// One thread per cell, walking its pixels in pixel order: a partial sum per host row block, added to
// the cell's total when the block changes - the host's arithmetic, step for step. A cell whose pixels
// are clipped (clip_limit != nullptr) skips the ones above its limit, as the host's clip rounds do.
__global__ void sum_kernel(int width, const int32_t *__restrict__ offset, const int32_t *__restrict__ index,
                           const int32_t *__restrict__ row_block, const float *__restrict__ mean,
                           const float *__restrict__ jac_x, const float *__restrict__ jac_y,
                           const float *__restrict__ clip_limit,
                           double *__restrict__ sum, double *__restrict__ sum_sq,
                           double *__restrict__ sum_jx, double *__restrict__ sum_jy,
                           int32_t *__restrict__ count) {
    const int c = blockIdx.x * blockDim.x + threadIdx.x;
    if (c >= BackgroundBand::CELLS)
        return;
    const bool clipped = clip_limit != nullptr;
    const float limit = clipped ? clip_limit[c] : 0.0f;

    double s = 0, ss = 0, jx = 0, jy = 0;
    double bs = 0, bss = 0, bjx = 0, bjy = 0;
    int32_t n = 0;
    int block = -1;
    for (int32_t p = offset[c]; p < offset[c + 1]; p++) {
        const int32_t i = index[p];
        const float value = mean[i];
        if (clipped && (limit < 0.0f || value > limit))
            continue;
        const int b = row_block[i / width];
        if (b != block) {
            s += bs; ss += bss; jx += bjx; jy += bjy;
            bs = bss = bjx = bjy = 0;
            block = b;
        }
        n++;
        bs += value;
        bss += static_cast<double>(value) * value;
        if (!clipped) {
            bjx += jac_x[i];
            bjy += jac_y[i];
        }
    }
    s += bs; ss += bss; jx += bjx; jy += bjy;
    sum[c] = s;
    sum_sq[c] = ss;
    count[c] = n;
    if (!clipped) {
        sum_jx[c] = jx;
        sum_jy[c] = jy;
    }
}

}   // namespace

struct BeamCenterBackgroundGPU::Impl {
    int width, height;
    size_t npixels;
    CudaStream stream;
    CudaDevicePtr<char> usable;
    CudaDevicePtr<float> mean;
    CudaDevicePtr<int32_t> row_block;
    CudaDevicePtr<BackgroundPanel> panels;
    int npanels = 0;
    CudaDevicePtr<int32_t> key, sorted_key, index, sorted_index;
    CudaDevicePtr<float> jac_x, jac_y;
    CudaDevicePtr<int32_t> offset;
    CudaDevicePtr<float> clip_limit;
    CudaDevicePtr<double> sum, sum_sq, sum_jx, sum_jy;
    CudaDevicePtr<int32_t> count;
    CudaDevicePtr<uint8_t> sort_scratch;
    size_t sort_scratch_bytes = 0;

    // Waits for the stream, which is destroyed after the buffers freed before it (see CudaDevicePtr).
    ~Impl() { cudaStreamSynchronize(stream); }

    Impl(int w, int h)
        : width(w), height(h), npixels(static_cast<size_t>(w) * h),
          usable(npixels), mean(npixels), row_block(h),
          key(npixels), sorted_key(npixels), index(npixels), sorted_index(npixels),
          jac_x(npixels), jac_y(npixels), offset(BackgroundBand::CELLS + 1),
          clip_limit(BackgroundBand::CELLS),
          sum(BackgroundBand::CELLS), sum_sq(BackgroundBand::CELLS),
          sum_jx(BackgroundBand::CELLS), sum_jy(BackgroundBand::CELLS),
          count(BackgroundBand::CELLS) {
        // Keys run to CELLS inclusive, the bin of everything outside the band.
        check(cub::DeviceRadixSort::SortPairs(nullptr, sort_scratch_bytes, key.get(), sorted_key.get(),
                                              index.get(), sorted_index.get(), npixels, 0, end_bit(), stream),
              "sort size");
        sort_scratch = CudaDevicePtr<uint8_t>(sort_scratch_bytes);
    }

    static int end_bit() {
        int bits = 0;
        while ((1 << bits) <= BackgroundBand::CELLS) bits++;
        return bits;
    }

    void Download(std::vector<double> &s, std::vector<double> &ss, std::vector<double> *jx,
                  std::vector<double> *jy, std::vector<int32_t> &n) {
        const size_t cells = BackgroundBand::CELLS;
        check(cudaMemcpyAsync(s.data(), sum.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
        check(cudaMemcpyAsync(ss.data(), sum_sq.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
        if (jx) check(cudaMemcpyAsync(jx->data(), sum_jx.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
        if (jy) check(cudaMemcpyAsync(jy->data(), sum_jy.get(), cells * sizeof(double), cudaMemcpyDeviceToHost, stream), "copy");
        check(cudaMemcpyAsync(n.data(), count.get(), cells * sizeof(int32_t), cudaMemcpyDeviceToHost, stream), "copy");
        check(cudaStreamSynchronize(stream), "sums");
    }
};

BeamCenterBackgroundGPU::BeamCenterBackgroundGPU(int width, int height, const std::vector<int> &block_row,
                                                 const std::vector<BackgroundPanel> &panels,
                                                 const char *usable, const float *mean)
    : impl(std::make_unique<Impl>(width, height)) {
    impl->npanels = static_cast<int>(panels.size());
    if (!panels.empty()) {
        impl->panels = CudaDevicePtr<BackgroundPanel>(panels.size());
        check(cudaMemcpyAsync(impl->panels.get(), panels.data(), panels.size() * sizeof(BackgroundPanel),
                              cudaMemcpyHostToDevice, impl->stream), "upload");
    }
    std::vector<int32_t> row_block(height);
    for (size_t b = 0; b + 1 < block_row.size(); b++)
        for (int y = block_row[b]; y < block_row[b + 1]; y++)
            row_block[y] = static_cast<int32_t>(b);
    check(cudaMemcpyAsync(impl->row_block.get(), row_block.data(), height * sizeof(int32_t),
                          cudaMemcpyHostToDevice, impl->stream), "upload");
    check(cudaMemcpyAsync(impl->usable.get(), usable, impl->npixels, cudaMemcpyHostToDevice, impl->stream), "upload");
    check(cudaMemcpyAsync(impl->mean.get(), mean, impl->npixels * sizeof(float), cudaMemcpyHostToDevice,
                          impl->stream), "upload");
    check(cudaStreamSynchronize(impl->stream), "upload");
}

BeamCenterBackgroundGPU::~BeamCenterBackgroundGPU() = default;

void BeamCenterBackgroundGPU::Bin(const BackgroundBand &band, float beam_x, float beam_y,
                                  std::vector<double> &sum, std::vector<double> &sum_sq,
                                  std::vector<double> &sum_jx, std::vector<double> &sum_jy,
                                  std::vector<int32_t> &count) {
    Impl &d = *impl;
    const auto blocks = static_cast<unsigned>((d.npixels + THREADS - 1) / THREADS);
    bin_kernel<<<blocks, THREADS, 0, d.stream>>>(band, d.width, d.npixels, beam_x, beam_y, d.panels.get(),
                                                 d.npanels, d.usable,
                                                 d.key, d.index, d.jac_x, d.jac_y);
    check(cudaGetLastError(), "bin");
    // A radix sort is stable, so each cell's pixels come out in pixel order.
    check(cub::DeviceRadixSort::SortPairs(d.sort_scratch.get(), d.sort_scratch_bytes, d.key.get(),
                                          d.sorted_key.get(), d.index.get(), d.sorted_index.get(),
                                          d.npixels, 0, Impl::end_bit(), d.stream), "sort");
    constexpr int cell_blocks = (BackgroundBand::CELLS + 1 + THREADS - 1) / THREADS;
    offset_kernel<<<cell_blocks, THREADS, 0, d.stream>>>(d.sorted_key, d.npixels, d.offset);
    check(cudaGetLastError(), "offsets");
    sum_kernel<<<cell_blocks, THREADS, 0, d.stream>>>(d.width, d.offset, d.sorted_index, d.row_block, d.mean,
                                                      d.jac_x, d.jac_y, nullptr, d.sum, d.sum_sq,
                                                      d.sum_jx, d.sum_jy, d.count);
    check(cudaGetLastError(), "sums");
    d.Download(sum, sum_sq, &sum_jx, &sum_jy, count);
}

void BeamCenterBackgroundGPU::Clip(const std::vector<float> &clip_limit,
                                   std::vector<double> &sum, std::vector<double> &sum_sq,
                                   std::vector<int32_t> &count) {
    Impl &d = *impl;
    check(cudaMemcpyAsync(d.clip_limit.get(), clip_limit.data(), BackgroundBand::CELLS * sizeof(float),
                          cudaMemcpyHostToDevice, d.stream), "upload");
    constexpr int cell_blocks = (BackgroundBand::CELLS + THREADS - 1) / THREADS;
    sum_kernel<<<cell_blocks, THREADS, 0, d.stream>>>(d.width, d.offset, d.sorted_index, d.row_block, d.mean,
                                                      d.jac_x, d.jac_y, d.clip_limit, d.sum, d.sum_sq,
                                                      d.sum_jx, d.sum_jy, d.count);
    check(cudaGetLastError(), "clip");
    d.Download(sum, sum_sq, nullptr, nullptr, count);
}