// SPDX-FileCopyrightText: 2025 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 "../image_analysis/spot_finding/ImageSpotFinderGPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinderCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"

static void fill_test_image(ImagePreprocessorBuffer& buffer, size_t width, size_t height) {
    for (size_t i = 0; i < width * height; i++)
        buffer[i] = (i % 2) * 5 + 5;
    buffer[width * 50 + 50] = 20;
    buffer[width * 25 + 26] = 16;
    buffer[width * 75 + 25] = 12;
}

// Helper to run the GPU finder and collect its spot list
static std::vector<DiffractionSpot> run_gpu_and_collect_spots(ImagePreprocessorBufferGPU &buffer,
                                                              size_t width, size_t height,
                                                              const SpotFindingSettings &settings,
                                                              const std::vector<bool> &res_mask) {
    auto stream = std::make_shared<CudaStream>();
    ImageSpotFinderGPU gpu(static_cast<int32_t>(width), static_cast<int32_t>(height), stream);
    REQUIRE(get_gpu_count() > 0);
    gpu.SetResolutionMask(res_mask);

    REQUIRE(cudaMemcpyAsync(buffer.getGPUBuffer(),
        buffer.getBuffer().data(),
        width * height * sizeof(int32_t),
        cudaMemcpyHostToDevice,
        *stream) == cudaSuccess);

    return gpu.Run(buffer, settings);
}

// Mirror of ImageSpotFinder_SignalToNoise
TEST_CASE("ImageSpotFinderGPU_SignalToNoise") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping ImageSpotFinderGPU_SignalToNoise");
        return;
    }

    const size_t width = 100, height = 100;
    std::vector<bool> res_mask(width * height, false);
    std::vector<bool> mask(width * height, false);

    ImagePreprocessorBufferGPU buffer(width * height);
    fill_test_image(buffer, width, height);

    SpotFindingSettings settings{
        .signal_to_noise_threshold = 3.0,
        .photon_count_threshold = 0,
        .min_pix_per_spot = 1,
        .max_pix_per_spot = 20,
        .high_resolution_limit = 0.5,
        .low_resolution_limit = 3.0,
    };

    // GPU produces strong pixels; the resolution mask is handed to the finder separately and the
    // connected-component search then runs on the device (SpotExtractorGPU). The spot-level filter
    // matches CPU ImageSpotFinder test behavior for these synthetic inputs.
    auto spots = run_gpu_and_collect_spots(buffer, width, height, settings, res_mask);

    REQUIRE(spots.size() == 2);
    REQUIRE(spots[0].RawCoord().y == 25);
    REQUIRE(spots[1].RawCoord().y == 50);
}

TEST_CASE("ImageSpotFinderGPU_CountThreshold") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping ImageSpotFinderGPU_CountThreshold");
        return;
    }

    const size_t width = 100, height = 100;
    std::vector<bool> res_mask(width * height, false);
    std::vector<bool> mask(width * height, false);

    ImagePreprocessorBufferGPU buffer(width * height);
    fill_test_image(buffer, width, height);

    SpotFindingSettings settings{
        .signal_to_noise_threshold = 0.0,
        .photon_count_threshold = 11,
        .min_pix_per_spot = 1,
        .max_pix_per_spot = 20,
        .high_resolution_limit = 0.5,
        .low_resolution_limit = 3.0,
    };

    // GPU produces strong pixels; the resolution mask is handed to the finder separately and the
    // connected-component search then runs on the device (SpotExtractorGPU). The spot-level filter
    // matches CPU ImageSpotFinder test behavior for these synthetic inputs.
    auto spots = run_gpu_and_collect_spots(buffer, width, height, settings, res_mask);

    REQUIRE(spots.size() == 3);
    REQUIRE(spots[0].RawCoord().y == 25);
    REQUIRE(spots[1].RawCoord().y == 50);
    REQUIRE(spots[2].RawCoord().y == 75);
}

