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

#include "ShadowAccumulatorGPU.h"

#include <limits>
#include <type_traits>

#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"

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

namespace {

// One frame folded into the projection. The sentinel test is predicated rather than branched: a
// warp straddling a module gap then costs the same as one that does not, against a per-frame memory
// bill of several hundred megabytes.
//
// The update is written exactly as the host writes it, including the "first value wins whatever it
// is" rule for the maximum - a pixel that has never been counted holds 0, which a genuinely negative
// maximum has to be allowed to replace. Each pixel is owned by one thread and the frames are
// separate launches on one stream, so there is no race and no atomic.
template <class T>
__global__ void accumulate_kernel(const T *__restrict__ frames, int nframes, size_t frame_stride,
                                  int64_t *__restrict__ max_value, int64_t *__restrict__ sum_value,
                                  uint32_t *__restrict__ valid_count, size_t npixels, T masked) {
    for (size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x; i < npixels;
         i += static_cast<size_t>(blockDim.x) * gridDim.x) {
        int64_t mx = max_value[i];
        int64_t sm = sum_value[i];
        uint32_t c = valid_count[i];
        for (int k = 0; k < nframes; k++) {
            const T v = frames[k * frame_stride + i];
            if (v == masked)
                continue;
            const int64_t vi = static_cast<int64_t>(v);
            if (c == 0 || vi > mx)
                mx = vi;
            sm += vi;
            c++;
        }
        max_value[i] = mx;
        sum_value[i] = sm;
        valid_count[i] = c;
    }
}

// Another card's projection folded into this one, with the update Add makes: the maximum is only taken
// from a projection that counted the pixel.
__global__ void fold_kernel(const int64_t *__restrict__ other_max, const int64_t *__restrict__ other_sum,
                            const uint32_t *__restrict__ other_count, int64_t *__restrict__ max_value,
                            int64_t *__restrict__ sum_value, uint32_t *__restrict__ valid_count, size_t npixels) {
    for (size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x; i < npixels;
         i += static_cast<size_t>(blockDim.x) * gridDim.x) {
        if (other_count[i] == 0)
            continue;
        if (valid_count[i] == 0 || other_max[i] > max_value[i])
            max_value[i] = other_max[i];
        sum_value[i] += other_sum[i];
        valid_count[i] += other_count[i];
    }
}

template <class T>
void launch(const uint8_t *raw, int nframes, size_t frame_bytes, int64_t *max_value,
            int64_t *sum_value, uint32_t *valid_count, size_t npixels, int blocks,
            cudaStream_t stream) {
    T masked;
    if constexpr (std::is_signed_v<T>)
        masked = std::numeric_limits<T>::min();
    else
        masked = std::numeric_limits<T>::max();
    accumulate_kernel<T><<<blocks, 256, 0, stream>>>(reinterpret_cast<const T *>(raw), nframes,
                                                     frame_bytes / sizeof(T), max_value, sum_value,
                                                     valid_count, npixels, masked);
    cuda_err(cudaGetLastError());
}

}  // namespace

ShadowAccumulatorGPU::ShadowAccumulatorGPU(size_t npixels)
    : npixels(npixels),
      stream(std::make_shared<CudaStream>()),
      gpu_max(npixels),
      gpu_sum(npixels),
      gpu_count(npixels) {
    cuda_err(cudaMemsetAsync(gpu_max, 0, sizeof(int64_t) * npixels, *stream));
    cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(int64_t) * npixels, *stream));
    cuda_err(cudaMemsetAsync(gpu_count, 0, sizeof(uint32_t) * npixels, *stream));

    cuda_err(cudaGetDevice(&device));
    cudaDeviceProp prop{};
    cuda_err(cudaGetDeviceProperties(&prop, device));
    blocks = 8 * prop.multiProcessorCount;

    // Size the batch for the widest pixel type there is, so nothing has to be allocated once frames
    // start arriving - a cudaMalloc then would stall every worker behind it.
    raw_capacity = npixels * sizeof(uint32_t) * BATCH;
    raw = CudaDevicePtr<uint8_t>(raw_capacity);
    cuda_err(cudaStreamSynchronize(*stream));
}

bool ShadowAccumulatorGPU::Supports(const CompressedImage &image) {
    if (!BSLZ4DecoderGPU::Supports(image))
        return false;
    switch (image.GetMode()) {
        case CompressedImageMode::Int8:
        case CompressedImageMode::Uint8:
        case CompressedImageMode::Int16:
        case CompressedImageMode::Uint16:
        case CompressedImageMode::Int32:
        case CompressedImageMode::Uint32:
            return true;
        default:
            return false;
    }
}

