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

// The background under a reflection is estimated from a ring that is CONCENTRIC with the signal disk,
// which is the whole reason a background varying across the reflection cancels between them. Within r3
// of the edge of the array the ring loses its outer part and stops being concentric: what is left sits
// off-centre along the radius, where a radial background has a different value, and the ring mean
// carries that difference into every disk pixel it is subtracted from. These tests put a known radial
// ramp under reflections at a range of distances from the border and ask for the background at the
// reflection's own centre back.

namespace {

constexpr double BKG_LEVEL = 200.0;      // counts/px at the beam centre
constexpr double BKG_GRADIENT = 2.0;     // counts/px per pixel of radius

Reflection MakeReflection(float x, float y, int hkl) {
    Reflection r{};
    r.h = hkl; r.k = hkl; r.l = hkl;
    r.predicted_x = x;
    r.predicted_y = y;
    r.d = 2.0f;
    r.prescaling_corr = 1.0f;
    r.partiality = 1.0f;
    return r;
}

double TrueBackground(double x, double y, double beam_x, double beam_y) {
    return BKG_LEVEL + BKG_GRADIENT * std::hypot(x - beam_x, y - beam_y);
}

DiffractionExperiment MakeExperiment(IntegratorMode mode, float beam_x, float beam_y) {
    DiffractionExperiment experiment(DetJF(2));
    experiment.DetectorDistance_mm(100.0f).IncidentEnergy_keV(WVL_1A_IN_KEV)
              .BeamX_pxl(beam_x).BeamY_pxl(beam_y);
    BraggIntegrationSettings settings;
    settings.Integrator(mode);
    experiment.ImportBraggIntegrationSettings(settings);
    return experiment;
}

}   // namespace

// The ramp is radial and there is no signal anywhere, so the answer is known exactly: the background
// the engine reports has to be the ramp's value at the reflection's own centre, at the border as much
// as in the middle of the array. Averaging the surviving ring instead reports its centroid's value,
// which at the border is several pixels of radius away.
TEST_CASE("BraggBackground_TruncatedRingIsNotBiased", "[Integration]") {
    const float beam_x = 400.0f, beam_y = 400.0f;
    const DiffractionExperiment experiment = MakeExperiment(IntegratorMode::BoxSum, beam_x, beam_y);
    const int W = static_cast<int>(experiment.GetXPixelsNum());
    const int H = static_cast<int>(experiment.GetYPixelsNum());

    ImagePreprocessorBuffer image(experiment.GetPixelsNum());
    for (int y = 0; y < H; ++y)
        for (int x = 0; x < W; ++x)
            image[static_cast<size_t>(y) * W + x] =
                    static_cast<int32_t>(std::lround(TrueBackground(x, y, beam_x, beam_y)));

    // Reflections marching in towards the array from the bottom edge, at three azimuths so the border
    // cuts the ring at a different angle to the radius each time, plus an interior control.
    std::vector<Reflection> predicted;
    int hkl = 1;
    for (int x : {400, 700, 1000})
        for (int inset : {2, 4, 6, 9, 12, 16, 40})
            predicted.push_back(MakeReflection(static_cast<float>(x) + 0.3f,
                                               static_cast<float>(H - 1 - inset) - 0.2f, hkl++));

    BraggIntegrationEngineCPU engine(experiment);
    const auto out = engine.Run(image, predicted, predicted.size(), 1);
    REQUIRE(out.size() == predicted.size());

    for (const auto &r : out) {
        const double expected = TrueBackground(r.predicted_x, r.predicted_y, beam_x, beam_y);
        INFO("reflection at " << r.predicted_x << "," << r.predicted_y
             << " expected bkg " << expected << " got " << r.bkg);
        CHECK(r.bkg == Catch::Approx(expected).margin(0.5));
        // A flat background under the disk means the box sum has nothing above it to report.
        CHECK(std::abs(r.I) < 6.0f * r.sigma);
    }
}

