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

#include "ErrorModelGPU.h"

#include <algorithm>
#include <string>

#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_segmented_radix_sort.cuh>

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

namespace {
    constexpr int THREADS = 256;
    constexpr int CHUNKS = 64;   // blocks per bin in the sum kernel; fixed, so the tree is too
    constexpr int NB = ERROR_MODEL_BINS;
    constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous;

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

    unsigned Blocks(size_t n) { return static_cast<unsigned>(std::max<size_t>(1, (n + THREADS - 1) / THREADS)); }

    // field 0..3: dev2, I2, s2, snr2 = I2 / s2 - the rank key's components, least significant first.
    __global__ void KeyKernel(int n, const ErrorModelSample *pool, const int32_t *idx, int field, double *key) {
        const int j = blockIdx.x * blockDim.x + threadIdx.x;
        if (j >= n) return;
        const ErrorModelSample &x = pool[idx[j]];
        key[j] = field == 0 ? x.dev2 : field == 1 ? x.I2 : field == 2 ? x.s2 : x.I2 / x.s2;
    }

    __global__ void IotaKernel(int n, int32_t *idx) {
        const int j = blockIdx.x * blockDim.x + threadIdx.x;
        if (j < n) idx[j] = j;
    }

    __global__ void GatherKernel(int n, const ErrorModelSample *pool, const int32_t *idx, double *s2, double *I2,
                                 double *dev2) {
        const int j = blockIdx.x * blockDim.x + threadIdx.x;
        if (j >= n) return;
        const ErrorModelSample &x = pool[idx[j]];
        s2[j] = x.s2; I2[j] = x.I2; dev2[j] = x.dev2;
    }

    // Block (chunk, bin): z2 of its samples, and its share of the bin's sums of s2/var and I2/var,
    // reduced over a fixed tree.
    __global__ void BinSumKernel(const int64_t *bin_start, const double *s2, const double *I2, const double *dev2,
                                 double a, double b2, double *z2, double *part_s, double *part_I) {
        __shared__ double sh_s[THREADS], sh_I[THREADS];
        const int bin = blockIdx.y, chunk = blockIdx.x;
        const int64_t lo = bin_start[bin], n = bin_start[bin + 1] - lo;
        const int64_t c0 = lo + n * chunk / CHUNKS, c1 = lo + n * (chunk + 1) / CHUNKS;
        double s = 0.0, I = 0.0;
        for (int64_t i = c0 + threadIdx.x; i < c1; i += THREADS) {
            const double v = __dadd_rn(__dmul_rn(a, s2[i]), __dmul_rn(b2, I2[i]));
            s += s2[i] / v; I += I2[i] / v;
            z2[i] = dev2[i] / v;
        }
        sh_s[threadIdx.x] = s; sh_I[threadIdx.x] = I;
        __syncthreads();
        for (int stride = THREADS / 2; stride > 0; stride >>= 1) {
            if (threadIdx.x < stride) {
                sh_s[threadIdx.x] += sh_s[threadIdx.x + stride];
                sh_I[threadIdx.x] += sh_I[threadIdx.x + stride];
            }
            __syncthreads();
        }
        if (threadIdx.x == 0) {
            part_s[bin * CHUNKS + chunk] = sh_s[0];
            part_I[bin * CHUNKS + chunk] = sh_I[0];
        }
    }

    __global__ void MedianKernel(const int64_t *bin_start, const double *z2_sorted, double *median) {
        const int bin = threadIdx.x;
        if (bin >= NB) return;
        const int64_t lo = bin_start[bin], n = bin_start[bin + 1] - lo;
        median[bin] = z2_sorted[lo + n / 2];
    }
}

