// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_all.hpp>

#ifdef JFJOCH_USE_CUDA

#include "../image_analysis/spot_finding//ImageAnalysisGPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinder.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)
{
    ImageSpotFinderGPU gpu(static_cast<int32_t>(width), static_cast<int32_t>(height));
    REQUIRE(ImageSpotFinderGPU::GPUPresent());

    // Set input buffer pointer to our CPU data, register and upload
    // Note: RegisterBuffer/UnregisterBuffer are optional for plain memcpy path,
    // but we use them to mirror intended flow.
    const_cast<ImageSpotFinderGPU&>(gpu).SetInputBuffer((void*)input.data());
    gpu.RegisterBuffer();
    gpu.LoadDataToGPU();

    // Run kernel
    gpu.RunSpotFinder(settings);

    // Collect strong pixels and convert to spots like CPU does
    StrongPixelSet strong;
    gpu.GetSpotFinderResults(strong);

    std::vector<DiffractionSpot> spots;
    strong.FindSpotsImage(settings, spots);

    gpu.UnregisterBuffer();
    return spots;
}

// Mirror of ImageSpotFinder_SignalToNoise
TEST_CASE("GPUImageAnalysis_SignalToNoise") {
    if (!ImageSpotFinderGPU::GPUPresent()) {
        WARN("No CUDA GPU present. Skipping GPUImageAnalysis_SignalToNoise");
        return;
    }

    const size_t width = 100, height = 100;
    std::vector<float> resolution(width * height, 2.0f);
    std::vector<bool> mask(width * height, false);
    resolution[width * 50 + 50] = 1.0f;

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

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

TEST_CASE("GPUImageAnalysis_CountThreshold") {
    if (!ImageSpotFinderGPU::GPUPresent()) {
        WARN("No CUDA GPU present. Skipping GPUImageAnalysis_SignalToNoise");
        return;
    }

    const size_t width = 100, height = 100;
    std::vector<float> resolution(width * height, 2.0f);
    std::vector<bool> mask(width * height, false);
    resolution[width * 50 + 50] = 1.0f;

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

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

#endif