// 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 <limits>
#include <random>
#include <vector>

#include "../common/CUDAWrapper.h"
#include "../common/PixelMask.h"
#include "../common/ScanResultGenerator.h"
#include "../compression/JFJochCompressor.h"
#include "../image_analysis/IndexAndRefine.h"
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorCPU.h"

#ifdef JFJOCH_USE_CUDA
#include "../compression/bitshuffle/bitshuffle.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorGPU.h"
#endif

// pixel_sum on the CPU and GPU preprocessors, to the card's definition (fpga/hls/pixel_calc.cpp): signed,
// and over the pixels that are not masked, not error and not saturated.
namespace {
    constexpr int64_t SATURATION = 110;

    template<class T>
    void CheckEveryPixelType(CompressedImageMode mode, T err_value) {
        DiffractionExperiment x(DetDECTRIS(64, 4, "Test", {}).BitDepthImage(8 * sizeof(T)).SaturationLimit(SATURATION));
        x.PixelSigned(std::is_signed_v<T>);

        // Pixels 0-63 hold a masked pixel, so they are tested one by one; the words after them are whole.
        std::vector<uint32_t> user_mask(x.GetPixelsNum(), 0);
        user_mask[3] = 1;
        user_mask[4] = 1;
        user_mask[5] = 1;
        PixelMask mask(x);
        mask.LoadUserMask(x, user_mask);

        std::vector<T> image(x.GetPixelsNum(), 0);
        image[1] = 100;
        image[3] = 50;                  // masked
        image[4] = err_value;           // masked before it is an error
        image[5] = SATURATION;          // masked before it is saturated
        image[70] = 20;
        image[72] = SATURATION - 1;
        image[64] = err_value;
        image[65] = SATURATION;
        image[66] = SATURATION + 1;

        int64_t expected = 100 + 20 + SATURATION - 1;
        if (std::is_signed_v<T>) {
            image[71] = static_cast<T>(-30);
            expected -= 30;
        }

        ImagePreprocessorCPU preprocessor(x, mask);
        ImagePreprocessorBuffer buffer(x.GetPixelsNum());
        const auto input = reinterpret_cast<const uint8_t *>(image.data());
        const ImageStatistics stats = preprocessor.Analyze(buffer, input, mode);
        CHECK(stats.pixel_sum == expected);
        CHECK(stats.masked_pixel_count == 3);
        CHECK(stats.error_pixel_count == 1);
        CHECK(stats.saturated_pixel_count == 2);

        // In blocks, as a bitshuffled image is decoded.
        ImageStatistics blocks{};
        preprocessor.AnalyzeBlock(buffer, 0, input, 40, mode, blocks);
        preprocessor.AnalyzeBlock(buffer, 40, input + 40 * sizeof(T), x.GetPixelsNum() - 40, mode, blocks);
        CHECK(blocks.pixel_sum == expected);
    }
}

TEST_CASE("ImagePreprocessorCPU_PixelSum", "[ImagePreprocessor]") {
    CheckEveryPixelType<int8_t>(CompressedImageMode::Int8, INT8_MIN);
    CheckEveryPixelType<uint8_t>(CompressedImageMode::Uint8, UINT8_MAX);
    CheckEveryPixelType<int16_t>(CompressedImageMode::Int16, INT16_MIN);
    CheckEveryPixelType<uint16_t>(CompressedImageMode::Uint16, UINT16_MAX);
    CheckEveryPixelType<int32_t>(CompressedImageMode::Int32, INT32_MIN);
    CheckEveryPixelType<uint32_t>(CompressedImageMode::Uint32, UINT32_MAX);
}

// Every pixel just under the limit of a large detector: far past 32 bits, in a sum of 32-bit pixels.
TEST_CASE("ImagePreprocessorCPU_PixelSum_NoOverflow", "[ImagePreprocessor]") {
    DiffractionExperiment x(DetDECTRIS(4096, 4096, "Test", {}).BitDepthImage(32));
    x.PixelSigned(true);
    const PixelMask mask(x);

    const int32_t value = static_cast<int32_t>(x.GetSaturationLimit()) - 1;
    const std::vector<int32_t> image(x.GetPixelsNum(), value);

    ImagePreprocessorCPU preprocessor(x, mask);
    ImagePreprocessorBuffer buffer(x.GetPixelsNum());
    const ImageStatistics stats = preprocessor.Analyze(buffer, reinterpret_cast<const uint8_t *>(image.data()),
                                                       CompressedImageMode::Int32);
    CHECK(stats.pixel_sum == static_cast<int64_t>(value) * static_cast<int64_t>(image.size()));
    CHECK(stats.pixel_sum > INT32_MAX);
}

