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

#include "PairSortGPU.h"

#include <algorithm>
#include <string>

#include <cub/device/device_radix_sort.cuh>

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

namespace {
    constexpr int THREADS = 256;
    constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside;

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

    __global__ void SplitKernel(int n, const double2 *pair, double *second, int32_t *index) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        second[i] = pair[i].y;
        index[i] = i;
    }

    __global__ void GatherFirstKernel(int n, const double2 *pair, const int32_t *index, double *first) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        first[i] = pair[index[i]].x;
    }

    __global__ void GatherPairKernel(int n, const double2 *pair, const int32_t *index, double2 *out) {
        const int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i >= n) return;
        out[i] = pair[index[i]];
    }
}

void SortPairsGPU(std::vector<std::pair<double, double>> &v) {
    static_assert(sizeof(std::pair<double, double>) == sizeof(double2));
    const int n = static_cast<int>(v.size());
    if (n < 2)
        return;
    constexpr size_t BYTES_PER_PAIR = 16 + 16 + 16 + 8;   // pairs in and out, keys, indices
    size_t free_bytes = 0;
    if (!cuda_memory_available(BYTES_PER_PAIR * v.size(), free_bytes))
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                              "Sorting " + std::to_string(v.size()) + " pairs needs "
                              + std::to_string(BYTES_PER_PAIR * v.size() >> 20) + " MiB of GPU memory, "
                              + std::to_string(free_bytes >> 20) + " MiB free - run on the CPU (rugnux --no-gpu)");
    const unsigned blocks = static_cast<unsigned>((n + THREADS - 1) / THREADS);
    CudaStream stream;
    CudaDevicePtr<double2> pair(n, ALLOC), out(n, ALLOC);
    CudaDevicePtr<double> key(n, ALLOC), key_sorted(n, ALLOC);
    CudaDevicePtr<int32_t> index(n, ALLOC), index_sorted(n, ALLOC);
    cuda_err(cudaMemcpyAsync(pair, v.data(), v.size() * sizeof(double2), cudaMemcpyHostToDevice, stream));
    // By .second, then - stably - by .first: the order of the pairs.
    SplitKernel<<<blocks, THREADS, 0, stream>>>(n, pair, key, index);
    cuda_err(cudaGetLastError());
    size_t bytes = 0;
    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key.get(), key_sorted.get(), index.get(),
                                             index_sorted.get(), n, 0, 64, stream));
    CudaDevicePtr<uint8_t> scratch(std::max<size_t>(bytes, 1), ALLOC);
    cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key.get(), key_sorted.get(), index.get(),
                                             index_sorted.get(), n, 0, 64, stream));
    GatherFirstKernel<<<blocks, THREADS, 0, stream>>>(n, pair, index_sorted, key);
    cuda_err(cudaGetLastError());
    cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key.get(), key_sorted.get(), index_sorted.get(),
                                             index.get(), n, 0, 64, stream));
    GatherPairKernel<<<blocks, THREADS, 0, stream>>>(n, pair, index, out);
    cuda_err(cudaGetLastError());
    cuda_err(cudaMemcpyAsync(v.data(), out, v.size() * sizeof(double2), cudaMemcpyDeviceToHost, stream));
    cuda_err(cudaStreamSynchronize(stream));
}