// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

// The GPU spot extractor (SpotExtractorGPU) has to reproduce the host connected-component search
// (StrongPixelSet::sparseccl, driven through ImageSpotFinder::ExtractComponentsHost) EXACTLY - the
// same partition of the strong pixels into components, the same number of spots, the same spot
// ORDER, and the same pixel count, photon sum, maximum and centroid sums for each of them. If the
// two ever drift apart, every downstream result (indexing, integration, merging) moves with them
// and nothing else in the suite would notice.
//
// Both extractors are driven from the SAME strong-pixel bit buffer here, so this compares only the
// extraction; the detection stages have their own parity cases
// (ImageSpotFinder_CPU_GPU_Parity, AdaptiveSpotFinderGPU_SpotFindingParity).
//
// The frames are generated, not recorded: dataset-derived fixtures are not allowed in the tree. The
// blob-size and density distributions were chosen to match what real 18 MP frames produce
// (~1e3 - 1.5e4 strong pixels, spots of a few to a few tens of pixels).

#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>

#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinder.h"
#include "../image_analysis/spot_finding/SpotExtractorGPU.h"

namespace {

// A detector large enough for the geometries below to be representative, small enough that a case
// costs milliseconds. The pathological shapes are what set the size, not the spot count.
constexpr int32_t WIDTH = 2048;
constexpr int32_t HEIGHT = 2048;

SpotFindingSettings ExtractionSettings(int64_t max_pix = 50) {
    SpotFindingSettings s{};
    s.min_pix_per_spot = 1;
    s.max_pix_per_spot = max_pix;
    return s;
}

// A finder whose detection stage is "use the bit buffer I was handed". It exists so both extractors
// can be given byte-identical input.
class PlantedSpotFinder : public ImageSpotFinder {
public:
    PlantedSpotFinder() : ImageSpotFinder(WIDTH, HEIGHT) {}
    void Detect(const ImagePreprocessorBuffer &, const SpotFindingSettings &) override {}
    void Plant(const std::vector<uint32_t> &strong_index) {
        std::fill(output_buffer.begin(), output_buffer.end(), 0);
        for (uint32_t flat : strong_index)
            output_buffer[flat / 32] |= 1u << (flat % 32);
    }
    const std::vector<uint32_t> &BitBuffer() const { return output_buffer; }
    size_t Words() const { return OutputSize(); }
};

// ---------------------------------------------------------------------------------------------
// Frame geometries
// ---------------------------------------------------------------------------------------------
struct Frame {
    std::string name;
    std::vector<uint32_t> index;       // flat pixel index, ascending
    std::vector<int32_t> value;
};

Frame Finish(std::string name, std::vector<uint32_t> index, uint64_t seed) {
    std::sort(index.begin(), index.end());
    index.erase(std::unique(index.begin(), index.end()), index.end());
    Frame f{std::move(name), std::move(index), {}};
    // Values spanning the range a preprocessed image really carries, including negatives (a
    // pedestal-subtracted pixel can go below zero) - the centroid sums must match on those too.
    std::mt19937_64 rng(seed);
    f.value.reserve(f.index.size());
    for (size_t i = 0; i < f.index.size(); i++)
        f.value.push_back(static_cast<int32_t>(rng() % 40000) - 200);
    return f;
}

// Compact blobs scattered over the detector: what a diffraction pattern looks like to the extractor.
Frame BlobFrame(const std::string &name, int nblob, uint64_t seed) {
    std::mt19937_64 rng(seed);
    std::vector<uint32_t> index;
    for (int b = 0; b < nblob; b++) {
        const int col = 20 + static_cast<int>(rng() % (WIDTH - 40));
        const int line = 20 + static_cast<int>(rng() % (HEIGHT - 40));
        const int radius = 1 + static_cast<int>(rng() % 3);     // spots of ~4 to ~30 pixels
        for (int dl = -radius; dl <= radius; dl++)
            for (int dc = -radius; dc <= radius; dc++)
                if (dl * dl + dc * dc <= radius * radius)
                    index.push_back(static_cast<uint32_t>((line + dl) * WIDTH + (col + dc)));
    }
    return Finish(name, std::move(index), seed ^ 0x9e37u);
}

// A lit band of whole detector rows - a hot module, a panel edge, and the limit a diffraction ring
// approaches where it runs tangent to a row. This is the shape the host's two-row sliding window
// turns quadratic on.
Frame BandFrame(const std::string &name, int nrows) {
    std::vector<uint32_t> index;
    for (int line = 500; line < 500 + nrows; line++)
        for (int col = 0; col < WIDTH; col++)
            index.push_back(static_cast<uint32_t>(line * WIDTH + col));
    return Finish(name, std::move(index), 0x51ed3u + nrows);
}

// Concentric diffraction rings. Unlike a hot column, a ring is SIGNAL and cannot be masked away, and
// at its top and bottom it produces exactly the long horizontal run that the band case models.
Frame RingFrame(const std::string &name, int radius, int half_width, int nring) {
    const double cx = WIDTH * 0.5, cy = HEIGHT * 0.5;
    std::vector<uint32_t> index;
    for (int line = 0; line < HEIGHT; line++) {
        for (int col = 0; col < WIDTH; col++) {
            const double d = std::sqrt((col - cx) * (col - cx) + (line - cy) * (line - cy));
            for (int k = 0; k < nring; k++) {
                if (std::fabs(d - (radius + k * 130)) <= half_width) {
                    index.push_back(static_cast<uint32_t>(line * WIDTH + col));
                    break;
                }
            }
        }
    }
    return Finish(name, std::move(index), 0x21f9u + radius * 31 + nring);
}

Frame SolidBlockFrame(const std::string &name, int side) {
    std::vector<uint32_t> index;
    for (int line = 300; line < 300 + side; line++)
        for (int col = 300; col < 300 + side; col++)
            index.push_back(static_cast<uint32_t>(line * WIDTH + col));
    return Finish(name, std::move(index), 0x7c1u + side);
}

// Isolated single pixels at the density a very noisy frame reaches: many components, all tiny.
Frame SparseNoiseFrame(const std::string &name, int npixel, uint64_t seed) {
    std::mt19937_64 rng(seed);
    std::vector<uint32_t> index;
    for (int i = 0; i < npixel; i++)
        index.push_back(static_cast<uint32_t>(rng() % (static_cast<uint64_t>(WIDTH) * HEIGHT)));
    return Finish(name, std::move(index), seed ^ 0x1234u);
}

// ---------------------------------------------------------------------------------------------
// Harness: run both extractors on one frame
// ---------------------------------------------------------------------------------------------
class ParityRunner {
    PlantedSpotFinder finder;
    ImagePreprocessorBufferGPU image;
    std::shared_ptr<CudaStream> stream;
    SpotExtractorGPU extractor;
    CudaDevicePtr<uint32_t> gpu_strong;

public:
    ParityRunner()
        : image(static_cast<size_t>(WIDTH) * HEIGHT),
          stream(std::make_shared<CudaStream>()),
          extractor(WIDTH, HEIGHT, stream),
          gpu_strong(finder.Words()) {}

