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

#include "FFTIndexerGPU.h"
#include <cufft.h>
#include <cmath>
#include <algorithm>

__device__ __host__ inline float complex_abs(const cufftComplex &z) {
    return sqrtf(z.x * z.x + z.y * z.y);
}

// The magnitudes of the transform, transposed to bin-major order: mag_t[bin * stride + direction], for the
// `directions_size` directions of one batch (mag_t points at the batch's first column, stride is all of them).
// The R2C output is direction-major, and the peak search gives each direction to one thread, so read
// straight from it neighbouring lanes are a whole spectrum apart - every load a separate sector, three
// streams per thread. Transposed, a warp reads 32 consecutive floats per bin. On an idle card that buys
// nothing (the search is latency-bound either way, and the transpose costs its own pass); with the
// stills indexer pool's streams sharing the card it is the difference between 67 and 98 frames/s
// (measured). A 32x32 tile in shared memory keeps both the read and the write coalesced. The search took
// complex_abs of the same values, so the magnitudes are bit for bit the ones it computed itself.
__global__ void magnitude_transpose(const cufftComplex *__restrict__ d_output,
                                    const int out_len,
                                    const int directions_size,
                                    const int stride,
                                    float *__restrict__ mag_t) {
    __shared__ float tile[32][33];
    const int bin0 = blockIdx.x * 32;
    const int dir0 = blockIdx.y * 32;
    for (int r = threadIdx.y; r < 32; r += blockDim.y) {
        const int dir = dir0 + r;
        const int bin = bin0 + threadIdx.x;
        if (dir < directions_size && bin < out_len)
            tile[r][threadIdx.x] = complex_abs(d_output[static_cast<size_t>(dir) * out_len + bin]);
    }
    __syncthreads();
    for (int r = threadIdx.y; r < 32; r += blockDim.y) {
        const int bin = bin0 + r;
        const int dir = dir0 + threadIdx.x;
        if (dir < directions_size && bin < out_len)
            mag_t[static_cast<size_t>(bin) * stride + dir] = tile[threadIdx.x][r];
    }
}

// The peak search splits each direction's spectrum into PEAK_CHUNKS consecutive runs of bins and gives each
// run to its own thread, which keeps a running window sum over its run; one thread per direction leaves the
// card nearly empty (16384 threads), and this search, a chain of dependent FP64 additions per thread, is
// then latency-bound. A block holds 32 directions (x, so the bin-major loads stay coalesced) times PEAK_CHUNKS.
// A run after the first sums its starting window afresh instead of carrying the sum in from the bins before
// it. That is the same double wherever the sums are exact, and they are: a window of float magnitudes fits
// a double's 53 bits unless one of them is below ~1e-9 of the window sum (no difference over 3000 stills
// frames and a rotation-style synthetic set). More runs buy nothing: the card's FP64 rate is then the limit.
constexpr int PEAK_CHUNKS = 4;

