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

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <cstring>
#include <random>
#include <vector>

#include "../common/PixelMask.h"
#include "../compression/JFJochCompressor.h"
#include "../image_analysis/image_preprocessing/BSLZ4DecoderGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"

// The images the device reads straight out of host memory where gpu_direct_host_read() holds (a
// unified-memory system), against the host decoder and the CPU preprocessor. The host buffers are
// page-locked, so a discrete card can read them too: run with JFJOCH_FORCE_DIRECT_HOST_READ=1 to take
// the direct path there; without it these check the copying path on the same inputs.
namespace {

bool SameStats(const ImageStatistics &a, const ImageStatistics &b) {
    return a.max_value == b.max_value && a.min_value == b.min_value
        && a.masked_pixel_count == b.masked_pixel_count
        && a.error_pixel_count == b.error_pixel_count
        && a.saturated_pixel_count == b.saturated_pixel_count;
}

std::vector<uint32_t> MakeImage(size_t npixels, uint32_t seed) {
    std::mt19937 rng(seed);
    std::vector<uint32_t> img(npixels, 0);
    for (size_t i = npixels / 4; i < npixels / 2; i++)
        img[i] = rng() % 11;
    for (size_t s = 0; s < 500; s++)
        img[rng() % npixels] = 30000 + rng() % 50000;
    for (size_t s = 0; s < 300; s++)
        img[rng() % npixels] = UINT32_MAX;
    return img;
}

void CheckSameAsCPU(const DiffractionExperiment &x, const PixelMask &mask, const ImagePreprocessorBufferGPU &gpu_buf,
                    const ImageStatistics &gpu_stats, const uint8_t *raw, CompressedImageMode mode) {
    const size_t npixels = x.GetPixelsNum();
    ImagePreprocessorCPU cpu_pre(x, mask);
    ImagePreprocessorBuffer cpu_buf(npixels);
    const ImageStatistics cpu_stats = cpu_pre.Analyze(cpu_buf, raw, mode);
    CHECK(SameStats(cpu_stats, gpu_stats));
    size_t ndiff = 0;
    for (size_t i = 0; i < npixels; i++)
        ndiff += gpu_buf[i] != cpu_buf[i];
    CHECK(ndiff == 0);
}

}  // namespace

TEST_CASE("DirectHostRead_BSLZ4Decode", "[direct_host_read]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");
    WARN("gpu_direct_host_read() = " << gpu_direct_host_read());

    // An odd pixel count, so the chunk ends in a verbatim tail that is copied rather than decoded.
    const size_t width = 1001, height = 997, npixels = width * height;
    const auto original = MakeImage(npixels, 3);
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    std::vector<uint8_t> compressed = compressor.Compress(original);
    CudaRegisteredVector<uint8_t> reg(compressed);
    const CompressedImage image(compressed.data(), compressed.size(), width, height, CompressedImageMode::Uint32,
                                CompressionAlgorithm::BSHUF_LZ4);

    auto stream = std::make_shared<CudaStream>();
    BSLZ4DecoderGPU decoder(npixels * sizeof(uint32_t), stream);
    CudaDevicePtr<uint8_t> gpu_out(npixels * sizeof(uint32_t));
    decoder.Decode(image, gpu_out.get());
    REQUIRE(cudaStreamSynchronize(*stream) == cudaSuccess);
    decoder.ThrowIfDecodeFailed();

    std::vector<uint32_t> result(npixels);
    REQUIRE(cudaMemcpy(result.data(), gpu_out.get(), npixels * sizeof(uint32_t), cudaMemcpyDeviceToHost)
            == cudaSuccess);
    CHECK(result == original);
}

TEST_CASE("DirectHostRead_Preprocessor", "[direct_host_read]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    DetectorSetup detector = DetJF4M();
    detector.SaturationLimit(65535);
    DiffractionExperiment x(detector);
    x.DetectorDistance_mm(80).BeamX_pxl(1030).BeamY_pxl(1080);
    const size_t npixels = x.GetPixelsNum();
    PixelMask mask(x);
    std::vector<uint32_t> user_mask(mask.GetMask().size(), 0);
    for (size_t i = 0; i < npixels; i += 997) user_mask[i] = 1;
    mask.LoadUserMask(x, user_mask);

    auto stream = std::make_shared<CudaStream>();
    ImagePreprocessorGPU gpu_pre(x, mask, stream, /*copy_image_to_host=*/true);

    // The bitshuffle/LZ4 chunk, decoded and preprocessed by the fused kernel.
    std::vector<uint32_t> img = MakeImage(npixels, 5);
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    std::vector<uint8_t> compressed = compressor.Compress(img);
    CudaRegisteredVector<uint8_t> compressed_reg(compressed);
    {
        const CompressedImage image(compressed.data(), compressed.size(), x.GetXPixelsNum(), x.GetYPixelsNum(),
                                    CompressedImageMode::Uint32, CompressionAlgorithm::BSHUF_LZ4);
        ImagePreprocessorBufferGPU gpu_buf(npixels);
        ImageStatistics gpu_stats{};
        REQUIRE(gpu_pre.AnalyzeCompressed(gpu_buf, image, gpu_stats));
        CheckSameAsCPU(x, mask, gpu_buf, gpu_stats, reinterpret_cast<const uint8_t *>(img.data()),
                       CompressedImageMode::Uint32);
    }

    // An uncompressed image, as the host-upload path receives it.
    {
        std::vector<uint8_t> raw(reinterpret_cast<const uint8_t *>(img.data()),
                                 reinterpret_cast<const uint8_t *>(img.data()) + npixels * sizeof(uint32_t));
        CudaRegisteredVector<uint8_t> raw_reg(raw);
        ImagePreprocessorBufferGPU gpu_buf(npixels);
        const ImageStatistics gpu_stats = gpu_pre.Analyze(gpu_buf, raw.data(), CompressedImageMode::Uint32);
        CheckSameAsCPU(x, mask, gpu_buf, gpu_stats, raw.data(), CompressedImageMode::Uint32);
    }

    // A marCCD frame's stored 16-bit pixels, widened by the kernel.
    {
        std::mt19937 rng(7);
        std::vector<uint8_t> stored(npixels * sizeof(uint16_t));
        for (size_t i = 0; i < npixels; i++) {
            const uint16_t v = (i % 1500 == 0) ? 65535 : static_cast<uint16_t>(rng() % 70000);
            std::memcpy(stored.data() + i * sizeof(uint16_t), &v, sizeof(v));
        }
        CudaRegisteredVector<uint8_t> stored_reg(stored);
        const CompressedImage image(stored.data(), stored.size(), x.GetXPixelsNum(), x.GetYPixelsNum(),
                                    CompressedImageMode::Int32, CompressionAlgorithm::UINT16_WIDENED);
        ImagePreprocessorBufferGPU gpu_buf(npixels);
        ImageStatistics gpu_stats{};
        REQUIRE(gpu_pre.AnalyzeCompressed(gpu_buf, image, gpu_stats));
        std::vector<uint8_t> decoded;
        CheckSameAsCPU(x, mask, gpu_buf, gpu_stats, image.GetUncompressedPtr(decoded), CompressedImageMode::Int32);
    }
}

#endif