// 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 <cmath>
#include <cstdint>
#include <vector>

#include "../common/BraggIntegrationSettings.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/Reflection.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h"

// One spot on a flat background, inside a grid of predictions 8 px apart that put no flux on the
// frame - the tails of reflections recorded on the frames either side of a finely sliced one. Their
// r2 regions cover every background ring, so whether the spot keeps a clean ring depends only on
// whether those predictions are allowed to mask it.
TEST_CASE("BraggIntegrationEngineCPU_NeighbourMaskFollowsPartiality", "[Integration][portable]") {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(400.0f).BeamY_pxl(400.0f);
    experiment.ImportBraggIntegrationSettings(BraggIntegrationSettings());
    const size_t width = experiment.GetXPixelsNum(), npixel = experiment.GetPixelsNum();

    const float cx = 600.3f, cy = 300.2f, amp = 800.0f, sigma = 1.3f;
    ImagePreprocessorBuffer image(npixel);
    for (size_t i = 0; i < npixel; ++i)
        image[i] = 12;
    for (int dy = -6; dy <= 6; ++dy)
        for (int dx = -6; dx <= 6; ++dx) {
            const int x = static_cast<int>(std::lround(cx)) + dx, y = static_cast<int>(std::lround(cy)) + dy;
            const float ex = x - cx, ey = y - cy;
            image[y * width + x] += static_cast<int32_t>(std::lround(amp * std::exp(-(ex * ex + ey * ey) / (2 * sigma * sigma))));
        }

    auto run = [&](float neighbour_partiality, BraggIntegrationCounts &counts) {
        std::vector<Reflection> predicted;
        for (int gy = -4; gy <= 4; ++gy)
            for (int gx = -4; gx <= 4; ++gx) {
                Reflection r{};
                r.h = gx; r.k = gy; r.l = 1;
                r.predicted_x = cx + 8.0f * gx;
                r.predicted_y = cy + 8.0f * gy;
                r.d = 2.0f;
                r.prescaling_corr = 1.0f;
                r.partiality = (gx == 0 && gy == 0) ? 0.5f : neighbour_partiality;
                predicted.push_back(r);
            }
        BraggIntegrationEngineCPU engine(experiment);
        const auto out = engine.Run(image, predicted, predicted.size(), 0);
        counts = engine.Counts();
        std::vector<Reflection> spot;
        for (const auto &r : out)
            if (r.h == 0 && r.k == 0)
                spot.push_back(r);
        return spot;
    };

    // Tails: no reflection is dropped, and the spot is measured against the clean flat background -
    // while the density the widened-radius guard reads still sees every prediction.
    BraggIntegrationCounts tails;
    const auto spot = run(0.01f, tails);
    CHECK(tails.bkg_starved == 0);
    CHECK(tails.bkg_starved_by_neighbour > 0);
    REQUIRE(spot.size() == 1);
    CHECK(spot[0].bkg == Catch::Approx(12.0f));
    CHECK(spot[0].I == Catch::Approx(6.2831853 * sigma * sigma * amp).epsilon(0.05));

    // The same predictions with their flux on this frame do take the rings, the spot's among them, and
    // it is the neighbours, not the detector, that starve them - the same rings the density counted.
    BraggIntegrationCounts on_frame;
    CHECK(run(1.0f, on_frame).empty());
    CHECK(on_frame.bkg_starved > 0);
    CHECK(on_frame.bkg_starved_by_neighbour == on_frame.bkg_starved);
    CHECK(on_frame.bkg_starved_by_neighbour == tails.bkg_starved_by_neighbour);
}