__global__ void calculate_fft_result(
    const float *__restrict__ mag_t,
    const float max_length_A,
    const float min_length_A,
    const int histogram_size,
    const int bg_half,
    const int directions_size,
    FFTResult *d_results) {
    __shared__ float run_prom[PEAK_CHUNKS][32];
    __shared__ float run_len[PEAK_CHUNKS][32];

    const int i = blockIdx.x * 32 + threadIdx.x;
    const int run = threadIdx.y;
    float best_prom = 0.0f;
    float best_len = -1.0f;

    if (i < directions_size) {
        const int out_len = (histogram_size / 2) + 1;
        // |transform| of bin k of this direction, from the bin-major array (magnitude_transpose)
        const auto mag_at = [&](int k) { return mag_t[static_cast<size_t>(k) * directions_size + i]; };
        const float len_coeff = 2.0f * max_length_A / static_cast<float>(histogram_size);

        // Pick the peak by PROMINENCE above a running-mean background of half-width bg_half:
        // the projected histogram has a broad low-frequency envelope whose magnitude can
        // exceed the true lattice peaks, so a plain argmax|spec| returns a short envelope
        // vector on weak/pink-beam frames. Subtracting the local mean removes that envelope
        // while sharp lattice peaks keep their height (mirrors FFTIndexerCPU).
        //
        // The window has a constant width, slid inward at the ends instead of truncated (see
        // FFTIndexerCPU): a peak within bg_half of either end - where the LONGEST cells sit - otherwise
        // gets a one-sided background and a biased prominence. A spectrum shorter than the window
        // averages over all of itself.
        const int width = min(2 * bg_half, out_len - 1);
        const int lo_max = out_len - 1 - width;
        const auto window_lo = [&](int j) { return min(max(j - bg_half, 0), lo_max); };

        // The window width is fixed, so the mean divides by one number throughout. A double division
        // is a long iteration on the FP64 units a consumer card has few of; multiplying by the correctly
        // rounded reciprocal and correcting the quotient once with its exact remainder gives the
        // correctly rounded quotient - the same double the division returns (Markstein (1990) IBM J.
        // Res. Dev. 34, 111-119).
        const double count = width + 1;
        const double inv_count = 1.0 / count;

        const int j0 = (out_len * run) / PEAK_CHUNKS;
        const int j1 = (out_len * (run + 1)) / PEAK_CHUNKS;
        int lo = window_lo(j0);
        double winsum = 0.0;
        for (int k = lo; k <= lo + width; ++k)
            winsum += mag_at(k);

        for (int j = j0; j < j1; ++j) {
            if (window_lo(j) > lo) {
                winsum += mag_at(lo + width + 1);
                winsum -= mag_at(lo);
                ++lo;
            }
            const float len = len_coeff * static_cast<float>(j);
            if (len <= min_length_A)
                continue;
            const double q = __dmul_rn(winsum, inv_count);
            const float bg = static_cast<float>(__fma_rn(__fma_rn(-q, count, winsum), inv_count, q));
            const float prom = mag_at(j) - bg;
            if (prom > best_prom) {
                best_prom = prom;
                best_len = len;
            }
        }
    }

    run_prom[run][threadIdx.x] = best_prom;
    run_len[run][threadIdx.x] = best_len;
    __syncthreads();

    // Each run kept its first bin of highest prominence; taking the runs in order with the same strict
    // comparison gives the direction's first bin of highest prominence.
    if (run == 0 && i < directions_size) {
        FFTResult result{.magnitude = 0.0f, .direction = i, .length = -1};
        for (int r = 0; r < PEAK_CHUNKS; ++r) {
            if (run_prom[r][threadIdx.x] > result.magnitude) {
                result.magnitude = run_prom[r][threadIdx.x];
                result.length = run_len[r][threadIdx.x];
            }
        }
        d_results[i] = result;
    }
}

__global__ void histogram_kernel(const float *__restrict__ coord_x,
                                 const float *__restrict__ coord_y,
                                 const float *__restrict__ coord_z,
                                 const float *__restrict__ dir_x,
                                 const float *__restrict__ dir_y,
                                 const float *__restrict__ dir_z,
                                 float histogram_spacing,
                                 int histogram_size,
                                 int coord_size,
                                 int direction_vectors_size,
                                 float *__restrict__ output) {
    int direction_idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (direction_idx < direction_vectors_size) {
        int base_offset = direction_idx * histogram_size;

        for (int i = 0; i < histogram_size; i++)
            output[base_offset + i] = 0;

        for (int i = 0; i < coord_size; i++) {
            float dot = fabsf(
                dir_x[direction_idx] * coord_x[i] + dir_y[direction_idx] * coord_y[i] + dir_z[direction_idx] * coord_z[
                    i]);
            int64_t bin = static_cast<int64_t>(dot / histogram_spacing);
            if (bin >= 0 && bin < histogram_size)
                output[base_offset + bin] += 1.0;
        }
    }
}

