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

#include "BeamCenterFFTGPU.h"

#include <algorithm>
#include <cmath>
#include <limits>

#include <cuda_runtime.h>
#include <cufft.h>

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

namespace {

constexpr int BLOCK = 256;
// Blocks in the greedy search's reduction. Each returns one candidate, and the last step of the
// reduction is done on the host over that many - a thousand is far below what a download costs and
// far above what the card needs to stay busy.
constexpr int SHORTLIST_BLOCKS = 1024;

// Every buffer here is read by work on the legacy NULL stream - the kernels, the cuFFT plans and the
// copies below are all queued there - so none of them may come from the pool. A pooled buffer is
// freed with cudaFreeAsync on the thread's allocation stream, which is non-blocking and so not
// ordered after the NULL stream: the free completed at once while a Suppress or an inverse transform
// was still writing, and the pool handed those bytes to the analysis engines being built on other
// workers at the same moment, or unmapped them. cudaFree synchronises the device first.
constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous;

int Blocks(int64_t n) { return static_cast<int>((n + BLOCK - 1) / BLOCK); }

void Check(cudaError_t err, const char *what) {
    if (err != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, what);
}

void Check(cufftResult res, const char *what) {
    if (res != CUFFT_SUCCESS)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, what);
}

std::vector<int> Dims(int64_t a) { return {static_cast<int>(a)}; }
std::vector<int> Dims(int64_t a, int64_t b) { return {static_cast<int>(a), static_cast<int>(b)}; }

// The transforms, as both the computation and its memory estimate take them: the point surfaces' 2D
// pair on an ny x nx grid.
CudaFFTShape PointForward(int64_t ny, int64_t nx) {
    const int64_t nxc = nx / 2 + 1;
    return {.n = Dims(ny, nx), .inembed = Dims(ny, nx), .idist = static_cast<int>(ny * nx),
            .onembed = Dims(ny, nxc), .odist = static_cast<int>(ny * nxc), .type = CUFFT_R2C};
}

CudaFFTShape PointInverse(int64_t ny, int64_t nx) {
    const int64_t nxc = nx / 2 + 1;
    return {.n = Dims(ny, nx), .inembed = Dims(ny, nxc), .idist = static_cast<int>(ny * nxc),
            .onembed = Dims(ny, nx), .odist = static_cast<int>(ny * nx), .type = CUFFT_C2R};
}

__global__ void SquareInPlace(float *data, int64_t n) {
    const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
    if (i < n)
        data[i] *= data[i];
}

__global__ void MultiplySpectra(const cufftComplex *u, const cufftComplex *v, cufftComplex *out,
                                int64_t n) {
    const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
    if (i < n)
        out[i] = cuCmulf(u[i], v[i]);
}

// The inverse transform leaves an unnormalised ny x nx result; only its 2h x 2w corner is the
// convolution surface. Scale that corner in place, for the one surface that stays where it is...
__global__ void ScaleCorner(float *data, int64_t nx, int64_t out_h, int64_t out_w, float norm) {
    const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
    if (i >= out_h * out_w)
        return;
    const int64_t y = i / out_w;
    data[y * nx + (i - y * out_w)] *= norm;
}

// ...and scale it into a surface of its own for the three that have to outlive the next inverse.
__global__ void CopyScaledCorner(const float *src, int64_t nx, int64_t out_h, int64_t out_w,
                                 float norm, float *dst) {
    const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
    if (i >= out_h * out_w)
        return;
    const int64_t y = i / out_w;
    dst[i] = src[y * nx + (i - y * out_w)] * norm;
}

