// 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 <vector>

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

// A 16-bit unsigned source keeps its width through the GPU pipeline (PreprocessedPixel.h): masked
// and saturated pixels ride on two reserved codes at the top of the range and are widened again on
// load. That has to give back exactly what the 32-bit path gives, for every value a 16-bit detector
// can produce and for every saturation limit it can declare - so both are run over the same frame
// and compared pixel for pixel, at both entry points (host upload and device decode).
namespace {

// The values the encoding turns on: the two reserved codes, their neighbours, and the pixels either
// side of the saturation limit.
std::vector<uint16_t> InterestingValues(int64_t saturation) {
    std::vector<uint16_t> v{0, 1, 2, 100, 0xFFFC, 0xFFFD, 0xFFFE, 0xFFFF};
    for (int64_t d = -2; d <= 2; d++) {
        const int64_t s = saturation + d;
        if (s >= 0 && s <= UINT16_MAX)
            v.push_back(static_cast<uint16_t>(s));
    }
    return v;
}

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

// Gather() reads through PixelView, so it reports the image in the pipeline's int32 convention
// whatever width it is stored in - which is the only way to read a narrow image back.
std::vector<int32_t> ReadBack(const ImagePreprocessorBuffer &buffer) {
    std::vector<int32_t> out(buffer.size());
    std::vector<uint32_t> index;
    std::vector<int32_t> values;
    for (size_t i = 0; i < buffer.size(); i += UINT16_MAX) {
        const size_t n = std::min<size_t>(UINT16_MAX, buffer.size() - i);
        index.resize(n);
        for (size_t k = 0; k < n; k++)
            index[k] = static_cast<uint32_t>(i + k);
        buffer.Gather(index, values);
        std::copy(values.begin(), values.end(), out.begin() + i);
    }
    return out;
}

void CheckNarrowMatchesWide(std::optional<int64_t> saturation, bool device_decode) {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(80).BeamX_pxl(1030).BeamY_pxl(1080);
    // A 16-bit unsigned container, which is the only source the narrow path is taken for.
    x.BitDepthImage(16).PixelSigned(false);
    x.Detector().SaturationLimit(saturation);
    const size_t npixels = x.GetPixelsNum();

    const std::vector<uint16_t> values = InterestingValues(x.GetSaturationLimit());
    std::vector<uint16_t> img(npixels);
    for (size_t i = 0; i < npixels; i++)
        img[i] = values[i % values.size()];

    // 11 is coprime with the number of values, so every value appears both masked and unmasked.
    PixelMask mask(x);
    auto &m = const_cast<std::vector<uint32_t> &>(mask.GetMask());
    for (size_t i = 3; i < npixels; i += 11)
        m[i] = 1;

    auto stream = std::make_shared<CudaStream>();
    // The 32-bit path, which is what every consumer saw before, is the reference.
    ImagePreprocessorGPU wide_pre(x, mask, stream, /*copy_image_to_host=*/true);
    ImagePreprocessorBufferGPU wide_buffer(npixels);
    // Telling the engine nothing has to come back to the host is what lets the image keep its width.
    ImagePreprocessorGPU narrow_pre(x, mask, stream, /*copy_image_to_host=*/false);
    ImagePreprocessorBufferGPU narrow_buffer(npixels, /*host_mirror=*/false);

    ImageStatistics wide_stats{};
    ImageStatistics narrow_stats{};
    if (device_decode) {
        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::Uint16, CompressionAlgorithm::BSHUF_LZ4);
        REQUIRE(wide_pre.AnalyzeCompressed(wide_buffer, image, wide_stats));
        REQUIRE(narrow_pre.AnalyzeCompressed(narrow_buffer, image, narrow_stats));
    } else {
        const auto *raw = reinterpret_cast<const uint8_t *>(img.data());
        wide_stats = wide_pre.Analyze(wide_buffer, raw, CompressedImageMode::Uint16);
        narrow_stats = narrow_pre.Analyze(narrow_buffer, raw, CompressedImageMode::Uint16);
    }

    INFO("saturation limit " << x.GetSaturationLimit() << (device_decode ? " device decode" : " host upload"));
    REQUIRE(narrow_buffer.IsNarrow());
    REQUIRE_FALSE(wide_buffer.IsNarrow());
    CHECK(SameStats(wide_stats, narrow_stats));

    const std::vector<int32_t> narrow_read = ReadBack(narrow_buffer);
    size_t ndiff = 0;
    size_t first = 0;
    for (size_t i = 0; i < npixels; i++) {
        if (narrow_read[i] != wide_buffer[i]) {
            if (ndiff == 0) first = i;
            ndiff++;
        }
    }
    INFO("first differing pixel " << first << " source value " << img[first]
         << (m[first] ? " (masked)" : "")
         << " wide " << wide_buffer[first] << " narrow " << narrow_read[first]
         << " of " << ndiff << " differing");
    CHECK(ndiff == 0);
}

} // namespace

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

    for (bool device_decode : {false, true}) {
        // No declared limit: the limit is the whole 16-bit range, and nothing can saturate.
        CheckNarrowMatchesWide(std::nullopt, device_decode);
        // A limit with plenty of room below the reserved codes - the ordinary case.
        CheckNarrowMatchesWide(5000, device_decode);
        // The largest limit that still leaves a code of its own, and the two above it.
        CheckNarrowMatchesWide(UINT16_MAX - 2, device_decode);
        CheckNarrowMatchesWide(UINT16_MAX - 1, device_decode);
        CheckNarrowMatchesWide(UINT16_MAX, device_decode);
    }
}

#endif