// The same histogram, one block per direction with the bins in shared memory.
//
// The kernel above gives a whole direction to a single thread, so neighbouring lanes write
// histogram_size floats apart - 12.6 kB at the default sizing. Every warp instruction then touches 32
// separate sectors of a buffer that is hundreds of megabytes (16384 directions x 3142 bins), with no
// hope of staying in a 4 MB L2, and the projection loop becomes 160 million scattered global
// read-modify-writes.
//
// The bins are counts, so here they are integers in shared memory. That matters twice: an integer
// atomicAdd is a real shared-memory instruction on Turing and Ada, where the float one compiles to a
// compare-and-swap retry loop; and a count below 2^24 converts to float exactly, so the output is bit
// for bit what the repeated `+= 1.0` above produces. The division by histogram_spacing stays a
// division - turning it into a multiply by the reciprocal would move a spot across a bin edge.
__global__ void histogram_shared_kernel(const float *__restrict__ coord_x,
                                        const float *__restrict__ coord_y,
                                        const float *__restrict__ coord_z,
                                        const float *__restrict__ dir_x,
                                        const float *__restrict__ dir_y,
                                        const float *__restrict__ dir_z,
                                        float histogram_spacing,
                                        int histogram_size,
                                        int coord_size,
                                        int direction_vectors_size,
                                        float *__restrict__ output) {
    extern __shared__ unsigned int bins[];
    const int direction_idx = blockIdx.x;
    if (direction_idx >= direction_vectors_size)
        return;

    for (int i = threadIdx.x; i < histogram_size; i += blockDim.x)
        bins[i] = 0;
    __syncthreads();

    const float dx = dir_x[direction_idx], dy = dir_y[direction_idx], dz = dir_z[direction_idx];
    for (int i = threadIdx.x; i < coord_size; i += blockDim.x) {
        const float dot = fabsf(dx * coord_x[i] + dy * coord_y[i] + dz * coord_z[i]);
        const int64_t bin = static_cast<int64_t>(dot / histogram_spacing);
        if (bin >= 0 && bin < histogram_size)
            atomicAdd(&bins[bin], 1u);
    }
    __syncthreads();

    float *out = output + static_cast<size_t>(direction_idx) * static_cast<size_t>(histogram_size);
    for (int i = threadIdx.x; i < histogram_size; i += blockDim.x)
        out[i] = static_cast<float>(bins[i]);
}

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

inline void cuda_err(cufftResult val) {
    if (val != cufftResult::CUFFT_SUCCESS)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "CuFFT error");
}

FFTIndexerGPU::FFTIndexerGPU(const IndexingSettings &settings)
    : FFTIndexer(settings), result_fft_reg(result_fft) {
    const int batch = std::min(nDirections, FFT_BATCH_DIRECTIONS);
    const int out_len = histogram_size / 2 + 1;
    d_input_fft = CudaDevicePtr<float>(static_cast<size_t>(batch) * histogram_size);
    d_output_fft = CudaDevicePtr<cufftComplex>(static_cast<size_t>(batch) * out_len);
    d_mag_t = CudaDevicePtr<float>(static_cast<size_t>(nDirections) * out_len);
    d_result_fft = CudaDevicePtr<FFTResult>(nDirections);

    d_spot_x = CudaDevicePtr<float>(FFT_MAX_SPOTS);
    d_spot_y = CudaDevicePtr<float>(FFT_MAX_SPOTS);
    d_spot_z = CudaDevicePtr<float>(FFT_MAX_SPOTS);

    spot_x = CudaHostPtr<float>(FFT_MAX_SPOTS);
    spot_y = CudaHostPtr<float>(FFT_MAX_SPOTS);
    spot_z = CudaHostPtr<float>(FFT_MAX_SPOTS);

    d_dir_x = CudaDevicePtr<float>(nDirections);
    d_dir_y = CudaDevicePtr<float>(nDirections);
    d_dir_z = CudaDevicePtr<float>(nDirections);

    DirectionsChanged();

    const auto shape = [&](int n) {
        return CudaFFTShape{.n = {static_cast<int>(histogram_size)},
                            .idist = static_cast<int>(histogram_size),
                            .odist = out_len,
                            .type = CUFFT_R2C,
                            .batch = n};
    };
    plan = CudaFFTPlan(shape(batch), stream);
    if (nDirections % batch != 0)
        plan_last = CudaFFTPlan(shape(nDirections % batch), stream);
}