    void SetResolutionMask(const std::vector<bool> &mask) {
        finder.SetResolutionMask(mask);
        std::vector<uint32_t> packed((static_cast<size_t>(WIDTH) * HEIGHT + 31) / 32, 0);
        for (size_t i = 0; i < mask.size(); i++)
            if (mask[i]) packed[i / 32] |= 1u << (i % 32);
        extractor.SetResolutionMask(packed);
    }

    void Load(const Frame &f) {
        auto &buffer = image.getBuffer();
        std::fill(buffer.begin(), buffer.end(), 0);
        for (size_t i = 0; i < f.index.size(); i++)
            buffer[f.index[i]] = f.value[i];
        REQUIRE(cudaMemcpy(image.getGPUBuffer(), buffer.data(), buffer.size() * sizeof(int32_t),
                           cudaMemcpyHostToDevice) == cudaSuccess);
        finder.Plant(f.index);
        REQUIRE(cudaMemcpy(gpu_strong.get(), finder.BitBuffer().data(),
                           finder.Words() * sizeof(uint32_t), cudaMemcpyHostToDevice) == cudaSuccess);
        REQUIRE(cudaDeviceSynchronize() == cudaSuccess);
    }

    std::vector<DiffractionSpot> Cpu(const SpotFindingSettings &settings) {
        return finder.ExtractComponents(image, settings);
    }
    std::vector<DiffractionSpot> Gpu(const SpotFindingSettings &settings) {
        std::vector<DiffractionSpot> spots;
        extractor.Extract(gpu_strong.get(), image.getGPUBuffer(), settings, spots);
        return spots;
    }
};

// Everything is compared EXACTLY: the number of spots, their order, and each one's pixel count,
// photon sum, maximum and centroid. All of them are built out of integer sums on both sides, so
// there is nothing here that a build flag can move - which is the point of accumulating in integers
// rather than in float, where gcc contracted the multiply-add under -march=x86-64-v3 and not at the
// baseline and left the last bit of the centroid a property of how the host was built.
void RequireIdentical(const std::string &what,
                      const std::vector<DiffractionSpot> &cpu,
                      const std::vector<DiffractionSpot> &gpu) {
    INFO(what);
    REQUIRE(gpu.size() == cpu.size());
    for (size_t i = 0; i < cpu.size(); i++) {
        INFO("spot " << i << " of " << cpu.size());
        REQUIRE(gpu[i].PixelCount() == cpu[i].PixelCount());
        REQUIRE(gpu[i].Count() == cpu[i].Count());
        REQUIRE(gpu[i].MaxCount() == cpu[i].MaxCount());
        // RawCoord divides the sums by the photon count, so comparing it compares the sums; a spot
        // whose photons sum to zero reports (0,0) on both sides by the same branch.
        const Coord c = cpu[i].RawCoord(), g = gpu[i].RawCoord();
        REQUIRE(memcmp(&g.x, &c.x, sizeof(float)) == 0);
        REQUIRE(memcmp(&g.y, &c.y, sizeof(float)) == 0);
    }
}

std::vector<Frame> RealisticFrames() {
    std::vector<Frame> frames;
    // Spot counts a real 18 MP frame produces, scaled to this detector: a weak still, a good
    // rotation frame, a heavily-diffracting one.
    frames.push_back(BlobFrame("blobs_sparse", 60, 11));
    frames.push_back(BlobFrame("blobs_typical", 600, 22));
    frames.push_back(BlobFrame("blobs_dense", 2500, 33));
    return frames;
}

std::vector<Frame> PathologicalFrames() {
    std::vector<Frame> frames;
    frames.push_back(BandFrame("band_2_rows", 2));
    frames.push_back(BandFrame("band_8_rows", 8));
    frames.push_back(BandFrame("band_31_rows", 31));       // just under the 65534-strong-pixel cap
    frames.push_back(SolidBlockFrame("solid_block_120", 120));
    frames.push_back(SolidBlockFrame("solid_block_240", 240));
    frames.push_back(RingFrame("ring_r700_w3", 700, 1, 1));
    frames.push_back(RingFrame("ring_r700_w9", 700, 4, 1));
    frames.push_back(RingFrame("ring_r300_w3_x4", 300, 1, 4));
    frames.push_back(RingFrame("ring_r950_w5", 950, 2, 1));
    frames.push_back(SparseNoiseFrame("noise_20k", 20000, 44));
    return frames;
}

} // namespace