TEST_CASE("ImageSpotFinderGPU_20M") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping ImageSpotFinderGPU_20M");
        return;
    }

    const size_t width = 4500, height = 4500;
    std::vector<bool> res_mask(width * height, false);
    std::vector<bool> mask(width * height, false);

    ImagePreprocessorBufferGPU buffer(width * height);
    fill_test_image(buffer, width, height);

    SpotFindingSettings settings{
        .signal_to_noise_threshold = 3.0,
        .photon_count_threshold = 0,
        .min_pix_per_spot = 1,
        .max_pix_per_spot = 20,
        .high_resolution_limit = 0.5,
        .low_resolution_limit = 3.0,
    };

    // GPU produces strong pixels; the resolution mask is handed to the finder separately and the
    // connected-component search then runs on the device (SpotExtractorGPU). The spot-level filter
    // matches CPU ImageSpotFinder test behavior for these synthetic inputs.
    auto spots = run_gpu_and_collect_spots(buffer, width, height, settings, res_mask);

    REQUIRE(spots.size() == 2);
    REQUIRE(spots[0].RawCoord().y == 25);
    REQUIRE(spots[1].RawCoord().y == 50);
}

// The two finders must return the same spots for the same frame - a dataset processed on a machine
// without a GPU has to give the same answer as one processed with it.
//
// The spots here are deliberately broad. Both finders measure a pixel against a 31x31 local
// background, so a blob several pixels across sits inside its own background window and inflates
// the mean and variance it is tested against. That is what the second pass exists to undo: it
// recomputes the background with the pixels found strong by the first pass excluded. A single-pass
// finder loses the outer pixels of every broad spot, so this comparison fails unless both sides run
// the same two passes.
TEST_CASE("ImageSpotFinder_CPU_GPU_Parity", "[ImageSpotFinder]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    const size_t width = 100, height = 100;

    ImagePreprocessorBufferGPU gpu_buffer(width * height);
    ImagePreprocessorBuffer cpu_buffer(width * height);

    // Background alternating 5/10 (mean 7.5, sd 2.5), plus two spots shaped to make the second pass
    // matter: a bright 5x5 core (300) with a thin one-pixel ring around it (25). The ring is well
    // above the clean background, but the core sitting inside the ring's own 31x31 window drags that
    // window's mean to ~16 and its sd to ~46, so on a single pass the ring fails the SNR test and the
    // spot comes out as the 25-pixel core. The second pass takes the core out of the background and
    // the ring passes, giving 49 pixels. The ring is kept thin on purpose: a wide halo would swamp
    // its own background and stay undetectable either way.
    auto fill = [&](ImagePreprocessorBuffer &b) {
        for (size_t i = 0; i < width * height; i++)
            b[i] = (i % 2) * 5 + 5;

        const struct { int cx, cy; } spots[] = {{50, 50}, {22, 74}};
        for (const auto &s : spots) {
            for (int dy = -3; dy <= 3; dy++) {
                for (int dx = -3; dx <= 3; dx++) {
                    const bool ring = std::abs(dx) == 3 || std::abs(dy) == 3;
                    b[(s.cy + dy) * width + (s.cx + dx)] = ring ? 25 : 300;
                }
            }
        }
    };
    fill(gpu_buffer);
    fill(cpu_buffer);

    SpotFindingSettings settings{
        .signal_to_noise_threshold = 3.0,
        .photon_count_threshold = 0,
        .min_pix_per_spot = 1,
        .max_pix_per_spot = 1000,
        .high_resolution_limit = 0.5,
        .low_resolution_limit = 3.0,
    };
    const std::vector<bool> res_mask(width * height, false);

    ImageSpotFinderCPU cpu(static_cast<int32_t>(width), static_cast<int32_t>(height));
    cpu.SetResolutionMask(res_mask);
    const auto cpu_spots = cpu.Run(cpu_buffer, settings);

    const auto gpu_spots = run_gpu_and_collect_spots(gpu_buffer, width, height, settings, res_mask);

    REQUIRE(cpu_spots.size() == gpu_spots.size());
    REQUIRE(cpu_spots.size() == 2);   // guard against both finding nothing and "agreeing"

    for (size_t i = 0; i < cpu_spots.size(); i++) {
        CHECK(cpu_spots[i].RawCoord().x == Catch::Approx(gpu_spots[i].RawCoord().x));
        CHECK(cpu_spots[i].RawCoord().y == Catch::Approx(gpu_spots[i].RawCoord().y));
        // The pixel set is what the second pass changes, so compare it rather than the centroid,
        // which stays put for a symmetric spot whether or not the halo was picked up.
        CHECK(cpu_spots[i].PixelCount() == gpu_spots[i].PixelCount());
        CHECK(cpu_spots[i].Count() == gpu_spots[i].Count());
    }

    // Both finders must reach past the bright core into the ring: 49 pixels, not the 25 of the core
    // alone. Without this the comparison above would still pass if both ran a single pass.
    CHECK(cpu_spots[0].PixelCount() == 49);
}