// A mask a run determines (a beam-stop shadow, hot pixels, inactive pixels) takes pixels out of the image, not out of the sum.
TEST_CASE("ImagePreprocessorCPU_PixelSum_RunMask", "[ImagePreprocessor]") {
    DiffractionExperiment x(DetDECTRIS(64, 4, "Test", {}).BitDepthImage(32));
    x.PixelSigned(true);

    std::vector<uint32_t> user_mask(x.GetPixelsNum(), 0), shadow(x.GetPixelsNum(), 0), hot(x.GetPixelsNum(), 0),
            inactive(x.GetPixelsNum(), 0);
    user_mask[1] = 1;               // the device's
    shadow[1] = 1;                  // and the run's: the device mask decides
    shadow[2] = 1;
    shadow[3] = 1;                  // error pixel under the shadow
    hot[4] = 1;
    hot[5] = 1;                     // saturated pixel, hot
    inactive[6] = 1;
    inactive[7] = 1;                // error pixel, inactive
    inactive[1] = 1;                // the device mask decides here too

    PixelMask device(x);
    device.LoadUserMask(x, user_mask);
    PixelMask with_run = device;
    with_run.LoadBeamStopMask(x, shadow);
    with_run.LoadHotPixelMask(x, hot);
    with_run.AddDeadPixelMask(x, inactive);

    std::vector<int32_t> image(x.GetPixelsNum(), 0);
    image[0] = 11;
    image[1] = 1000;
    image[2] = 22;
    image[3] = INT32_MIN;
    image[4] = -33;
    image[5] = INT32_MAX;
    image[6] = 44;
    image[7] = INT32_MIN;
    const auto input = reinterpret_cast<const uint8_t *>(image.data());

    ImagePreprocessorBuffer buffer(x.GetPixelsNum());
    ImagePreprocessorCPU device_only(x, device);
    CHECK(device_only.Analyze(buffer, input, CompressedImageMode::Int32).pixel_sum == 11 + 22 - 33 + 44);

    ImagePreprocessorCPU run(x, with_run);
    const ImageStatistics stats = run.Analyze(buffer, input, CompressedImageMode::Int32);
    CHECK(stats.pixel_sum == 11 + 22 - 33 + 44);
    CHECK(stats.masked_pixel_count == 7);
    CHECK(stats.max_value == 11);
    CHECK(buffer[2] == INT32_MIN);
    CHECK(buffer[6] == INT32_MIN);
}

// CPU or GPU, whichever this machine runs, from the image to the vector the master file is written from.
TEST_CASE("MXAnalysisWithoutFPGA_PixelSum", "[ImagePreprocessor]") {
    constexpr size_t width = 1030, height = 1064;

    DiffractionExperiment experiment(DetDECTRIS(width, height, "Test", {}).BitDepthImage(32));
    experiment.ImagesPerTrigger(1).NumTriggers(1).PixelSigned(true)
            .DetectorDistance_mm(75).BeamX_pxl(515).BeamY_pxl(532).IncidentEnergy_keV(12.4);

    std::vector<uint32_t> user_mask(width * height, 0), hot(width * height, 0);
    user_mask[1000] = 1;
    hot[10] = 1;
    PixelMask pixel_mask(experiment);
    pixel_mask.LoadUserMask(experiment, user_mask);
    pixel_mask.LoadHotPixelMask(experiment, hot);

    AzimuthalIntegrationMapping mapping(experiment, pixel_mask);
    IndexAndRefine index_and_refine(experiment, nullptr);
    MXAnalysisWithoutFPGA analysis(experiment, mapping, pixel_mask, index_and_refine);

    std::vector<int32_t> pixels(width * height, 0);
    pixels[10] = 3'000'000;
    pixels[11] = 2'000'000'000;
    pixels[12] = -40;
    pixels[1000] = 777;
    pixels[2000] = INT32_MIN;
    pixels[3000] = INT32_MAX;
    const int64_t expected = int64_t{3'000'000} + 2'000'000'000 - 40;

    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    const std::vector<uint8_t> compressed = compressor.Compress(pixels);

    DataMessage message{};
    message.number = 0;
    message.image = CompressedImage(compressed.data(), compressed.size(), width, height,
                                    CompressedImageMode::Int32, CompressionAlgorithm::BSHUF_LZ4);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.enable = false;

    AzimuthalIntegrationProfile profile(mapping);
    analysis.Analyze(message, profile, settings);
    REQUIRE(message.pixel_sum);
    CHECK(message.pixel_sum.value() == expected);

    ScanResultGenerator scan_result(experiment);
    scan_result.Add(message);
    EndMessage end{};
    scan_result.FillEndMessage(end);
    CHECK(end.pixel_sum == std::vector<int64_t>{expected});
}