TEST_CASE("SpotExtractorGPU_Parity_RealisticFrames", "[SpotExtractorGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_RealisticFrames");
        return;
    }
    ParityRunner runner;
    const auto settings = ExtractionSettings();
    for (const auto &f : RealisticFrames()) {
        runner.Load(f);
        RequireIdentical(f.name + " (" + std::to_string(f.index.size()) + " strong pixels)",
                         runner.Cpu(settings), runner.Gpu(settings));
    }
}

// The occupancy range the extractor has to hold up over, from a nearly empty frame to one pixel
// short of the give-up threshold, plus the first frame past it (where both sides must return
// nothing rather than disagree).
TEST_CASE("SpotExtractorGPU_Parity_OccupancySweep", "[SpotExtractorGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_OccupancySweep");
        return;
    }
    ParityRunner runner;
    const auto settings = ExtractionSettings();
    for (int target : {100, 300, 1000, 3000, 10000, 30000, 65000, 70000}) {
        // Isolated pixels give the most components per strong pixel, so this is also the worst case
        // for the labelling and the compaction.
        Frame f = SparseNoiseFrame("occupancy_" + std::to_string(target), target, 900u + target);
        runner.Load(f);
        RequireIdentical(f.name + " (" + std::to_string(f.index.size()) + " strong pixels)",
                         runner.Cpu(settings), runner.Gpu(settings));
    }
}