void ShadowAccumulatorGPU::EnsureRawCapacity(size_t bytes_per_frame) {
    if (bytes_per_frame * BATCH > raw_capacity) {
        // The constructor already sized this for the widest pixel type, so in practice this only
        // runs if that guess was too small. cudaMalloc and cudaFree both synchronise the whole
        // device, which is why it is never done per frame.
        FoldPending();
        cuda_err(cudaStreamSynchronize(*stream));
        raw_capacity = bytes_per_frame * BATCH;
        raw = CudaDevicePtr<uint8_t>(raw_capacity);
    }
    // Only once the pending batch has been folded: FoldPending strides the buffer by frame_bytes,
    // so it has to keep describing the frames already in it until they are gone.
    frame_bytes = bytes_per_frame;
}

// Fold the frames decoded so far into the projection. One pass over the accumulator for the whole
// batch rather than one per frame.
void ShadowAccumulatorGPU::FoldPending() {
    if (pending == 0)
        return;
    switch (pending_mode) {
        case CompressedImageMode::Int8:
            launch<int8_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
        case CompressedImageMode::Uint8:
            launch<uint8_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
        case CompressedImageMode::Int16:
            launch<int16_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
        case CompressedImageMode::Uint16:
            launch<uint16_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
        case CompressedImageMode::Int32:
            launch<int32_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
        case CompressedImageMode::Uint32:
            launch<uint32_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "ShadowAccumulatorGPU: unsupported image mode");
    }
    pending = 0;
}

void ShadowAccumulatorGPU::Add(const CompressedImage &image) {
    if (!Supports(image))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "ShadowAccumulatorGPU: image cannot be decoded on the device");
    if (static_cast<size_t>(image.GetWidth()) * image.GetHeight() != npixels)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "ShadowAccumulatorGPU: image size does not match the detector");

    // A batch holds one pixel type and one frame size; a change of either closes the batch first,
    // while frame_bytes and pending_mode still describe the frames already in it.
    if (pending > 0 && (image.GetMode() != pending_mode
                        || image.GetUncompressedSize() != frame_bytes))
        FoldPending();

    if (!decoder || image.GetUncompressedSize() != frame_bytes)
        decoder = std::make_unique<BSLZ4DecoderGPU>(image.GetUncompressedSize(), stream);
    EnsureRawCapacity(image.GetUncompressedSize());
    pending_mode = image.GetMode();

    decoder->Decode(image, raw.get() + static_cast<size_t>(pending) * frame_bytes);
    // The decode has to be known good before its frame is counted, and the flag only arrives once
    // the stream has drained.
    cuda_err(cudaStreamSynchronize(*stream));
    decoder->ThrowIfDecodeFailed();

    pending++;
    frames++;
    if (pending == BATCH)
        FoldPending();
}

void ShadowAccumulatorGPU::Fold(ShadowAccumulatorGPU &other) {
    {
        // The other card's last batch is folded on its own stream, with its card current.
        ScopedGpuPin pin(other.device);
        other.FoldPending();
        cuda_err(cudaStreamSynchronize(*other.stream));
    }
    FoldPending();
    CudaDevicePtr<int64_t> other_max(npixels), other_sum(npixels);
    CudaDevicePtr<uint32_t> other_count(npixels);
    cuda_err(cudaMemcpyPeerAsync(other_max, device, other.gpu_max, other.device, sizeof(int64_t) * npixels, *stream));
    cuda_err(cudaMemcpyPeerAsync(other_sum, device, other.gpu_sum, other.device, sizeof(int64_t) * npixels, *stream));
    cuda_err(cudaMemcpyPeerAsync(other_count, device, other.gpu_count, other.device, sizeof(uint32_t) * npixels,
                                 *stream));
    fold_kernel<<<blocks, 256, 0, *stream>>>(other_max, other_sum, other_count, gpu_max, gpu_sum, gpu_count, npixels);
    cuda_err(cudaGetLastError());
    // The temporaries are pooled: they go once the stream has used them (see CudaDevicePtr).
    cuda_err(cudaStreamSynchronize(*stream));
    frames += other.frames;
}

void ShadowAccumulatorGPU::Download(std::vector<int64_t> &max_value, std::vector<int64_t> &sum_value,
                                    std::vector<uint32_t> &valid_count) {
    FoldPending();
    max_value.resize(npixels);
    sum_value.resize(npixels);
    valid_count.resize(npixels);
    cuda_err(cudaMemcpyAsync(max_value.data(), gpu_max.get(), sizeof(int64_t) * npixels,
                             cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(sum_value.data(), gpu_sum.get(), sizeof(int64_t) * npixels,
                             cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(valid_count.data(), gpu_count.get(), sizeof(uint32_t) * npixels,
                             cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaStreamSynchronize(*stream));
}

std::vector<uint32_t> ShadowAccumulatorGPU::Mask(const ShadowMaskSetup &setup, const std::vector<uint32_t> &pixel_mask) {
    FoldPending();
    return ShadowMaskOnDevice(setup, pixel_mask, gpu_max, gpu_sum, gpu_count, frames, *stream);
}

std::vector<float> ShadowAccumulatorGPU::MeanProjection(const std::vector<uint32_t> &pixel_mask) {
    FoldPending();
    return MeanProjectionOnDevice(pixel_mask, gpu_sum, gpu_count, npixels, *stream);
}