ErrorModelFit FitErrorModelGPU(const std::vector<ErrorModelSample> &pool) {
    ErrorModelFit fit;
    if (pool.size() < static_cast<size_t>(8 * NB))
        return fit;
    const int n = static_cast<int>(pool.size());

    // The rank sort is the peak: the pool, two key and two index arrays and the sort scratch.
    constexpr size_t BYTES_PER_SAMPLE = 32 + 2 * 8 + 2 * 4 + 16;
    size_t free_bytes = 0;
    if (!cuda_memory_available(BYTES_PER_SAMPLE * pool.size(), free_bytes))
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                              "Error model: " + std::to_string(pool.size()) + " samples need "
                              + std::to_string(BYTES_PER_SAMPLE * pool.size() >> 20) + " MiB of GPU memory, "
                              + std::to_string(free_bytes >> 20) + " MiB free - run the CPU build");

    CudaStream stream;
    CudaDevicePtr<ErrorModelSample> d_pool(pool.size(), ALLOC);
    cuda_err(cudaMemcpyAsync(d_pool, pool.data(), pool.size() * sizeof(ErrorModelSample), cudaMemcpyHostToDevice,
                             stream));

    // The samples in the host's rank order - (I2/s2, s2, I2, dev2), a total order - by four stable
    // passes, least significant key first. The equal-count bins are then consecutive stretches.
    CudaDevicePtr<int32_t> idx(n, ALLOC), idx_next(n, ALLOC);
    CudaDevicePtr<double> key(n, ALLOC), key_sorted(n, ALLOC);
    IotaKernel<<<Blocks(n), THREADS, 0, stream>>>(n, idx);
    cuda_err(cudaGetLastError());
    size_t bytes = 0;
    cuda_err(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key.get(), key_sorted.get(), idx.get(), idx_next.get(),
                                             n, 0, 64, stream));
    CudaDevicePtr<uint8_t> scratch(std::max<size_t>(bytes, 1), ALLOC);
    for (int field = 0; field < 4; ++field) {
        KeyKernel<<<Blocks(n), THREADS, 0, stream>>>(n, d_pool, idx, field, key);
        cuda_err(cudaGetLastError());
        cuda_err(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key.get(), key_sorted.get(), idx.get(),
                                                 idx_next.get(), n, 0, 64, stream));
        std::swap(idx, idx_next);
    }
    scratch = CudaDevicePtr<uint8_t>(); idx_next = CudaDevicePtr<int32_t>(); key = CudaDevicePtr<double>();
    // After the last pass key_sorted holds I2/s2 in rank order.
    CudaDevicePtr<double> s2(n, ALLOC), I2(n, ALLOC), dev2(n, ALLOC);
    GatherKernel<<<Blocks(n), THREADS, 0, stream>>>(n, d_pool, idx, s2, I2, dev2);
    cuda_err(cudaGetLastError());
    d_pool = CudaDevicePtr<ErrorModelSample>(); idx = CudaDevicePtr<int32_t>();

    const size_t per = pool.size() / NB;
    std::vector<int64_t> bin_start(NB + 1);
    for (int b = 0; b < NB; ++b) bin_start[b] = static_cast<int64_t>(b * per);
    bin_start[NB] = n;
    CudaDevicePtr<int64_t> d_bin_start(NB + 1, ALLOC);
    cuda_err(cudaMemcpyAsync(d_bin_start, bin_start.data(), (NB + 1) * sizeof(int64_t), cudaMemcpyHostToDevice,
                             stream));
    {
        double top_mid = 0.0;
        cuda_err(cudaMemcpyAsync(&top_mid, key_sorted.get() + (bin_start[NB - 1] + n) / 2, sizeof(double),
                                 cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaStreamSynchronize(stream));
        fit.b_measured = top_mid > ERROR_MODEL_B_LEVER_MIN;
    }

    // The iterations: per bin the two sums (each chunk's partial added on the host in chunk order) and
    // the median z2, read off each bin's sorted values.
    CudaDevicePtr<double> z2(n, ALLOC), z2_sorted(n, ALLOC), part_s(NB * CHUNKS, ALLOC), part_I(NB * CHUNKS, ALLOC),
            median(NB, ALLOC);
    size_t seg_bytes = 0;
    cuda_err(cub::DeviceSegmentedRadixSort::SortKeys(nullptr, seg_bytes, z2.get(), z2_sorted.get(), n, NB,
                                                     d_bin_start.get(), d_bin_start.get() + 1, 0, 64, stream));
    CudaDevicePtr<uint8_t> seg_scratch(std::max<size_t>(seg_bytes, 1), ALLOC);
    std::vector<double> h_part_s(NB * CHUNKS), h_part_I(NB * CHUNKS), h_median(NB);
    std::vector<double> bs2(NB), bI2(NB), bd2(NB);
    double a = 1.0, b2 = 0.0;
    for (int iter = 0; iter < 30; ++iter) {
        BinSumKernel<<<dim3(CHUNKS, NB), THREADS, 0, stream>>>(d_bin_start, s2, I2, dev2, a, b2, z2, part_s, part_I);
        cuda_err(cudaGetLastError());
        cuda_err(cub::DeviceSegmentedRadixSort::SortKeys(seg_scratch.get(), seg_bytes, z2.get(), z2_sorted.get(), n,
                                                         NB, d_bin_start.get(), d_bin_start.get() + 1, 0, 64,
                                                         stream));
        MedianKernel<<<1, NB, 0, stream>>>(d_bin_start, z2_sorted, median);
        cuda_err(cudaGetLastError());
        cuda_err(cudaMemcpyAsync(h_part_s.data(), part_s, NB * CHUNKS * sizeof(double), cudaMemcpyDeviceToHost,
                                 stream));
        cuda_err(cudaMemcpyAsync(h_part_I.data(), part_I, NB * CHUNKS * sizeof(double), cudaMemcpyDeviceToHost,
                                 stream));
        cuda_err(cudaMemcpyAsync(h_median.data(), median, NB * sizeof(double), cudaMemcpyDeviceToHost, stream));
        cuda_err(cudaStreamSynchronize(stream));
        for (int bin = 0; bin < NB; ++bin) {
            double s = 0.0, I = 0.0;
            for (int c = 0; c < CHUNKS; ++c) {
                s += h_part_s[bin * CHUNKS + c];
                I += h_part_I[bin * CHUNKS + c];
            }
            const double nb = static_cast<double>(bin_start[bin + 1] - bin_start[bin]);
            bs2[bin] = s; bI2[bin] = I; bd2[bin] = h_median[bin] / ERROR_MODEL_CHI2_1_MEDIAN * nb;
        }
        if (!ErrorModelUpdate(bs2, bI2, bd2, a, b2, fit))
            break;
    }
    return fit;
}