// The masked Pearson, elementwise in double, exactly as BeamCenterPointScore does it on the host -
// the numerator and the denominator are cancellations of large near-equal terms. Q is read out of
// the transform buffer where the last inverse left it, so it needs no surface of its own. The score
// replaces C, which nothing wants afterwards.
__global__ void PearsonKernel(float *C, const float *S, const float *D, const float *Q_padded,
                              int64_t nx, int64_t out_h, int64_t out_w, double lim,
                              double variance_limit) {
    const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
    if (i >= out_h * out_w)
        return;
    const double d = D[i];
    if (d <= lim) {
        C[i] = -INFINITY;
        return;
    }
    const int64_t y = i / out_w;
    const double q = Q_padded[y * nx + (i - y * out_w)];
    const double s = S[i];
    const double num = d * C[i] - s * s;
    const double den = d * q - s * s;
    C[i] = den > d * d * variance_limit ? static_cast<float>(num / den) : -INFINITY;
}

// The host's rule for the greedy search, so the device finds the same peak: a strictly greater
// value wins, among equal values the lower index does, and an index of -1 means "nothing found"
// (which is what a surface of -infinity leaves, and what ends the shortlist).
__device__ bool TakeRight(float lv, int64_t li, float rv, int64_t ri) {
    if (ri < 0)
        return false;
    if (li < 0)
        return true;
    return rv > lv || (rv == lv && ri < li);
}

__global__ void BlockBest(const float *r, int64_t n, float *block_value, int64_t *block_index) {
    __shared__ float sv[BLOCK];
    __shared__ int64_t si[BLOCK];
    float bv = -INFINITY;
    int64_t bi = -1;
    for (int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x; i < n;
         i += static_cast<int64_t>(gridDim.x) * blockDim.x)
        if (r[i] > bv) {   // increasing i, so the lowest index among equal values is the one kept
            bv = r[i];
            bi = i;
        }
    sv[threadIdx.x] = bv;
    si[threadIdx.x] = bi;
    for (unsigned stride = blockDim.x / 2; stride > 0; stride /= 2) {
        __syncthreads();
        if (threadIdx.x < stride
            && TakeRight(sv[threadIdx.x], si[threadIdx.x], sv[threadIdx.x + stride],
                         si[threadIdx.x + stride])) {
            sv[threadIdx.x] = sv[threadIdx.x + stride];
            si[threadIdx.x] = si[threadIdx.x + stride];
        }
    }
    if (threadIdx.x == 0) {
        block_value[blockIdx.x] = sv[0];
        block_index[blockIdx.x] = si[0];
    }
}

// The non-maximum suppression square around a peak that has been taken.
__global__ void Suppress(float *r, int64_t w, int64_t x0, int64_t x1, int64_t y0, int64_t y1) {
    const int64_t cols = x1 - x0 + 1;
    const int64_t i = blockIdx.x * static_cast<int64_t>(blockDim.x) + threadIdx.x;
    if (i >= (y1 - y0 + 1) * cols)
        return;
    const int64_t y = i / cols;
    r[(y0 + y) * w + x0 + (i - y * cols)] = -INFINITY;
}

void CheckLastKernel(const char *what) {
    Check(cudaGetLastError(), what);
}

// The 2D transform set, left on the device. Only THREE 2h x 2w surfaces are ever allocated: the
// fourth is the last inverse's output, and nothing overwrites the transform buffer afterwards, so
// Q is read where it lies.
struct PointSurfacesDevice {
    CudaDevicePtr<float> C, S, D;   // 2h x 2w, normalised
    CudaDevicePtr<float> real;      // ny x nx; its 2h x 2w corner is Q, normalised
    int64_t nx = 0;
};