// The same reflections read against a FLAT background: the fit must not invent a correction where
// there is no gradient to correct, at the border or anywhere else.
TEST_CASE("BraggBackground_FlatBackgroundIsUnchangedAtTheBorder", "[Integration]") {
    const float beam_x = 400.0f, beam_y = 400.0f;
    const DiffractionExperiment experiment = MakeExperiment(IntegratorMode::BoxSum, beam_x, beam_y);
    const int H = static_cast<int>(experiment.GetYPixelsNum());

    ImagePreprocessorBuffer image(experiment.GetPixelsNum());
    for (size_t i = 0; i < experiment.GetPixelsNum(); ++i)
        image[i] = static_cast<int32_t>(BKG_LEVEL);

    std::vector<Reflection> predicted;
    int hkl = 1;
    for (int inset : {2, 4, 6, 9, 12, 16, 40}) {
        predicted.push_back(MakeReflection(700.3f, static_cast<float>(H - 1 - inset) - 0.2f, hkl++));
        predicted.push_back(MakeReflection(static_cast<float>(inset) + 0.3f, 700.2f, hkl++));
    }

    BraggIntegrationEngineCPU engine(experiment);
    const auto out = engine.Run(image, predicted, predicted.size(), 1);
    REQUIRE(out.size() == predicted.size());
    for (const auto &r : out) {
        INFO("reflection at " << r.predicted_x << "," << r.predicted_y);
        CHECK(r.bkg == Catch::Approx(BKG_LEVEL).margin(1e-3));
    }
}

// MINPK is a fraction of the profile the fit can see, and the denominator has to be the WHOLE profile:
// counted over the grid cells that land on the array it is 1 by construction for a reflection whose
// peak is off the sensor, and then no threshold on it can fire.
TEST_CASE("BraggBackground_ProfileMassOffTheArrayIsRejected", "[Integration]") {
    const float beam_x = 400.0f, beam_y = 400.0f;
    const DiffractionExperiment experiment =
            MakeExperiment(IntegratorMode::ProfileGaussian, beam_x, beam_y);
    const int W = static_cast<int>(experiment.GetXPixelsNum());
    const int H = static_cast<int>(experiment.GetYPixelsNum());

    ImagePreprocessorBuffer image(experiment.GetPixelsNum());
    for (size_t i = 0; i < experiment.GetPixelsNum(); ++i)
        image[i] = static_cast<int32_t>(BKG_LEVEL);
    // Strong, well-formed spots so the profile is learned and the fit has something to work on.
    auto add_spot = [&](float cx, float cy) {
        for (int dy = -6; dy <= 6; ++dy)
            for (int dx = -6; dx <= 6; ++dx) {
                const int x = static_cast<int>(std::lround(cx)) + dx;
                const int y = static_cast<int>(std::lround(cy)) + dy;
                if (x < 0 || y < 0 || x >= W || y >= H) continue;
                const double ex = x - cx, ey = y - cy;
                image[static_cast<size_t>(y) * W + x] +=
                        static_cast<int32_t>(std::lround(4000.0 * std::exp(-(ex * ex + ey * ey) / (2.0 * 1.3 * 1.3))));
            }
    };

    std::vector<Reflection> predicted;
    int hkl = 1;
    for (int gy = 0; gy < 8; ++gy)
        for (int gx = 0; gx < 8; ++gx) {
            const float cx = 100.0f + 60.0f * gx, cy = 100.0f + 60.0f * gy;
            add_spot(cx, cy);
            predicted.push_back(MakeReflection(cx, cy, hkl++));
        }
    // A reflection whose predicted centre sits just outside the array: nearly all of its profile,
    // its peak included, is off the sensor.
    const size_t off_array = predicted.size();
    predicted.push_back(MakeReflection(700.0f, static_cast<float>(H) + 2.0f, hkl++));

    BraggIntegrationEngineCPU engine(experiment);
    const auto out = engine.Run(image, predicted, predicted.size(), 1);
    for (const auto &r : out)
        CHECK(r.h != static_cast<int>(off_array) + 1);
    CHECK(out.size() == off_array);
}