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

#include "../image_analysis/geom_refinement/BeamCenterFromBackground.h"
#include "../common/DetectorSetup.h"
#include "../common/JFJochMath.h"

namespace {

// Solvent and air scatter: a decaying continuum with the water ring on it. The ring is where the
// leverage comes from - the continuum here is a pure exponential, on which g' is proportional to g
// and a shift and an amplitude are the same thing - so `ring` is how much there is to fit.
float background(float two_theta_rad, float ring) {
    const float ring_two_theta = 0.3239f;   // ~3.1 A at 1 A
    const float t = (two_theta_rad - ring_two_theta) / 0.035f;
    return 140.0f * std::exp(-two_theta_rad / 0.25f) + ring * std::exp(-0.5f * t * t);
}

// The projection the pre-scan hands over: the mean of a few tens of frames, laid out about
// geom_true, NAN where the detector has nothing. `shadow_sector` multiplies one sextant, the way a
// holder arm or a cryostream does.
std::vector<float> SynthesiseProjection(const DiffractionExperiment &experiment,
                                        const PixelMask &mask,
                                        const DiffractionGeometry &geom_true,
                                        float ring, float shadow_sector) {
    const auto W = static_cast<int>(experiment.GetXPixelsNumConv());
    const auto H = static_cast<int>(experiment.GetYPixelsNumConv());
    const auto &pixel_mask = mask.GetMask(experiment);

    std::vector<float> mean(static_cast<size_t>(W) * H, NAN);
    std::mt19937 rng(20260812);
    std::normal_distribution<float> gauss(0.0f, 1.0f);
    constexpr float FRAMES = 60.0f;   // the mean of this many frames, so the noise is that far down

    for (int y = 0; y < H; y++) {
        for (int x = 0; x < W; x++) {
            const size_t i = static_cast<size_t>(y) * W + x;
            if (pixel_mask[i] != 0)
                continue;
            float value = background(geom_true.TwoTheta_rad(static_cast<float>(x), static_cast<float>(y)), ring);
            const float phi = geom_true.Phi_rad(static_cast<float>(x), static_cast<float>(y));
            if (phi > 0.0f && phi < static_cast<float>(PI) / 3.0f)
                value *= shadow_sector;
            mean[i] = value + gauss(rng) * std::sqrt(value / FRAMES);
        }
    }
    return mean;
}

DiffractionExperiment TestExperiment() {
    DiffractionExperiment x(DetJF4M());
    x.IncidentEnergy_keV(WVL_1A_IN_KEV).DetectorDistance_mm(100.0f);
    // The band the estimator fits, 12-2.2 A, has to be on the detector, so start from its centre.
    x.BeamX_pxl(static_cast<float>(x.GetXPixelsNumConv()) / 2.0f)
     .BeamY_pxl(static_cast<float>(x.GetYPixelsNumConv()) / 2.0f);
    return x;
}

DiffractionGeometry OffsetBy(const DiffractionGeometry &geom, float dx, float dy) {
    DiffractionGeometry out = geom;
    out.BeamX_pxl(geom.GetBeamX_pxl() + dx).BeamY_pxl(geom.GetBeamY_pxl() + dy);
    return out;
}

}   // namespace

// The measurement: the background is isotropic in 2-theta about the beam, so a centre that is off
// shifts each azimuthal sector's radial profile by a different amount, and the shifts give the
// centre back. Nothing here is indexed, so this is what a de-novo run has to start from - and the
// estimate has to arrive, because a routine that quietly returns "not measurable" is
// indistinguishable from a careful refusal in every log line and every merging statistic.
TEST_CASE("BeamCenterFromBackground_RecoversAnInjectedOffset", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f);
    const auto estimate = FindBeamCenterFromBackground(x, pixel_mask, projection);

    REQUIRE(estimate.has_value());
    CHECK(estimate->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(0.5));
    CHECK(estimate->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(0.5));
    // And it has to say so precisely enough to be used: the caller commits at 1 px.
    CHECK(estimate->sigma_pxl < 1.0f);
}

// The sigma is the only thing standing between a bad background and a wrong geometry, so it has to
// grow when the ring it is fitting does not. With the ring at 1.4% of the continuum the centre is
// still found, but the fit says it is an order of magnitude less sure of it.
TEST_CASE("BeamCenterFromBackground_SigmaTracksTheLeverage", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const auto strong = FindBeamCenterFromBackground(
            x, pixel_mask, SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f));
    const auto weak = FindBeamCenterFromBackground(
            x, pixel_mask, SynthesiseProjection(x, pixel_mask, geom_true, 2.0f, 1.0f));

    REQUIRE(strong.has_value());
    REQUIRE(weak.has_value());
    CHECK(weak->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(1.0));
    CHECK(weak->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(1.0));
    CHECK(weak->sigma_pxl > 5.0f * strong->sigma_pxl);
}