// The same comparison on a TALL image, which is what makes it bite.
//
// The GPU kernel splits the image into ImageSpotFinderGPU::numberOfWaves (32) horizontal waves and
// each wave walks its rows in three stages: a priming loop, a main loop, and a drain loop. The main
// loop only runs while front < rmax, and front starts NBX+1 rows ahead of the wave's first row - so
// for an image of 100 rows, where a wave owns ceil(100/32) = 4 rows, THE MAIN LOOP NEVER EXECUTES.
// Every row goes through priming/drain. On a real detector frame (4362 rows -> 137 rows per wave)
// the main loop instead carries ~121 of every 137 rows, i.e. almost the whole image.
//
// That is why a bug living only in the main loop survived: it was unreachable at the test's size.
// This case is 1024 rows (32 per wave, ~16 through the main loop) so the path is covered, and the
// spots are placed deep inside a wave rather than at its edges. Values above 46340 are included
// deliberately: the window keeps a sum of squares, and squaring in 32 bits overflows past that,
// while a real detector saturates around 1e6.
TEST_CASE("ImageSpotFinder_CPU_GPU_Parity_TallImage", "[ImageSpotFinder]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    const size_t width = 128, height = 1024;

    ImagePreprocessorBufferGPU gpu_buffer(width * height);
    ImagePreprocessorBuffer cpu_buffer(width * height);

    // Spots at rows a wave reaches through its MAIN loop, not its priming or drain rows.
    const struct { int cx, cy; int32_t core; } spots[] = {
        {64,  80, 300},        // ordinary bright spot
        {64, 400, 100000},     // above 46340: 32-bit squaring wraps, 64-bit does not
        {64, 720, 300},
    };
    auto fill = [&](ImagePreprocessorBuffer &b) {
        for (size_t i = 0; i < width * height; i++)
            b[i] = (i % 2) * 5 + 5;
        for (const auto &s : spots)
            for (int dy = -3; dy <= 3; dy++)
                for (int dx = -3; dx <= 3; dx++) {
                    const bool ring = std::abs(dx) == 3 || std::abs(dy) == 3;
                    b[(s.cy + dy) * width + (s.cx + dx)] = ring ? 25 : s.core;
                }
    };
    fill(gpu_buffer);
    fill(cpu_buffer);

    SpotFindingSettings settings{
        .signal_to_noise_threshold = 3.0,
        .photon_count_threshold = 0,
        .min_pix_per_spot = 1,
        .max_pix_per_spot = 1000,
        .high_resolution_limit = 0.5,
        .low_resolution_limit = 3.0,
    };
    const std::vector<bool> res_mask(width * height, false);

    ImageSpotFinderCPU cpu(static_cast<int32_t>(width), static_cast<int32_t>(height));
    cpu.SetResolutionMask(res_mask);
    const auto cpu_spots = cpu.Run(cpu_buffer, settings);
    const auto gpu_spots = run_gpu_and_collect_spots(gpu_buffer, width, height, settings, res_mask);

    REQUIRE(cpu_spots.size() == 3);   // all three found, so "agreeing on nothing" cannot pass
    REQUIRE(gpu_spots.size() == cpu_spots.size());

    for (size_t i = 0; i < cpu_spots.size(); i++) {
        INFO("spot " << i);
        CHECK(cpu_spots[i].RawCoord().x == Catch::Approx(gpu_spots[i].RawCoord().x));
        CHECK(cpu_spots[i].RawCoord().y == Catch::Approx(gpu_spots[i].RawCoord().y));
        CHECK(cpu_spots[i].PixelCount() == gpu_spots[i].PixelCount());
        CHECK(cpu_spots[i].Count() == gpu_spots[i].Count());
        CHECK(cpu_spots[i].PixelCount() == 49);   // core + ring, i.e. the second pass did its job
    }
}