PointSurfacesDevice TransformPoint(const std::vector<float> &a, const std::vector<float> &m,
                                   int64_t h, int64_t w) {
    const int64_t ny = BeamCenterFFTPadSize(2 * h);
    const int64_t nx = BeamCenterFFTPadSize(2 * w);
    const int64_t nxc = nx / 2 + 1;
    const int64_t nreal = ny * nx;
    const int64_t ncomplex = ny * nxc;
    const int64_t ncrop = 4 * h * w;

    PointSurfacesDevice out;
    out.nx = nx;
    out.real = CudaDevicePtr<float>(nreal, ALLOC);
    // Two spectra, a product taking the place of one that is finished with (see below); and one work
    // area for both plans, which never run at the same time.
    CudaDevicePtr<cufftComplex> d_a(ncomplex, ALLOC), d_m(ncomplex, ALLOC);
    CudaDevicePtr<char> work(std::max<size_t>({CudaFFTPlan::WorkBytes(PointForward(ny, nx)),
                                               CudaFFTPlan::WorkBytes(PointInverse(ny, nx)), 1}), ALLOC);
    CudaFFTPlan fwd(PointForward(ny, nx), nullptr, work.get());
    CudaFFTPlan inv(PointInverse(ny, nx), nullptr, work.get());

    auto upload = [&](const std::vector<float> &src) {
        Check(cudaMemset(out.real.get(), 0, nreal * sizeof(float)),
              "BeamCenterFFT: memset failed");
        Check(cudaMemcpy2D(out.real.get(), nx * sizeof(float), src.data(), w * sizeof(float),
                           w * sizeof(float), h, cudaMemcpyHostToDevice),
              "BeamCenterFFT: image upload failed");
    };

    const float norm = 1.0f / (static_cast<float>(ny) * static_cast<float>(nx));
    // The product goes to `dst`, which may be u itself, and the inverse - which leaves an unnormalised
    // ny x nx result whose 2h x 2w corner is the surface - may overwrite it, as cuFFT's C2R does with
    // its input.
    auto inverse_product = [&](const cufftComplex *u, const cufftComplex *v, cufftComplex *dst) {
        MultiplySpectra<<<Blocks(ncomplex), BLOCK>>>(u, v, dst, ncomplex);
        CheckLastKernel("BeamCenterFFT: spectrum product failed");
        Check(inv.C2R(dst, out.real), "BeamCenterFFT: inverse transform failed");
    };
    auto take_corner = [&](CudaDevicePtr<float> &dst) {
        dst = CudaDevicePtr<float>(ncrop, ALLOC);
        CopyScaledCorner<<<Blocks(ncrop), BLOCK>>>(out.real.get(), nx, 2 * h, 2 * w, norm,
                                                   dst.get());
        CheckLastKernel("BeamCenterFFT: surface scaling failed");
    };

    // a.a goes where m's spectrum will be, before it is there; a.m in the place of a's spectrum, the last
    // to read it; m.m there too; and the spectrum of a^2 there after that.
    upload(a);
    Check(fwd.R2C(out.real, d_a), "BeamCenterFFT: forward transform failed");
    inverse_product(d_a, d_a, d_m);
    take_corner(out.C);
    upload(m);
    Check(fwd.R2C(out.real, d_m), "BeamCenterFFT: forward transform failed");
    inverse_product(d_a, d_m, d_a);
    take_corner(out.S);
    inverse_product(d_m, d_m, d_a);
    take_corner(out.D);
    upload(a);
    SquareInPlace<<<Blocks(nreal), BLOCK>>>(out.real.get(), nreal);
    CheckLastKernel("BeamCenterFFT: squaring failed");
    Check(fwd.R2C(out.real, d_a), "BeamCenterFFT: forward transform failed");
    inverse_product(d_a, d_m, d_a);
    ScaleCorner<<<Blocks(ncrop), BLOCK>>>(out.real.get(), nx, 2 * h, 2 * w, norm);
    CheckLastKernel("BeamCenterFFT: surface scaling failed");
    return out;
}