// A holder arm or a cryostream is multiplicative and azimuthal, and a sector that is simply darker
// looks exactly like a sector whose profile has moved. The per-sector amplitude is what tells them
// apart: without it half a sextant of shadow reads as tens of pixels of centre error.
TEST_CASE("BeamCenterFromBackground_AnAzimuthalShadowIsNotACentreError", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = x.GetDiffractionGeometry();   // the centre is already right
    const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 0.5f);
    const auto estimate = FindBeamCenterFromBackground(x, pixel_mask, projection);

    REQUIRE(estimate.has_value());
    CHECK(estimate->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(1.0));
    CHECK(estimate->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(1.0));
}

// The same, with the detector tilted. Every pixel's 2-theta and azimuth, and the derivative of
// 2-theta with respect to the centre that the fit is built on, go through the detector rotation
// matrix, so the tilt is not a detail of the geometry here - it is in the Jacobian.
TEST_CASE("BeamCenterFromBackground_SurvivesADetectorTilt", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    x.PoniRot1_rad(0.005f).PoniRot2_rad(-0.003f);
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const auto [direct_x, direct_y] = geom_true.GetDirectBeam_pxl();
    REQUIRE(std::hypot(direct_x - geom_true.GetBeamX_pxl(), direct_y - geom_true.GetBeamY_pxl()) > 5.0f);

    const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f);
    const auto estimate = FindBeamCenterFromBackground(x, pixel_mask, projection);

    REQUIRE(estimate.has_value());
    CHECK(estimate->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(0.5));
    CHECK(estimate->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(0.5));
    CHECK(estimate->sigma_pxl < 1.0f);
}

// A centre hundreds of pixels from the detector's middle - a header that names the middle when the
// detector was raised - is the case the fit exists for, and it is the case an iteration cap decides.
// The shift a sector's regression can report is bounded by the width of the features it reads, so
// the walk advances by a bounded distance per iteration however far it still has to go: the count
// is a travel budget, and too small a one leaves the fit part way there, still walking, reporting
// the precision of its last step as though it had arrived.
TEST_CASE("BeamCenterFromBackground_RecoversACentreFarFromTheDetectorMiddle", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 40.0f, 300.0f);
    const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f);
    const auto estimate = FindBeamCenterFromBackground(x, pixel_mask, projection);

    REQUIRE(estimate.has_value());
    CHECK(estimate->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(2.0));
    CHECK(estimate->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(2.0));
}

// The fit is accumulated over row blocks so that a machine with more cores gives the same answer as
// one with fewer: the blocks are a property of the image and the fold runs over them in order.
TEST_CASE("BeamCenterFromBackground_DoesNotDependOnTheThreadCount", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f);
    const auto one = FindBeamCenterFromBackground(x, pixel_mask, projection, 1);
    const auto many = FindBeamCenterFromBackground(x, pixel_mask, projection, 8);

    REQUIRE(one.has_value());
    REQUIRE(many.has_value());
    CHECK(one->beam_x_pxl == many->beam_x_pxl);
    CHECK(one->beam_y_pxl == many->beam_y_pxl);
    CHECK(one->sigma_pxl == many->sigma_pxl);
}

// FindBeamCenter is the walk with an FFT capture in front of it: the capture chooses the basin and
// the walk finishes inside it and reports what it knows the answer to. The tests below are about
// the composition - that the walk really is seeded where the capture landed, that a decline falls
// back the way it is documented to, and that the beam stop is taken out of the image the capture
// scores.

namespace {

// A projection that is NAN everywhere except one exactly centrosymmetric patch. The capture has an
// unambiguous answer on it (a correlation coefficient cannot exceed 1, and the patch reaches 1 at
// its own centre) while the ring fit has nothing at all: no ring of the 12-2.2 A band is covered
// by a patch that small, from any start.
std::vector<float> SymmetricPatchOnly(const DiffractionExperiment &experiment, float cx, float cy,
                                      int half_size) {
    const auto W = static_cast<int>(experiment.GetXPixelsNumConv());
    const auto H = static_cast<int>(experiment.GetYPixelsNumConv());
    std::vector<float> mean(static_cast<size_t>(W) * H, NAN);
    for (int y = static_cast<int>(cy) - half_size; y <= static_cast<int>(cy) + half_size; y++)
        for (int x = static_cast<int>(cx) - half_size; x <= static_cast<int>(cx) + half_size; x++) {
            // Keyed on the canonical member of the +- pair about (cx, cy), so the patch is exactly
            // its own mirror there.
            int64_t dx = std::lround(2.0f * (static_cast<float>(x) - cx));
            int64_t dy = std::lround(2.0f * (static_cast<float>(y) - cy));
            if (dy < 0 || (dy == 0 && dx < 0)) {
                dx = -dx;
                dy = -dy;
            }
            uint64_t k = static_cast<uint64_t>(dx) * 0x9e3779b97f4a7c15ULL
                         + static_cast<uint64_t>(dy) * 0xc2b2ae3d27d4eb4fULL;
            k ^= k >> 31;
            k *= 0xbf58476d1ce4e5b9ULL;
            k ^= k >> 29;
            mean[static_cast<size_t>(y) * W + x] = 100.0f + static_cast<float>(k % 4096u) * 0.25f;
        }
    return mean;
}

}   // namespace

