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

#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));
}