// The masked Pearson over the four surfaces, leaving the score in `surf.C`.
void ScorePoint(PointSurfacesDevice &surf, int64_t h, int64_t w, float min_pair_fraction,
                double global_variance) {
    const int64_t ncrop = 4 * h * w;
    CudaDevicePtr<float> block_value(SHORTLIST_BLOCKS, ALLOC);
    CudaDevicePtr<int64_t> block_index(SHORTLIST_BLOCKS, ALLOC);
    BlockBest<<<SHORTLIST_BLOCKS, BLOCK>>>(surf.D.get(), ncrop, block_value.get(),
                                           block_index.get());
    CheckLastKernel("BeamCenterFFT: pair-count maximum failed");
    std::vector<float> values(SHORTLIST_BLOCKS);
    Check(cudaMemcpy(values.data(), block_value.get(), values.size() * sizeof(float),
                     cudaMemcpyDeviceToHost),
          "BeamCenterFFT: pair-count maximum download failed");
    float d_max = 0.0f;   // as on the host, where the running maximum starts at zero
    for (float v : values)
        d_max = std::max(d_max, v);

    PearsonKernel<<<Blocks(ncrop), BLOCK>>>(surf.C.get(), surf.S.get(), surf.D.get(),
                                            surf.real.get(), surf.nx, 2 * h, 2 * w,
                                            min_pair_fraction * static_cast<double>(d_max),
                                            BEAM_CENTER_VARIANCE_FLOOR * global_variance);
    CheckLastKernel("BeamCenterFFT: masked Pearson failed");
}

} // namespace

BeamCenterConvSurfaces2D BeamCenterFFTGPU::PointSurfaces(const std::vector<float> &a,
                                                         const std::vector<float> &m, int64_t h,
                                                         int64_t w) {
    // The four surfaces brought back as the interface describes them. Nothing in the pipeline asks
    // for this any more - PointShortlist below keeps them on the device - but the shared host code
    // is written on them and the parity test compares them, so the engine still supplies them.
    PointSurfacesDevice surf = TransformPoint(a, m, h, w);
    const size_t ncrop = static_cast<size_t>(4 * h * w);

    BeamCenterConvSurfaces2D out;
    auto download = [&](const CudaDevicePtr<float> &src, std::vector<float> &dst) {
        dst.resize(ncrop);
        Check(cudaMemcpy(dst.data(), src.get(), ncrop * sizeof(float), cudaMemcpyDeviceToHost),
              "BeamCenterFFT: surface download failed");
    };
    download(surf.C, out.C);
    download(surf.S, out.S);
    download(surf.D, out.D);
    out.Q.resize(ncrop);
    Check(cudaMemcpy2D(out.Q.data(), 2 * w * sizeof(float), surf.real.get(),
                       surf.nx * sizeof(float), 2 * w * sizeof(float), 2 * h,
                       cudaMemcpyDeviceToHost),
          "BeamCenterFFT: surface download failed");
    return out;
}

std::vector<float> BeamCenterFFTGPU::PointScoreSurface(const std::vector<float> &a,
                                                       const std::vector<float> &m, int64_t h,
                                                       int64_t w, float min_pair_fraction,
                                                       double global_variance) {
    PointSurfacesDevice surf = TransformPoint(a, m, h, w);
    ScorePoint(surf, h, w, min_pair_fraction, global_variance);
    std::vector<float> out(static_cast<size_t>(4 * h * w));
    Check(cudaMemcpy(out.data(), surf.C.get(), out.size() * sizeof(float), cudaMemcpyDeviceToHost),
          "BeamCenterFFT: score download failed");
    return out;
}