// The case the composition exists for. With the centre in the file destroyed - the value a
// beamline writes when it has nothing to write - the walk has no ring of the fitted band covered
// from where it starts, so it declines and says nothing at all. The capture does not care how
// wrong the file is: it scores every centre on the detector in one transform set, and the walk
// then finishes from there with a real fit sigma, low enough that a caller would commit it.
TEST_CASE("FindBeamCenter_AnswersWhereTheWalkDeclines", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = x.GetDiffractionGeometry();
    const auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f);
    x.BeamX_pxl(0.0f).BeamY_pxl(0.0f);   // the destroyed header

    CHECK(!FindBeamCenterFromBackground(x, pixel_mask, projection).has_value());

    const auto composed = FindBeamCenter(x, pixel_mask, projection);
    REQUIRE(composed.has_value());
    CHECK(composed->beam_x_pxl == Catch::Approx(geom_true.GetBeamX_pxl()).margin(1.0));
    CHECK(composed->beam_y_pxl == Catch::Approx(geom_true.GetBeamY_pxl()).margin(1.0));
    CHECK(composed->sigma_pxl < 1.0f);
}

// The end of the fallback chain. Where the walk has nothing to fit from either start - neither from
// the capture nor from the centre in the file - the capture stands alone, and it is disclosed as a
// capture: BEAM_CENTER_CAPTURE_SIGMA_PXL sits above every ceiling a caller commits geometry on, so
// the answer reaches the consumers that only need a hypothesis to test and is refused by the two
// that would change the geometry with it.
TEST_CASE("FindBeamCenter_FallsBackToTheCaptureAlone", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const float cx = 700.5f, cy = 640.0f;
    const auto projection = SymmetricPatchOnly(x, cx, cy, 90);

    CHECK(!FindBeamCenterFromBackground(x, pixel_mask, projection).has_value());

    const auto composed = FindBeamCenter(x, pixel_mask, projection);
    REQUIRE(composed.has_value());
    CHECK(composed->beam_x_pxl == cx);
    CHECK(composed->beam_y_pxl == cy);
    CHECK(composed->sigma_pxl == BEAM_CENTER_CAPTURE_SIGMA_PXL);
}

// The beam stop is blanked out of the image the capture scores. A large one-sided opaque region is
// centrosymmetric about ITS own centre, and where it is big enough that preference beats the
// background's - measured on a real umbra over 9.6 % of the detector, the capture landed 48 px out
// and masking it put it back to 1.1 px. The mask that reaches the capture already has the stop in
// it, so this costs nothing but has to keep working.
TEST_CASE("FindBeamCenter_BlanksTheBeamStopOutOfTheCapture", "[BeamCenter]") {
    DiffractionExperiment x = TestExperiment();
    PixelMask pixel_mask(x);

    const DiffractionGeometry geom_true = x.GetDiffractionGeometry();
    auto projection = SynthesiseProjection(x, pixel_mask, geom_true, 60.0f, 1.0f);

    // An umbra on one side of the beam: opaque, so the pixels under it carry no background.
    const auto W = static_cast<int>(x.GetXPixelsNumConv());
    const auto H = static_cast<int>(x.GetYPixelsNumConv());
    const float sx = geom_true.GetBeamX_pxl() + 330.0f, sy = geom_true.GetBeamY_pxl() + 70.0f;
    std::vector<uint32_t> stop(static_cast<size_t>(W) * H, 0);
    for (int y = 0; y < H; y++)
        for (int x_pxl = 0; x_pxl < W; x_pxl++)
            if (std::hypot(x_pxl - sx, y - sy) < 300.0f) {
                stop[static_cast<size_t>(y) * W + x_pxl] = 1;
                projection[static_cast<size_t>(y) * W + x_pxl] = 0.0f;
            }

    const auto unmasked = BeamCenterFFTScore(W, H, projection);
    REQUIRE(!unmasked.point.empty());

    pixel_mask.LoadBeamStopMask(x, stop);
    BeamCenterFFTResult capture;
    const auto composed = FindBeamCenter(x, pixel_mask, projection, 0, &capture);
    REQUIRE(composed.has_value());
    REQUIRE(!capture.point.empty());

    const float masked_error = std::hypot(capture.point[0].beam_x_pxl - geom_true.GetBeamX_pxl(),
                                          capture.point[0].beam_y_pxl - geom_true.GetBeamY_pxl());
    const float unmasked_error = std::hypot(unmasked.point[0].beam_x_pxl - geom_true.GetBeamX_pxl(),
                                            unmasked.point[0].beam_y_pxl - geom_true.GetBeamY_pxl());
    CHECK(masked_error < 12.0f);
    CHECK(masked_error <= unmasked_error);
}