// 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 <cstring>
#include <atomic>
#include <random>
#include <thread>
#include <vector>

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

// The miniCBF and marCCD frames reach the GPU as their files store them, and the device has to turn
// them into exactly the int32 image the host decoder would have: the same pixels, the same refusal of
// a stream that ends early.
namespace {

// x-CBF_BYTE_OFFSET with the smallest container for each delta; force64 puts a delta in the 15-byte
// form whatever its size, which no real writer does but the format allows.
void AppendDelta(std::vector<uint8_t> &out, int64_t delta, bool force64) {
    auto put = [&out](int64_t v, int bytes) {
        for (int i = 0; i < bytes; i++) out.push_back(static_cast<uint8_t>(v >> (8 * i)));
    };
    if (!force64 && delta > -128 && delta <= 127) { put(delta, 1); return; }
    put(-128, 1);
    if (!force64 && delta > -32768 && delta <= 32767) { put(delta, 2); return; }
    put(-32768, 2);
    if (!force64 && delta > INT32_MIN && delta <= INT32_MAX) { put(delta, 4); return; }
    put(INT32_MIN, 4);
    put(delta, 8);
}

std::vector<uint8_t> EncodeByteOffset(const std::vector<int32_t> &pixels, std::mt19937 &rng) {
    std::vector<uint8_t> out;
    int64_t prev = 0;
    for (const int32_t p : pixels) {
        AppendDelta(out, p - prev, rng() % 5000 == 0);
        prev = p;
    }
    return out;
}

// Mostly small deltas, with runs of every escape: int16 and int32 deltas whose payload bytes include
// 0x80 (the escape byte, which must not be taken for one inside a payload), and negative pixels.
std::vector<int32_t> MakePixels(size_t n, uint32_t seed) {
    std::mt19937 rng(seed);
    std::vector<int32_t> p(n);
    for (size_t i = 0; i < n; i++) {
        const uint32_t r = rng() % 100;
        if (r < 80) p[i] = static_cast<int32_t>(rng() % 40);
        else if (r < 90) p[i] = 128 + static_cast<int32_t>(rng() % 1000) * 256;   // payload bytes 0x80
        else if (r < 95) p[i] = static_cast<int32_t>(rng() % 2000000) - 1000000;
        else if (r < 98) p[i] = -1;
        else p[i] = static_cast<int32_t>(rng());
    }
    // A long run of escapes back to back, so whole tiles hold nothing but multi-byte tokens.
    for (size_t i = n / 3; i < n / 3 + 5000 && i < n; i++)
        p[i] = (i % 2) ? 1000000 : -1000000;
    return p;
}

std::vector<int32_t> DecodeOnDevice(const std::vector<uint8_t> &stream_bytes, size_t nx, size_t ny,
                                    CBFDecoderGPU &decoder, CudaStream &stream) {
    const CompressedImage image(stream_bytes.data(), stream_bytes.size(), nx, ny, CompressedImageMode::Int32,
                                CompressionAlgorithm::CBF_BYTE_OFFSET);
    CudaDevicePtr<int32_t> out(nx * ny);
    decoder.Decode(image, out);
    std::vector<int32_t> ret(nx * ny);
    REQUIRE(cudaMemcpyAsync(ret.data(), out.get(), ret.size() * sizeof(int32_t), cudaMemcpyDeviceToHost, stream)
            == cudaSuccess);
    REQUIRE(cudaStreamSynchronize(stream) == cudaSuccess);
    decoder.ThrowIfDecodeFailed();
    return ret;
}

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

}  // namespace

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

    auto stream = std::make_shared<CudaStream>();
    CBFDecoderGPU decoder(stream);
    std::mt19937 rng(11);
    for (const auto &[nx, ny] : std::vector<std::pair<size_t, size_t>>{{1, 1}, {7, 3}, {1000, 1013}, {2463, 1000}}) {
        const auto pixels = MakePixels(nx * ny, static_cast<uint32_t>(nx * 31 + ny));
        auto bytes = EncodeByteOffset(pixels, rng);
        // A real file carries the CIF trailer after the stream; nothing past the last pixel is read.
        for (const char c : std::string("\r\n--CIF-BINARY-FORMAT-SECTION----\r\n;\r\n\x80\x80\x80"))
            bytes.push_back(static_cast<uint8_t>(c));

        std::vector<int32_t> host(nx * ny);
        DecodeCBFByteOffset(bytes.data(), bytes.size(), host.data(), host.size());
        REQUIRE(host == pixels);
        INFO(nx << " x " << ny);
        CHECK(DecodeOnDevice(bytes, nx, ny, decoder, *stream) == host);
    }
}

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

    auto stream = std::make_shared<CudaStream>();
    CBFDecoderGPU decoder(stream);
    std::mt19937 rng(5);
    const size_t nx = 300, ny = 200;
    auto pixels = MakePixels(nx * ny, 3);
    pixels.back() = pixels[pixels.size() - 2] + 1000000;   // the last delta takes the 7-byte form
    const auto bytes = EncodeByteOffset(pixels, rng);

    // One token short, and cut inside the last (multi-byte) token: both decoders refuse both.
    const std::vector<uint8_t> short_stream(bytes.begin(), bytes.end() - 2);
    for (const auto &s : {std::vector<uint8_t>(bytes.begin(), bytes.end() - 7), short_stream}) {
        std::vector<int32_t> host(nx * ny);
        CHECK_THROWS(DecodeCBFByteOffset(s.data(), s.size(), host.data(), host.size()));
        CHECK_THROWS(DecodeOnDevice(s, nx, ny, decoder, *stream));
    }
}

