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

static void fill_test_image(std::vector<int32_t>& input, size_t width, size_t height) {
    input.resize(width * height);
    for (size_t i = 0; i < width * height; i++)
        input[i] = (i % 2) * 5 + 5;
    input[width * 50 + 50] = 20;
    input[width * 25 + 26] = 16;
    input[width * 75 + 25] = 12;
}

// Helper to run GPU and get DiffractionSpot list via StrongPixelSet -> FindSpotsImage
static std::vector<DiffractionSpot> run_gpu_and_collect_spots(const std::vector<int32_t>& input,
                                                              size_t width, size_t height,
                                                              const SpotFindingSettings& settings,
                                                              const std::vector<bool>& res_mask)
{
    ImageSpotFinderGPU gpu(static_cast<int32_t>(width), static_cast<int32_t>(height));
    REQUIRE(get_gpu_count() > 0);

    memcpy(gpu.GetInputBuffer().data(), input.data(), width * height * sizeof(int32_t));
    return gpu.Run(settings, res_mask);
}

// 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);

    std::vector<int32_t> input;
    fill_test_image(input, 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; FindSpotsImage uses mask/resolution implicit in StrongPixelSet.
    // StrongPixelSet doesn't carry resolution/mask by itself, but FindSpotsImage(settings, vec)
    // matches CPU ImageSpotFinder test behavior for these synthetic inputs.
    auto spots = run_gpu_and_collect_spots(input, 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);

    std::vector<int32_t> input;
    fill_test_image(input, 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; FindSpotsImage uses mask/resolution implicit in StrongPixelSet.
    // StrongPixelSet doesn't carry resolution/mask by itself, but FindSpotsImage(settings, vec)
    // matches CPU ImageSpotFinder test behavior for these synthetic inputs.
    auto spots = run_gpu_and_collect_spots(input, 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);


    std::vector<int32_t> input;
    fill_test_image(input, 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; FindSpotsImage uses mask/resolution implicit in StrongPixelSet.
    // StrongPixelSet doesn't carry resolution/mask by itself, but FindSpotsImage(settings, vec)
    // matches CPU ImageSpotFinder test behavior for these synthetic inputs.
    auto spots = run_gpu_and_collect_spots(input, width, height, settings, res_mask);

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

#endif