void FFTIndexerGPU::DirectionsChanged() {
    std::vector<float> dir_x(nDirections), dir_y(nDirections), dir_z(nDirections);
    for (int i = 0; i < nDirections; i++) {
        dir_x[i] = direction_vectors.at(i).x;
        dir_y[i] = direction_vectors.at(i).y;
        dir_z[i] = direction_vectors.at(i).z;
    }
    // Checked like every other CUDA call here: a silently failed upload leaves the device holding the
    // PREVIOUS grid, so the directions the host reads the results against are not the ones the kernel
    // used - wrong rows reported as good ones, and nothing says so. On the engine's stream, which the
    // kernels reading them run on: a cudaMemcpy on the NULL stream is not ordered before those.
    cuda_err(cudaMemcpyAsync(d_dir_x, dir_x.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice, stream));
    cuda_err(cudaMemcpyAsync(d_dir_y, dir_y.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice, stream));
    cuda_err(cudaMemcpyAsync(d_dir_z, dir_z.data(), nDirections * sizeof(float), cudaMemcpyHostToDevice, stream));
}

void FFTIndexerGPU::ExecuteFFT(const std::vector<Coord> &coord, size_t nspots) {
    const int l_blockDim = 128;

    for (int i = 0; i < nspots; i++) {
        spot_x[i] = coord[i].x;
        spot_y[i] = coord[i].y;
        spot_z[i] = coord[i].z;
    }

    cuda_err(cudaMemcpyAsync(d_spot_x, spot_x, nspots * sizeof(float), cudaMemcpyHostToDevice, stream));
    cuda_err(cudaMemcpyAsync(d_spot_y, spot_y, nspots * sizeof(float), cudaMemcpyHostToDevice, stream));
    cuda_err(cudaMemcpyAsync(d_spot_z, spot_z, nspots * sizeof(float), cudaMemcpyHostToDevice, stream));

    // The histograms and their transforms, FFT_BATCH_DIRECTIONS directions at a time (see the header); the
    // magnitudes of all of them go to d_mag_t, which the peak search reads in one go.
    const int ndir = static_cast<int>(direction_vectors.size());
    const int out_len = static_cast<int>(histogram_size / 2 + 1);
    // Shared-memory bins where they fit (they do at any sane sizing - 12.6 kB at the defaults), the
    // thread-per-direction kernel where they do not.
    const size_t hist_shared_bytes = static_cast<size_t>(histogram_size) * sizeof(unsigned int);
    for (int first = 0; first < ndir; first += FFT_BATCH_DIRECTIONS) {
        const int n = std::min(FFT_BATCH_DIRECTIONS, ndir - first);
        if (hist_shared_bytes <= 48 * 1024) {
            histogram_shared_kernel<<<n, 256, hist_shared_bytes, stream>>>(
                    d_spot_x, d_spot_y, d_spot_z, d_dir_x + first, d_dir_y + first, d_dir_z + first,
                    histogram_spacing, histogram_size, nspots, n, d_input_fft);
        } else {
            histogram_kernel<<<(n + l_blockDim - 1) / l_blockDim, l_blockDim, 0, stream>>>(
                    d_spot_x, d_spot_y, d_spot_z, d_dir_x + first, d_dir_y + first, d_dir_z + first,
                    histogram_spacing, histogram_size, nspots, n, d_input_fft);
        }
        cuda_err(cudaGetLastError());

        cuda_err((n == FFT_BATCH_DIRECTIONS || n == nDirections ? plan : plan_last).R2C(
                              d_input_fft, d_output_fft));

        const dim3 tr_grid((out_len + 31) / 32, (n + 31) / 32);
        magnitude_transpose<<<tr_grid, dim3(32, 8), 0, stream>>>(d_output_fft, out_len, n, ndir,
                                                                 d_mag_t + first);
        cuda_err(cudaGetLastError());
    }

    // Background half-window ~15 A (length-based, so independent of histogram sizing); see
    // FFTIndexerCPU for the prominence-vs-envelope rationale and the validated optimum.
    const double len_coeff = 2.0 * static_cast<double>(max_length_A) / static_cast<double>(histogram_size);
    const int bg_half = std::max(1, static_cast<int>(std::lround(15.0 / len_coeff)));
    const int peak_grid = (static_cast<int>(direction_vectors.size()) + 31) / 32;
    calculate_fft_result<<<peak_grid, dim3(32, PEAK_CHUNKS), 0, stream>>>(d_mag_t,
                                                               max_length_A, min_length_A, histogram_size,
                                                               bg_half,
                                                               direction_vectors.size(), d_result_fft);
    cuda_err(cudaGetLastError());

    cuda_err(cudaMemcpyAsync(result_fft.data(), d_result_fft, direction_vectors.size() * sizeof(FFTResult),
                             cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
}