// 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 <algorithm>
#include <chrono>

#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
#include "../image_analysis/azint/AzIntEngineGPU.h"
#include "../image_analysis/spot_finding/AdaptiveSpotFinderCPU.h"
#include "../image_analysis/spot_finding/AdaptiveSpotFinderGPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinderGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"

namespace {

// Build a realistic full-detector azimuthal-integration mapping (JF4M, ~4.5 MP) whose q-range spans
// most of the detector, so the timing runs over a representative pixel count.
DiffractionExperiment MakeExperiment() {
    DiffractionExperiment x(DetJF4M());
    x.DetectorDistance_mm(80).BeamX_pxl(1030).BeamY_pxl(1080);
    x.QSpacingForAzimInt_recipA(0.05).QRangeForAzimInt_recipA(0.05, 5.0);
    return x;
}

// Deterministic image: a low, slightly rippled background (well below any adaptive threshold) plus a
// grid of bright multi-pixel blobs that both finders must recover identically.
void FillTestImage(ImagePreprocessorBuffer &buffer, const DiffractionExperiment &x) {
    const size_t w = x.GetXPixelsNum();
    const size_t h = x.GetYPixelsNum();
    for (size_t i = 0; i < w * h; i++)
        buffer[i] = 8 + static_cast<int32_t>(i % 5);   // background 8..12 (mean 10)

    // Bright 3x3 blobs on a coarse grid, kept clear of the edges and the beam centre.
    for (size_t row = 300; row < h - 300; row += 450) {
        for (size_t col = 300; col < w - 300; col += 450) {
            for (int dr = -1; dr <= 1; dr++)
                for (int dc = -1; dc <= 1; dc++)
                    buffer[(row + dr) * w + (col + dc)] = 200;
        }
    }
}

SpotFindingSettings AdaptiveSettings() {
    SpotFindingSettings s{};
    s.adaptive_threshold = true;
    s.false_pixels_per_frame = 100.0f;
    s.min_pix_per_spot = 1;
    s.max_pix_per_spot = 50;
    s.high_resolution_limit = 0.0f;      // no resolution gate for the parity test
    s.low_resolution_limit = 1.0e6f;
    s.high_res_gap_Q_recipA = std::nullopt;
    return s;
}

std::vector<std::pair<int, int>> SortedCoords(const std::vector<DiffractionSpot> &spots) {
    std::vector<std::pair<int, int>> out;
    out.reserve(spots.size());
    for (const auto &s : spots)
        out.emplace_back(static_cast<int>(std::lround(s.RawCoord().y)),
                         static_cast<int>(std::lround(s.RawCoord().x)));
    std::sort(out.begin(), out.end());
    return out;
}

} // namespace

// Spot-finding functionality: the fused GPU engine must reproduce the reference CPU adaptive finder's
// spot list. The two share AdaptiveThreshold.h and the host connected-component extractor. They do
// NOT accumulate identically: the CPU sums each ring serially in double, while the GPU stages a
// block's contribution in float before reducing across blocks in double (see the comment on the
// kernel). So ring sigma can differ in the last bits, and since detection compares integer pixel
// values against the threshold, a threshold that crosses an integer flips every pixel of that value
// in the ring at once. That is the difference this test is bounding.
TEST_CASE("AdaptiveSpotFinderGPU_SpotFindingParity", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_SpotFindingParity");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    AdaptiveSpotFinderCPU cpu(mapping);
    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);

    cpu.SetResolutionMask(res_mask);
    gpu.SetResolutionMask(res_mask);
    const auto cpu_spots = cpu.Run(buffer, settings);
    const auto gpu_spots = gpu.Run(buffer, settings);

    INFO("cpu spots=" << cpu_spots.size() << " gpu spots=" << gpu_spots.size());
    REQUIRE(cpu_spots.size() > 0);
    REQUIRE(cpu_spots.size() == gpu_spots.size());
    CHECK(SortedCoords(cpu_spots) == SortedCoords(gpu_spots));
}

// Azimuthal-integration functionality: the profile the fused engine computes as a byproduct of the
// same pass must match a standalone GPU azimuthal integrator over the same image.
TEST_CASE("AdaptiveSpotFinderGPU_AzimuthalIntegration", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_AzimuthalIntegration");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);
    gpu.Run(buffer, settings);

    AzIntEngineGPU azint(mapping, stream);
    AzimuthalIntegrationProfile ref_profile(mapping);
    azint.Run(buffer, ref_profile);

    const auto ref = ref_profile.GetResult();
    const auto got = gpu.GetProfile().GetResult();
    const auto ref_count = ref_profile.GetPixelCount();
    const auto got_count = gpu.GetProfile().GetPixelCount();
    REQUIRE(ref.size() == got.size());
    REQUIRE(ref_count == got_count);   // identical per-ring pixel counts (same valid-pixel binning)
    for (size_t b = 0; b < ref.size(); b++) {
        if (std::isnan(ref[b])) {
            CHECK(std::isnan(got[b]));
        } else {
            CHECK(got[b] == Catch::Approx(ref[b]).epsilon(0.01).margin(0.02));
        }
    }
}