std::vector<BeamCenterFFTCandidate>
BeamCenterFFTGPU::PointShortlist(const std::vector<float> &a, const std::vector<float> &m,
                                 int64_t h, int64_t w, const BeamCenterFFTSettings &settings,
                                 double global_variance) {
    PointSurfacesDevice surf = TransformPoint(a, m, h, w);
    ScorePoint(surf, h, w, settings.min_pair_fraction, global_variance);

    // The greedy search, peak by peak, where the surface is. Each round reduces it to one candidate
    // per block and the host picks between those by the same rule, so the answer is the host's and
    // does not depend on how the reduction was split.
    const int64_t surface_h = 2 * h, surface_w = 2 * w;
    const int64_t n = surface_h * surface_w;
    const int64_t d = std::lround(2.0f * settings.nms_radius_pxl);
    CudaDevicePtr<float> block_value(SHORTLIST_BLOCKS, ALLOC);
    CudaDevicePtr<int64_t> block_index(SHORTLIST_BLOCKS, ALLOC);
    std::vector<float> values(SHORTLIST_BLOCKS);
    std::vector<int64_t> indices(SHORTLIST_BLOCKS);

    std::vector<BeamCenterFFTCandidate> out;
    for (int k = 0; k < settings.candidates_point; k++) {
        BlockBest<<<SHORTLIST_BLOCKS, BLOCK>>>(surf.C.get(), n, block_value.get(),
                                               block_index.get());
        CheckLastKernel("BeamCenterFFT: peak search failed");
        Check(cudaMemcpy(values.data(), block_value.get(), values.size() * sizeof(float),
                         cudaMemcpyDeviceToHost),
              "BeamCenterFFT: peak download failed");
        Check(cudaMemcpy(indices.data(), block_index.get(), indices.size() * sizeof(int64_t),
                         cudaMemcpyDeviceToHost),
              "BeamCenterFFT: peak download failed");

        int64_t best = -1;
        float best_v = -std::numeric_limits<float>::infinity();
        for (int b = 0; b < SHORTLIST_BLOCKS; b++)
            if (indices[b] >= 0
                && (values[b] > best_v || (values[b] == best_v && indices[b] < best)))
                best_v = values[b], best = indices[b];
        if (best < 0 || !std::isfinite(best_v))
            break;

        const int64_t iy = best / surface_w, ix = best % surface_w;
        out.push_back({static_cast<float>(ix) / 2.0f, static_cast<float>(iy) / 2.0f, best_v});
        const int64_t y0 = std::max<int64_t>(0, iy - d), y1 = std::min(surface_h - 1, iy + d);
        const int64_t x0 = std::max<int64_t>(0, ix - d), x1 = std::min(surface_w - 1, ix + d);
        Suppress<<<Blocks((y1 - y0 + 1) * (x1 - x0 + 1)), BLOCK>>>(surf.C.get(), surface_w, x0, x1,
                                                                   y0, y1);
        CheckLastKernel("BeamCenterFFT: peak suppression failed");
    }
    return out;
}

size_t BeamCenterFFTGPU::DeviceMemoryNeeded(int64_t width, int64_t height) {
    const int64_t ny = BeamCenterFFTPadSize(2 * height);
    const int64_t nx = BeamCenterFFTPadSize(2 * width);
    const int64_t nxc = nx / 2 + 1;
    const size_t work_fwd = CudaFFTPlan::WorkBytes(PointForward(ny, nx));
    const size_t work_inv = CudaFFTPlan::WorkBytes(PointInverse(ny, nx));
    // One transform buffer, two spectra, three of the four convolution surfaces - the fourth stays in
    // the transform buffer, where the last inverse leaves it - and one work area for both plans.
    return ny * nx * sizeof(float) + 2 * ny * nxc * sizeof(cufftComplex)
           + 3 * 4 * static_cast<size_t>(width) * height * sizeof(float)
           + std::max(work_fwd, work_inv);
}

bool BeamCenterFFTGPU::FitsInDeviceMemory(int64_t width, int64_t height) {
    // The work areas are read from the plans themselves, and a transform cuFFT cannot plan does not fit.
    size_t needed = 0;
    try {
        needed = DeviceMemoryNeeded(width, height);
    } catch (const JFJochException &) {
        cuda_clear_error();
        return false;
    }
    size_t free_bytes = 0, total_bytes = 0;
    Check(cudaMemGetInfo(&free_bytes, &total_bytes), "Beam centre FFT: reading the device memory");
    return needed <= total_bytes / 2;
}