// Through the preprocessor, as the image loop calls it: the stored encodings against the CPU
// preprocessor fed what the host decodes them to. The marCCD frame carries 65535, its saturation
// marker, which a Uint16 image would call an error pixel - widened to int32 it is saturated.
TEST_CASE("ImagePreprocessorGPU_StoredEncodingsMatchCPU", "[ImagePreprocessorGPU]") {
    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);

    std::mt19937 rng(17);
    const auto pixels = MakePixels(npixels, 9);
    const auto cbf = EncodeByteOffset(pixels, rng);
    std::vector<uint16_t> marccd(npixels);
    for (size_t i = 0; i < npixels; i++)
        marccd[i] = (i % 1500 == 0) ? 65535 : static_cast<uint16_t>(rng() % 70000);

    const std::vector<CompressedImage> images = {
        CompressedImage(cbf.data(), cbf.size(), x.GetXPixelsNum(), x.GetYPixelsNum(), CompressedImageMode::Int32,
                        CompressionAlgorithm::CBF_BYTE_OFFSET),
        CompressedImage(marccd.data(), marccd.size() * sizeof(uint16_t), x.GetXPixelsNum(), x.GetYPixelsNum(),
                        CompressedImageMode::Int32, CompressionAlgorithm::UINT16_WIDENED)
    };
    auto stream = std::make_shared<CudaStream>();
    ImagePreprocessorGPU gpu_pre(x, mask, stream, /*copy_image_to_host=*/true);
    ImagePreprocessorCPU cpu_pre(x, mask);
    for (const auto &image : images) {
        std::vector<uint8_t> decoded;
        const uint8_t *raw = image.GetUncompressedPtr(decoded);
        ImagePreprocessorBuffer cpu_buf(npixels);
        const ImageStatistics cpu_stats = cpu_pre.Analyze(cpu_buf, raw, CompressedImageMode::Int32);

        ImagePreprocessorBufferGPU gpu_buf(npixels);
        ImageStatistics gpu_stats{};
        REQUIRE(gpu_pre.AnalyzeCompressed(gpu_buf, image, gpu_stats));
        INFO("algorithm " << static_cast<int>(image.GetCompressionAlgorithm()));
        CHECK(SameStats(cpu_stats, gpu_stats));
        CHECK(cpu_stats.saturated_pixel_count > 0);
        size_t ndiff = 0;
        for (size_t i = 0; i < npixels; i++)
            ndiff += gpu_buf[i] != cpu_buf[i];
        CHECK(ndiff == 0);
    }
}

// The pre-scan hands the frames it kept to two consumers at once - the spot measurement and the
// defective-pixel search - each with its own engines, so one marCCD frame is uploaded by several
// preprocessors concurrently while each of them is also moving on to other frames.
TEST_CASE("ImagePreprocessorGPU_SharedStoredFramesManyEngines", "[ImagePreprocessorGPU]") {
    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::mt19937 rng(23);
    std::vector<std::vector<uint16_t>> frames(4, std::vector<uint16_t>(npixels));
    for (auto &f : frames)
        for (auto &v : f)
            v = static_cast<uint16_t>(rng() % 60000);

    std::vector<ImageStatistics> reference(frames.size());
    {
        ImagePreprocessorCPU cpu_pre(x, mask);
        ImagePreprocessorBuffer cpu_buf(npixels);
        for (size_t f = 0; f < frames.size(); f++) {
            std::vector<int32_t> widened(frames[f].begin(), frames[f].end());
            reference[f] = cpu_pre.Analyze(cpu_buf, reinterpret_cast<const uint8_t *>(widened.data()),
                                           CompressedImageMode::Int32);
        }
    }

    std::atomic<size_t> mismatches{0};
    std::vector<std::thread> threads;
    for (int t = 0; t < 4; t++)
        threads.emplace_back([&, t] {
            auto stream = std::make_shared<CudaStream>();
            ImagePreprocessorGPU gpu_pre(x, mask, stream, /*copy_image_to_host=*/false);
            ImagePreprocessorBufferGPU gpu_buf(npixels, /*host_mirror=*/false);
            for (int i = 0; i < 200; i++) {
                const size_t f = (i + t) % frames.size();
                const CompressedImage image(frames[f].data(), npixels * sizeof(uint16_t), x.GetXPixelsNum(),
                                            x.GetYPixelsNum(), CompressedImageMode::Int32,
                                            CompressionAlgorithm::UINT16_WIDENED);
                ImageStatistics stats{};
                gpu_pre.AnalyzeCompressed(gpu_buf, image, stats);
                if (!SameStats(stats, reference[f]))
                    mismatches++;
            }
        });
    for (auto &t : threads)
        t.join();
    CHECK(mismatches == 0);
}

#endif