TEST_CASE("SpotExtractorGPU_Parity_Pathological", "[SpotExtractorGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_Pathological");
        return;
    }
    ParityRunner runner;
    for (const auto &f : PathologicalFrames()) {
        runner.Load(f);
        // Both the production max-pix and an effectively unlimited one: the second forces the big
        // components to be summed as well, which the first deliberately skips.
        for (int64_t max_pix : {int64_t(50), int64_t(1000000)}) {
            const auto settings = ExtractionSettings(max_pix);
            RequireIdentical(f.name + " (" + std::to_string(f.index.size()) + " strong pixels, max_pix "
                                     + std::to_string(max_pix) + ")",
                             runner.Cpu(settings), runner.Gpu(settings));
        }
    }
}

TEST_CASE("SpotExtractorGPU_Parity_ResolutionMask", "[SpotExtractorGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Parity_ResolutionMask");
        return;
    }
    ParityRunner runner;
    const auto settings = ExtractionSettings();
    // A resolution band, the shape UpdateMaskResolution really produces: excluding it has to split
    // and drop components identically on both sides.
    std::vector<bool> mask(static_cast<size_t>(WIDTH) * HEIGHT, false);
    const double cx = WIDTH * 0.5, cy = HEIGHT * 0.5;
    for (int line = 0; line < HEIGHT; line++)
        for (int col = 0; col < WIDTH; col++) {
            const double d = std::sqrt((col - cx) * (col - cx) + (line - cy) * (line - cy));
            if (d < 250.0 || d > 900.0)
                mask[static_cast<size_t>(line) * WIDTH + col] = true;
        }
    runner.SetResolutionMask(mask);
    for (const auto &f : RealisticFrames()) {
        runner.Load(f);
        RequireIdentical(f.name + " with resolution mask", runner.Cpu(settings), runner.Gpu(settings));
    }
    runner.Load(RingFrame("ring_r700_w3_masked", 700, 1, 1));
    RequireIdentical("ring with resolution mask", runner.Cpu(settings), runner.Gpu(settings));
}

// The extraction is full of atomics and a lock-free union-find, so the same frame has to come back
// byte-identical however the blocks happen to interleave. Anything less and the spot list - and
// every result downstream of it - would wander between runs of the same data.
TEST_CASE("SpotExtractorGPU_Determinism", "[SpotExtractorGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping SpotExtractorGPU_Determinism");
        return;
    }
    ParityRunner runner;
    const auto settings = ExtractionSettings();
    for (const auto &f : {BlobFrame("blobs_typical", 600, 22), RingFrame("ring_r700_w3", 700, 1, 1),
                          BandFrame("band_8_rows", 8)}) {
        runner.Load(f);
        const auto reference = runner.Gpu(settings);
        for (int repeat = 0; repeat < 100; repeat++) {
            const auto again = runner.Gpu(settings);
            INFO(f.name << ", repeat " << repeat);
            REQUIRE(again.size() == reference.size());
            REQUIRE(memcmp(again.data(), reference.data(), again.size() * sizeof(DiffractionSpot)) == 0);
        }
    }
}

#endif