// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-only

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

#ifdef JFJOCH_USE_CUDA

#include <random>
#include <vector>

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

// The device-decode path does NOT decompress into a buffer and then preprocess it: one kernel
// un-transposes the bitshuffle blocks and applies the mask, the error marker, the saturation cap and
// the statistics as it goes, so the decompressed image never exists. That is a different code path
// from the host-upload one, not a reordering of it, and the thing it has to reproduce is the whole
// observable output - every preprocessed pixel AND every counter - against the CPU preprocessor fed
// the host-decompressed image.
//
// Masked, error and saturated pixels are the interesting part: their priority (masked > error >
// saturated) and their sentinel outputs (INT32_MIN / INT32_MIN / INT32_MAX) are decided in the fused
// kernel now, so the image below deliberately contains all three, and the mask deliberately covers
// some of them.
namespace {

DiffractionExperiment MakeExperiment(size_t saturation) {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(80).BeamX_pxl(1030).BeamY_pxl(1080);
    return x;
}

template <class T>
std::vector<T> MakeImage(size_t npixels, T err_value, uint32_t seed) {
    std::mt19937 rng(seed);
    std::vector<T> img(npixels, 0);
    // Sparse background with long runs, so LZ4 produces overlapping matches.
    for (size_t i = npixels / 4; i < npixels / 2; i++)
        img[i] = static_cast<T>(rng() % 11);
    // Bright spots, some above any plausible saturation cap.
    for (size_t s = 0; s < 500; s++) {
        const size_t c = rng() % npixels;
        for (size_t d = 0; d < 5 && c + d < npixels; d++)
            img[c + d] = static_cast<T>(30000 + (rng() % 5000));
    }
    // Explicit error markers, scattered.
    for (size_t s = 0; s < 300; s++)
        img[rng() % npixels] = err_value;
    return img;
}

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;
}

// One element size end to end: compress, decode+preprocess on the device, and compare against the
// host decompression fed through the CPU preprocessor.
template <class T>
void CheckFusedMatchesCPU(CompressedImageMode mode, T err_value, uint32_t seed) {
    DiffractionExperiment x = MakeExperiment(32000);
    const size_t npixels = x.GetPixelsNum();

    PixelMask mask(x);
    // Mask a deterministic scatter of pixels, so masked-vs-error-vs-saturated priority is exercised
    // rather than assumed.
    auto &m = const_cast<std::vector<uint32_t> &>(mask.GetMask());
    for (size_t i = 0; i < npixels; i += 997) m[i] = 1;
    for (size_t i = 13; i < npixels; i += 4001) m[i] = 1;

    const auto img = MakeImage<T>(npixels, err_value, seed);
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    const std::vector<uint8_t> compressed = compressor.Compress(img);
    const CompressedImage image(compressed.data(), compressed.size(),
                                x.GetXPixelsNum(), x.GetYPixelsNum(), mode,
                                CompressionAlgorithm::BSHUF_LZ4);
    REQUIRE(BSLZ4DecoderGPU::Supports(image));

    // Reference: host decompression + CPU preprocessing.
    ImagePreprocessorCPU cpu_pre(x, mask);
    ImagePreprocessorBuffer cpu_buf(npixels);
    std::vector<uint8_t> decompression_buffer;
    const uint8_t *raw = image.GetUncompressedPtr(decompression_buffer);
    const ImageStatistics cpu_stats = cpu_pre.Analyze(cpu_buf, raw, mode);

    // Under test: compressed chunk straight to the device, decoded and preprocessed in one pass.
    auto stream = std::make_shared<CudaStream>();
    ImagePreprocessorGPU gpu_pre(x, mask, stream, /*copy_image_to_host=*/true);
    ImagePreprocessorBufferGPU gpu_buf(npixels);
    ImageStatistics gpu_stats{};
    REQUIRE(gpu_pre.AnalyzeCompressed(gpu_buf, image, gpu_stats));

    INFO("mode " << static_cast<int>(mode));
    CHECK(SameStats(cpu_stats, gpu_stats));

    size_t ndiff = 0, first = 0;
    for (size_t i = 0; i < npixels; i++) {
        if (gpu_buf[i] != cpu_buf[i]) {
            if (ndiff == 0) first = i;
            ndiff++;
        }
    }
    INFO("first differing pixel " << first << " cpu " << cpu_buf[first] << " gpu " << gpu_buf[first]
         << " of " << ndiff << " differing");
    CHECK(ndiff == 0);
}

} // namespace

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

    CheckFusedMatchesCPU<uint32_t>(CompressedImageMode::Uint32, UINT32_MAX, 1);
    CheckFusedMatchesCPU<uint16_t>(CompressedImageMode::Uint16, UINT16_MAX, 2);
    CheckFusedMatchesCPU<int32_t>(CompressedImageMode::Int32, INT32_MIN, 3);
    CheckFusedMatchesCPU<int16_t>(CompressedImageMode::Int16, INT16_MIN, 4);
}

// The host-upload path must keep producing exactly what it did - it is still what every non-LZ4
// image takes - so the two entry points are held against each other on the same frame.
TEST_CASE("ImagePreprocessorGPU_FusedMatchesHostUpload", "[ImagePreprocessorGPU]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    DiffractionExperiment x = MakeExperiment(32000);
    const size_t npixels = x.GetPixelsNum();
    PixelMask mask(x);
    auto &m = const_cast<std::vector<uint32_t> &>(mask.GetMask());
    for (size_t i = 0; i < npixels; i += 1301) m[i] = 1;

    const auto img = MakeImage<uint32_t>(npixels, UINT32_MAX, 77);
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    const std::vector<uint8_t> compressed = compressor.Compress(img);
    const CompressedImage image(compressed.data(), compressed.size(),
                                x.GetXPixelsNum(), x.GetYPixelsNum(), CompressedImageMode::Uint32,
                                CompressionAlgorithm::BSHUF_LZ4);

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

    ImagePreprocessorBufferGPU fused_buf(npixels);
    ImageStatistics fused_stats{};
    REQUIRE(pre.AnalyzeCompressed(fused_buf, image, fused_stats));

    // Same engine, same frame, but decompressed on the host and uploaded.
    std::vector<uint8_t> decompression_buffer;
    const uint8_t *raw = image.GetUncompressedPtr(decompression_buffer);
    ImagePreprocessorBufferGPU upload_buf(npixels);
    const ImageStatistics upload_stats = pre.Analyze(upload_buf, raw, CompressedImageMode::Uint32);

    CHECK(SameStats(fused_stats, upload_stats));
    size_t ndiff = 0;
    for (size_t i = 0; i < npixels; i++) if (fused_buf[i] != upload_buf[i]) ndiff++;
    CHECK(ndiff == 0);

    // Decoding on the device replaced a host decompression, so the cost is still reported as one.
    CHECK(pre.GetLastDecompressionTime_s() > 0.0f);
}

#endif