#ifdef JFJOCH_USE_CUDA
// All three routes through the kernels against the CPU, with a run's shadow over the device mask: the upload, the decode fused with the preprocessing,
// and the two-kernel route of a bitshuffle block too large for the fused one.
TEST_CASE("ImagePreprocessorGPU_PixelSum", "[ImagePreprocessorGPU]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    DiffractionExperiment x(DetDECTRIS(4096, 4096, "Test", {}).BitDepthImage(32));
    x.PixelSigned(true);
    const size_t npixels = x.GetPixelsNum();

    std::vector<uint32_t> user_mask(npixels, 0);
    std::vector<uint32_t> shadow(npixels, 0);
    for (size_t i = 0; i < npixels; i += 997)
        user_mask[i] = 1;
    for (size_t i = 0; i < npixels; i += 501)
        shadow[i] = 1;
    PixelMask mask(x);
    mask.LoadUserMask(x, user_mask);
    mask.LoadBeamStopMask(x, shadow);
    std::vector<uint32_t> inactive(npixels, 0);
    for (size_t i = 0; i < npixels; i += 1003)
        inactive[i] = 1;
    mask.AddDeadPixelMask(x, inactive);

    // Mostly near the limit, so the sum needs far more than 32 bits; negative pixels and markers among them.
    std::mt19937 rng(11);
    const int32_t limit = static_cast<int32_t>(x.GetSaturationLimit());
    std::vector<int32_t> image(npixels);
    for (auto &v: image) {
        const uint32_t r = rng() % 1000;
        v = r == 0 ? INT32_MIN : r == 1 ? limit : r < 100 ? -static_cast<int32_t>(rng() % 100000) : limit - 1 - static_cast<int32_t>(r);
    }

    int64_t expected = 0;
    for (size_t i = 0; i < npixels; i++)
        if (user_mask[i] == 0 && image[i] != INT32_MIN && image[i] < limit)
            expected += image[i];
    REQUIRE(expected > INT32_MAX);

    const auto input = reinterpret_cast<const uint8_t *>(image.data());
    ImagePreprocessorCPU cpu(x, mask);
    ImagePreprocessorBuffer cpu_buffer(npixels);
    CHECK(cpu.Analyze(cpu_buffer, input, CompressedImageMode::Int32).pixel_sum == expected);

    auto stream = std::make_shared<CudaStream>();
    ImagePreprocessorGPU gpu(x, mask, stream, /*copy_image_to_host=*/false);
    ImagePreprocessorBufferGPU gpu_buffer(npixels);

    CHECK(gpu.Analyze(gpu_buffer, input, CompressedImageMode::Int32).pixel_sum == expected);

    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    const std::vector<uint8_t> compressed = compressor.Compress(image);
    ImageStatistics fused{};
    REQUIRE(gpu.AnalyzeCompressed(gpu_buffer, CompressedImage(compressed.data(), compressed.size(), x.GetXPixelsNum(),
                                                              x.GetYPixelsNum(), CompressedImageMode::Int32,
                                                              CompressionAlgorithm::BSHUF_LZ4), fused));
    CHECK(fused.pixel_sum == expected);

    // The container by hand, with 32 kB blocks: 64-bit size and 32-bit block size in bytes, big-endian.
    constexpr size_t block_elements = 8192;
    std::vector<uint8_t> large_blocks(12 + bshuf_compress_lz4_bound(npixels, sizeof(int32_t), block_elements));
    const uint64_t total_bytes = npixels * sizeof(int32_t);
    for (int i = 0; i < 8; i++)
        large_blocks[i] = static_cast<uint8_t>(total_bytes >> (56 - 8 * i));
    const uint32_t block_bytes = block_elements * sizeof(int32_t);
    for (int i = 0; i < 4; i++)
        large_blocks[8 + i] = static_cast<uint8_t>(block_bytes >> (24 - 8 * i));
    const int64_t written = bshuf_compress_lz4(image.data(), large_blocks.data() + 12, npixels, sizeof(int32_t),
                                               block_elements);
    REQUIRE(written > 0);
    ImageStatistics two_kernels{};
    REQUIRE(gpu.AnalyzeCompressed(gpu_buffer, CompressedImage(large_blocks.data(), 12 + written, x.GetXPixelsNum(),
                                                              x.GetYPixelsNum(), CompressedImageMode::Int32,
                                                              CompressionAlgorithm::BSHUF_LZ4), two_kernels));
    CHECK(two_kernels.pixel_sum == expected);
}
#endif