// The local-box SNR test on the knife edge, where the GPU used to answer differently from the CPU.
// Both compare (in_minus_mean)^2 against var * strong2 computed in float; the CPU compares the
// integer square with the ceiling of that float, the GPU compared the square converted to float -
// which, for squares past 2^24, can round onto the threshold and fail a pixel the CPU accepts.
//
// The centre of a checkerboard of 0 and 57 holding 171 is such a pixel at 5 sigma: its 960
// neighbours are 480 zeros and 480 of 57, so var = 748569600 and in_minus_mean = 136800, whose
// square 18714240000 is var * 25 exactly. var * 25 rounds down in float to 18714238976, 1024 below
// the square - half a float spacing - so the square rounds onto the threshold too. The CPU takes
// the pixel (18714240000 > 18714238976), the GPU used not to.
//
// The second image is the count threshold: at photon_count_threshold = 0 the CPU still asks for a
// positive value, the GPU took 0 to mean no threshold. A 0 above a negative background (the
// checkerboard of -3 and -1, as a pedestal-subtracted detector gives) is 2 sigma above its mean.
TEST_CASE("ImageSpotFinder_CPU_GPU_Parity_KnifeEdge", "[ImageSpotFinder]") {
    if (get_gpu_count() == 0)
        SKIP("No CUDA GPU present");

    const size_t width = 64, height = 64;
    const std::vector<bool> res_mask(width * height, false);

    auto run_both = [&](int32_t even, int32_t odd, int32_t centre, float snr, size_t expected) {
        ImagePreprocessorBufferGPU gpu_buffer(width * height);
        ImagePreprocessorBuffer cpu_buffer(width * height);
        for (size_t row = 0; row < height; row++)
            for (size_t col = 0; col < width; col++) {
                const int32_t v = ((row + col) % 2 == 0) ? even : odd;
                gpu_buffer[row * width + col] = v;
                cpu_buffer[row * width + col] = v;
            }
        gpu_buffer[32 * width + 32] = centre;
        cpu_buffer[32 * width + 32] = centre;

        SpotFindingSettings settings{
            .signal_to_noise_threshold = snr,
            .photon_count_threshold = 0,
            .min_pix_per_spot = 1,
            .max_pix_per_spot = 1000,
            .high_resolution_limit = 0.5,
            .low_resolution_limit = 3.0,
        };
        ImageSpotFinderCPU cpu(static_cast<int32_t>(width), static_cast<int32_t>(height));
        cpu.SetResolutionMask(res_mask);
        const auto cpu_spots = cpu.Run(cpu_buffer, settings);
        const auto gpu_spots = run_gpu_and_collect_spots(gpu_buffer, width, height, settings, res_mask);
        CHECK(cpu_spots.size() == expected);
        CHECK(gpu_spots.size() == cpu_spots.size());
    };

    SECTION("SNR test with the square past 2^24") {
        run_both(0, 57, 171, 5.0f, 1);
    }
    SECTION("Zero count threshold on a negative background") {
        run_both(-3, -1, 0, 1.5f, 0);
    }
}

#endif