// The ring sums are built by atomics, which arrive in an arbitrary order, so the same frame has to be
// re-run to show the engine agrees with itself: detection is a hard "value >= threshold" on integer
// counts, and a threshold that wobbles between runs flips pixels on the boundary and with them the size
// of a connected component. Two runs, same spot list.
TEST_CASE("AdaptiveSpotFinderGPU_RunToRunReproducible", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_RunToRunReproducible");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);

    // The ring accumulators are exact integers, so the threshold does not depend on the order the
    // block atomics arrive in and the spot list has to be bit-identical every time - not merely
    // close. Repeat enough times to give a scheduling-dependent threshold a chance to show itself:
    // the effect it used to have was ~1 changed observation in a million, so a handful of repeats on
    // a quiet background would not have caught it.
    const auto first = gpu.Run(buffer, settings);
    REQUIRE(first.size() > 0);
    for (int repeat = 0; repeat < 50; repeat++) {
        const auto again = gpu.Run(buffer, settings);
        REQUIRE(again.size() == first.size());
        REQUIRE(SortedCoords(again) == SortedCoords(first));
    }
}

// The threshold is computed from sums of int32 pixel values, so the two engines can agree EXACTLY
// rather than approximately - and that is the property worth locking, because it is what makes the
// GPU path's spot list independent of how the reduction happened to be scheduled.
TEST_CASE("AdaptiveSpotFinderGPU_RingStatsMatchCPUExactly", "[AdaptiveSpotFinderGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_RingStatsMatchCPUExactly");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderGPU gpu(mapping, stream);
    AdaptiveSpotFinderCPU cpu(mapping);

    const auto gpu_spots = gpu.Run(buffer, settings);
    const auto cpu_spots = cpu.Run(buffer, settings);
    REQUIRE(gpu_spots.size() == cpu_spots.size());
    REQUIRE(SortedCoords(gpu_spots) == SortedCoords(cpu_spots));
}

TEST_CASE("AdaptiveSpotFinderGPU_Speed", "[AdaptiveSpotFinderGPU][.benchmark]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_Speed");
        return;
    }

    DiffractionExperiment x = MakeExperiment();
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    ImagePreprocessorBufferGPU buffer(x.GetPixelsNum());
    FillTestImage(buffer, x);
    REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
                       x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
    // The engines run on non-blocking streams, which do not wait for this NULL-stream copy: a pageable
    // H2D cudaMemcpy returns once the source is staged, with the DMA still in flight.
    REQUIRE(cudaDeviceSynchronize() == cudaSuccess);

    std::vector<bool> res_mask(x.GetPixelsNum(), false);
    const SpotFindingSettings settings = AdaptiveSettings();

    auto stream = std::make_shared<CudaStream>();
    AdaptiveSpotFinderCPU cpu(mapping);
    AdaptiveSpotFinderGPU gpu_fused(mapping, stream);
    ImageSpotFinderGPU gpu_classic(x.GetXPixelsNum(), x.GetYPixelsNum(), stream);
    AzIntEngineGPU azint(mapping, stream);
    AzimuthalIntegrationProfile profile(mapping);

    const int warmup = 5, iters = 40;
    auto bench = [&](const char *name, auto &&fn) {
        for (int i = 0; i < warmup; i++) fn();
        const auto t0 = std::chrono::steady_clock::now();
        for (int i = 0; i < iters; i++) fn();
        const auto t1 = std::chrono::steady_clock::now();
        const double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / iters;
        WARN(name << ": " << ms << " ms/frame");
        return ms;
    };

    const double t_azint = bench("GPU azint (standalone)", [&] { azint.Run(buffer, profile); });
    const double t_cpu = bench("CPU adaptive spot finding", [&] { cpu.Run(buffer, settings); });
    const double t_classic = bench("GPU classic spot finding (local-box)", [&] { gpu_classic.Run(buffer, settings); });
    const double t_fused = bench("GPU adaptive FUSED (azint + spot finding)", [&] { gpu_fused.Run(buffer, settings); });

    WARN("standard adaptive path (GPU azint + CPU adaptive) = " << (t_azint + t_cpu)
         << " ms/frame   vs   fused GPU = " << t_fused << " ms/frame   (speedup "
         << (t_azint + t_cpu) / t_fused << "x)");
    WARN("fused GPU vs GPU classic finder alone (no azint): " << t_fused << " vs " << t_classic << " ms/frame");
}

#endif