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

#include <catch2/catch_test_macros.hpp>

#include <cmath>
#include <random>

#include "../image_analysis/bragg_integration/BackgroundAperture.h"
#include "../common/JFJochMath.h"

namespace {

constexpr int W = 600, H = 600;

// A background of `level` counts a pixel with Poisson noise, modulated by `ripple` (relative) with a
// period of 20 px along x - the structure a ring further from the disk reads wrong.
ImagePreprocessorBuffer Frame(double level, double ripple, uint32_t seed) {
    std::mt19937 rng(seed);
    ImagePreprocessorBuffer img(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            std::poisson_distribution<int> p(level * (1.0 + ripple * std::sin(2.0 * PI * x / 20.0)));
            img[static_cast<size_t>(y) * W + x] = p(rng);
        }
    return img;
}

const RingBackground no_rings({}, 0, 0.0f, 0.0f);

DiffractionGeometry Geometry() {
    DiffractionGeometry g;
    g.BeamX_pxl(300.0f).BeamY_pxl(300.0f).DetectorDistance_mm(60.0f).PixelSize_mm(0.075f).Wavelength_A(1.0f);
    return g;
}

int BestR3(double level, double ripple) {
    const auto r3 = background_aperture::CandidateR3(8.0f, 14.0f);
    background_aperture::Tally tally;
    for (uint32_t seed = 1; seed <= 4; seed++) {
        auto img = Frame(level, ripple, seed);
        background_aperture::MeasureFrame(img, W, H, {}, 6.0f, 8.0f, r3, 4.0f, 1.0f, no_rings, Geometry(),
                                          DetectorModules(), tally);
    }
    REQUIRE(tally.apertures > 100);
    CHECK(tally.apertures_radial == 0);
    return static_cast<int>(r3[background_aperture::Best(tally.MeanSquare(false))]);
}

}   // namespace

TEST_CASE("BackgroundAperture_CandidateR3", "[Integration]") {
    const auto r3 = background_aperture::CandidateR3(8.0f, 14.0f);
    REQUIRE(r3.size() == 5);
    CHECK(r3.front() == 10.0f);
    CHECK(r3.back() == 14.0f);
}

// A few counts a pixel on a flat background: the noise wins, and the widest ring is kept.
TEST_CASE("BackgroundAperture_FlatNoisyKeepsWidestRing", "[Integration]") {
    CHECK(BestR3(0.5, 0.0) == 14);
}

// Thousands of counts a pixel on a background that varies over the stencil's own scale: the structure
// wins, and the narrowest ring is kept.
TEST_CASE("BackgroundAperture_StructuredBrightTakesNarrowestRing", "[Integration]") {
    CHECK(BestR3(2000.0, 0.05) == 10);
}

// An aperture over an unreadable pixel is not used.
TEST_CASE("BackgroundAperture_SkipsUnreadablePixels", "[Integration]") {
    const auto r3 = background_aperture::CandidateR3(8.0f, 14.0f);
    auto img = Frame(1.0, 0.0, 7);
    background_aperture::Tally all;
    background_aperture::MeasureFrame(img, W, H, {}, 6.0f, 8.0f, r3, 4.0f, 1.0f, no_rings, Geometry(),
                                      DetectorModules(), all);
    for (size_t i = 0; i < static_cast<size_t>(W) * H; i += 97)
        img[i] = INT32_MIN;
    background_aperture::Tally masked;
    background_aperture::MeasureFrame(img, W, H, {}, 6.0f, 8.0f, r3, 4.0f, 1.0f, no_rings, Geometry(),
                                      DetectorModules(), masked);
    CHECK(masked.apertures == 0);
    CHECK(all.apertures > 0);
}

// A bright background that is a curved function of the radius alone: the radial correction, read off
// the frame's ring background, makes the ring predict the disk better than it does without it.
TEST_CASE("BackgroundAperture_RadialCorrectionHelpsOnARadialBackground", "[Integration]") {
    const DiffractionGeometry g = Geometry();
    const auto background = [](double r) { return 2000.0 + 0.2 * (r - 150.0) * (r - 150.0); };
    std::mt19937 rng(3);
    ImagePreprocessorBuffer img(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            std::poisson_distribution<int> p(background(std::hypot(x - 300.0, y - 300.0)));
            img[static_cast<size_t>(y) * W + x] = p(rng);
        }
    const float low_q = 0.05f, dq = 0.005f;
    std::vector<float> profile(1000);
    for (size_t i = 0; i < profile.size(); i++) {
        const double q = low_q + (i + 0.5) * dq;
        const double r = 60.0 / 0.075 * std::tan(2.0 * std::asin(q * 1.0 / (4.0 * PI)));
        profile[i] = static_cast<float>(background(r));
    }
    const RingBackground rings(profile, static_cast<int32_t>(profile.size()), low_q, dq);
    const auto r3 = background_aperture::CandidateR3(8.0f, 14.0f);
    background_aperture::Tally tally;
    // Unclipped: a gradient this steep would put the high side of the ring above a photon clip.
    background_aperture::MeasureFrame(img, W, H, {}, 6.0f, 8.0f, r3, 0.0f, 1.0f, rings, g, DetectorModules(), tally);
    REQUIRE(tally.apertures_radial == tally.apertures);
    REQUIRE(tally.apertures > 100);
    const auto plain = tally.MeanSquare(false), radial = tally.MeanSquare(true);
    // The pre-scan's question: at the ring the plain differences choose, does the correction help.
    const int best = background_aperture::Best(plain);
    CHECK(radial[best] < 0.8 * plain[best]);
    CHECK(background_aperture::RadialHelps(tally, best));
    CHECK(background_aperture::RadialHelps(tally, best, background_aperture::ROTATION_RADIAL_MIN_GAIN));
    CHECK(!background_aperture::RadialHelps(tally, best, 0.99));
    // With the curvature taken out, the noise wins again and a wide ring is best.
    CHECK(r3[background_aperture::Best(radial)] >= 12.0f);
}

// A background falling steeply away from the beam: the flat clip cuts the beam-side half of every ring
// and the ring reads low. The integrator's clip follows the curve when the correction is on, and so does
// the apertures' corrected measurement - which then predicts the disk far better.
TEST_CASE("BackgroundAperture_CorrectedClipFollowsTheCurve", "[Integration]") {
    const DiffractionGeometry g = Geometry();
    const auto background = [](double r) { return 50000.0 - 100.0 * r; };
    std::mt19937 rng(5);
    ImagePreprocessorBuffer img(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            std::poisson_distribution<int> p(background(std::hypot(x - 300.0, y - 300.0)));
            img[static_cast<size_t>(y) * W + x] = p(rng);
        }
    const float low_q = 0.05f, dq = 0.005f;
    std::vector<float> profile(1000);
    for (size_t i = 0; i < profile.size(); i++) {
        const double q = low_q + (i + 0.5) * dq;
        const double r = 60.0 / 0.075 * std::tan(2.0 * std::asin(q * 1.0 / (4.0 * PI)));
        profile[i] = static_cast<float>(background(r));
    }
    const RingBackground rings(profile, static_cast<int32_t>(profile.size()), low_q, dq);
    background_aperture::Tally tally;
    background_aperture::MeasureFrame(img, W, H, {}, 6.0f, 8.0f, {14.0f}, 4.0f, 1.0f, rings, g, DetectorModules(),
                                      tally);
    REQUIRE(tally.apertures > 100);
    const auto plain = tally.MeanSquare(false), radial = tally.MeanSquare(true);
    CHECK(radial[0] < 0.1 * plain[0]);
    CHECK(background_aperture::RadialHelps(tally, 0));
}