// 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/BeamCenterFromSpots.h"
#include "../common/DetectorSetup.h"
#include "../common/JFJochMath.h"

namespace {

constexpr float WEDGE_DEG = 1.0f;
constexpr int PAIRS = 30;
constexpr float PAIR_SPACING_DEG = 180.0f / PAIRS;
constexpr float POSITION_NOISE_PXL = 0.15f;

struct Sweep {
    std::vector<float> angle_deg;
    std::vector<BeamCenterSpot> spots;
};

DiffractionExperiment TestExperiment(const Coord &spindle) {
    DiffractionExperiment x(DetDECTRIS(1500, 1600, "Test detector", ""));
    x.IncidentEnergy_keV(WVL_1A_IN_KEV).DetectorDistance_mm(180.0f);
    x.BeamX_pxl(static_cast<float>(x.GetXPixelsNumConv()) / 2.0f)
     .BeamY_pxl(static_cast<float>(x.GetYPixelsNumConv()) / 2.0f);
    x.Goniometer(GoniometerAxis("omega", 0.0f, WEDGE_DEG, spindle, {}));
    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;
}

// |F(h)| = |F(-h)|, which is what tells a true Friedel match from a lattice-shifted one, so the
// intensity has to depend on the reflection through |h| |k| |l| alone.
float StructureFactor(int h, int k, int l) {
    const uint32_t key = 2654435761u * (std::abs(h) + 41u * std::abs(k) + 1723u * std::abs(l) + 1u);
    return 50.0f + static_cast<float>(key % 4000u);
}

// A synthetic rotation sweep. An orthorhombic lattice in a fixed orientation is turned about the
// spindle; every reflection meets the Ewald sphere exactly twice, and each meeting is put on the
// detector at the geometry the estimator is asked to find. Only the frames of the pre-scan sample
// are kept, so what comes out is what the pre-pass sees: pairs half a turn apart, and about a
// thirtieth of the sweep.
Sweep MakeSweep(const DiffractionExperiment &experiment, const DiffractionGeometry &truth, int pairs) {
    const Coord spindle = experiment.GetGoniometer()->GetAxis().Normalize();
    const Coord S0 = truth.GetScatteringVector();

    Sweep sweep;
    for (int i = 0; i < pairs; i++) {
        sweep.angle_deg.push_back(i * PAIR_SPACING_DEG);
        sweep.angle_deg.push_back(i * PAIR_SPACING_DEG + 180.0f);
    }

    const RotMatrix orientation(0.7f, Coord(0.3f, 0.5f, 0.8f));
    const Coord astar = orientation * Coord(1.0f / 55.0f, 0, 0);
    const Coord bstar = orientation * Coord(0, 1.0f / 65.0f, 0);
    const Coord cstar = orientation * Coord(0, 0, 1.0f / 75.0f);
    const float q_max = 1.0f / 2.4f;

    std::mt19937 rng(20260812);
    std::normal_distribution<float> noise(0.0f, POSITION_NOISE_PXL);

    for (int h = -30; h <= 30; h++)
        for (int k = -35; k <= 35; k++)
            for (int l = -40; l <= 40; l++) {
                if (h == 0 && k == 0 && l == 0)
                    continue;
                const Coord q0 = astar * static_cast<float>(h) + bstar * static_cast<float>(k)
                                 + cstar * static_cast<float>(l);
                const float q_sq = q0 * q0;
                if (q_sq > q_max * q_max)
                    continue;

                // Turning about the spindle leaves q.m alone and rotates the rest, so the Ewald
                // condition q.S0 = -|q|^2/2 is one sinusoid in the rotation angle: two solutions,
                // which are the reflection's two crossings.
                //
                // Which WAY the angle turns the crystal is rugnux's convention, not a free choice:
                // BraggPredictionRot takes the offset to the diffracting condition as
                // -atan2(sin, cos) in this same frame, i.e. dq/dphi = -m x q. The sweep this
                // generator makes is otherwise time-reversed with respect to every real dataset,
                // which no estimator that reads only positions can notice.
                const Coord along = spindle * (q0 * spindle);
                const Coord across = q0 - along;
                const Coord turned = across % spindle;
                const float a = across * S0, b = turned * S0;
                const float radius = std::hypot(a, b);
                const float target = -0.5f * q_sq - along * S0;
                if (radius <= 0.0f || std::abs(target) > radius)
                    continue;
                const float phase = std::atan2(b, a);
                const float offset = std::acos(target / radius);

                for (const float sign: {-1.0f, 1.0f}) {
                    float phi_deg = (phase + sign * offset) * 180.0f / static_cast<float>(PI);
                    while (phi_deg < 0.0f) phi_deg += 360.0f;
                    while (phi_deg >= 360.0f) phi_deg -= 360.0f;
                    for (size_t frame = 0; frame < sweep.angle_deg.size(); frame++) {
                        if (phi_deg < sweep.angle_deg[frame] || phi_deg >= sweep.angle_deg[frame] + WEDGE_DEG)
                            continue;
                        const float phi = phi_deg * static_cast<float>(PI) / 180.0f;
                        const Coord q = along + across * std::cos(phi) + turned * std::sin(phi);
                        const auto [x, y] = truth.RecipToDetector(q);
                        if (!std::isfinite(x) || x < 0 || y < 0
                            || x >= experiment.GetXPixelsNumConv() || y >= experiment.GetYPixelsNumConv())
                            continue;
                        sweep.spots.push_back({x + noise(rng), y + noise(rng), StructureFactor(h, k, l),
                                               static_cast<int>(frame)});
                    }
                }
            }
    return sweep;
}

}   // namespace

// The measurement: the sweep at phi+180 is the mirror image of the sweep at phi in the coordinate
// along the spindle, and every reflection crosses the Ewald sphere twice, mirrored in the other
// one. Between them the two put the beam centre where nothing has been indexed yet. Both a
// horizontal spindle and a vertical one, because which estimator supplies which coordinate is
// decided by the goniometer axis and one of the two is a real detector's arrangement.
TEST_CASE("BeamCenterFromSpots_RecoversAnInjectedOffset", "[BeamCenter]") {
    const Coord spindle = GENERATE(Coord(1, 0, 0), Coord(0, 1, 0));
    const DiffractionExperiment x = TestExperiment(spindle);
    const DiffractionGeometry truth = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const Sweep sweep = MakeSweep(x, truth, PAIRS);

    const auto estimate = FindBeamCenterFromSpotSymmetry(x, sweep.angle_deg, sweep.spots);

    REQUIRE(estimate.has_value());
    CHECK(estimate->beam_x_pxl == Catch::Approx(truth.GetBeamX_pxl()).margin(0.5));
    CHECK(estimate->beam_y_pxl == Catch::Approx(truth.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 header is where the search starts and nothing more. A file whose centre is 15 px out has
    // to give the same answer as one whose centre is right, or the estimate is partly the header's.
    DiffractionExperiment moved = x;
    moved.BeamX_pxl(x.GetBeamX_pxl() + 15.0f).BeamY_pxl(x.GetBeamY_pxl() - 15.0f);
    const auto from_elsewhere = FindBeamCenterFromSpotSymmetry(moved, sweep.angle_deg, sweep.spots);
    REQUIRE(from_elsewhere.has_value());
    CHECK(from_elsewhere->beam_x_pxl == Catch::Approx(estimate->beam_x_pxl).margin(0.1));
    CHECK(from_elsewhere->beam_y_pxl == Catch::Approx(estimate->beam_y_pxl).margin(0.1));
}

// A hot pixel, the beam-stop halo or the edge of a mask sits at the SAME place on every frame, so
// it pairs with itself and votes at twice its own position - one vote per frame pair, which is
// enough to out-vote the real peak. On real data that was two crystals wrong by 12 and 1.9 px.
TEST_CASE("BeamCenterFromSpots_APersistentArtefactIsNotABeamCentre", "[BeamCenter]") {
    const DiffractionExperiment x = TestExperiment(Coord(1, 0, 0));
    const DiffractionGeometry truth = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    Sweep sweep = MakeSweep(x, truth, PAIRS);

    for (size_t frame = 0; frame < sweep.angle_deg.size(); frame++)
        for (int i = 0; i < 12; i++)
            sweep.spots.push_back({truth.GetBeamX_pxl() + 40.0f + 3.0f * i,
                                   truth.GetBeamY_pxl() - 25.0f - 2.0f * i,
                                   9000.0f, static_cast<int>(frame)});

    const auto estimate = FindBeamCenterFromSpotSymmetry(x, sweep.angle_deg, sweep.spots);

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

// The vote reaches 30 px about wherever the search starts, so a header further out than that leaves
// the true peak outside the window and a lattice-shifted one wins it - and wins it cleanly, with a
// frame-pair scatter as small as a correct answer's. What separates them is that the wrong one
// depends on where the search began: started elsewhere, the estimator finds something else. The
// reported sigma has to carry that, or the caller commits a centre tens of pixels out.
TEST_CASE("BeamCenterFromSpots_AnAnswerThatDependsOnTheHeaderIsNotACentre", "[BeamCenter]") {
    const DiffractionExperiment x = TestExperiment(Coord(1, 0, 0));
    const DiffractionGeometry truth = OffsetBy(x.GetDiffractionGeometry(), 35.0f, 35.0f);
    const Sweep sweep = MakeSweep(x, truth, PAIRS);

    const auto estimate = FindBeamCenterFromSpotSymmetry(x, sweep.angle_deg, sweep.spots);

    REQUIRE(estimate.has_value());
    // It answers, and the answer is nowhere near the truth - that is the failure mode, not a bug.
    CHECK(std::hypot(estimate->beam_x_pxl - truth.GetBeamX_pxl(),
                     estimate->beam_y_pxl - truth.GetBeamY_pxl()) > 10.0f);
    // What must not happen is that it says so quietly: the caller commits at 1 px.
    CHECK(estimate->sigma_pxl > 1.0f);
}

// Not every spot is one of a reflection's two crossings. A second read-out, a satellite, a family
// of detector artefacts - anything that shadows the real spots at a fixed displacement pairs with
// them at that displacement and votes as a tooth of its own, and here that tooth is TWICE the true
// one's height. What the decoys cannot copy is when their partner appears: the sweep angle between
// two crossings is fixed by where the first one is, and a decoy's is not.
TEST_CASE("BeamCenterFromSpots_ATallerDecoyToothDoesNotWin", "[BeamCenter]") {
    const DiffractionExperiment x = TestExperiment(Coord(1, 0, 0));
    const DiffractionGeometry truth = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    Sweep sweep = MakeSweep(x, truth, PAIRS);

    std::mt19937 rng(20260813);
    const size_t recorded = sweep.spots.size();
    for (size_t i = 0; i < recorded; i++) {
        const BeamCenterSpot &spot = sweep.spots[i];
        sweep.spots.push_back({spot.x, spot.y + 9.0f, spot.intensity,
                               static_cast<int>(rng() % sweep.angle_deg.size())});
    }

    const auto estimate = FindBeamCenterFromSpotSymmetry(x, sweep.angle_deg, sweep.spots);

    REQUIRE(estimate.has_value());
    // The decoy tooth is 4.5 px away in the coordinate the second crossing supplies.
    CHECK(estimate->beam_y_pxl == Catch::Approx(truth.GetBeamY_pxl()).margin(0.5));
}

// The Friedel mirror needs half a turn. On a screening wedge, on stills, on anything short there is
// no partner frame to mirror onto, and the estimator has to say so rather than answer from the
// pairs it has not got - the background estimator is what covers those.
TEST_CASE("BeamCenterFromSpots_DeclinesOnAShortSweep", "[BeamCenter]") {
    const DiffractionExperiment x = TestExperiment(Coord(1, 0, 0));
    const DiffractionGeometry truth = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const Sweep sweep = MakeSweep(x, truth, PAIRS);

    // The first half of the sample only: 30 frames spread over 174 deg, none of them a pair.
    std::vector<float> angle_deg;
    std::vector<BeamCenterSpot> spots;
    for (size_t frame = 0; frame < sweep.angle_deg.size(); frame += 2)
        angle_deg.push_back(sweep.angle_deg[frame]);
    for (const auto &spot: sweep.spots)
        if (spot.frame % 2 == 0)
            spots.push_back({spot.x, spot.y, spot.intensity, spot.frame / 2});

    CHECK_FALSE(FindBeamCenterFromSpotSymmetry(x, angle_deg, spots).has_value());
}

// No beamline can hold the spindle exactly perpendicular to the beam, or exactly on a lab axis -
// and the file always claims it does: every master of the regression set writes the goniometer axis
// as an exact lab vector, so the deviation has to come from the spots or from nowhere. What the
// estimator cannot survive is the part turned ABOUT THE BEAM: both mirror lines turn with it, and a
// mirror taken about the lab axis instead smears the vote by twice that angle times the other
// coordinate until a neighbouring comb tooth outvotes the true one. It does not decline when that
// happens - it answers, at the same sigma as a good fit.
//
// The other component, the spindle tipped towards the beam, is here too and must not be corrected
// in the same way: mirroring about the spindle itself would move the direct beam by twice that
// angle times the distance. Both mirror planes contain the beam whatever the spindle does.
TEST_CASE("BeamCenterFromSpots_SurvivesASpindleOffPerpendicular", "[BeamCenter]") {
    const float azimuth = 0.002f;   // turned about the beam
    const float tipped = 0.002f;    // and out of the plane perpendicular to it
    const bool vertical = GENERATE(false, true);
    const Coord in_plane = vertical ? Coord(-std::sin(azimuth), std::cos(azimuth), 0)
                                    : Coord(std::cos(azimuth), std::sin(azimuth), 0);
    const Coord spindle = in_plane * std::cos(tipped) + Coord(0, 0, 1) * std::sin(tipped);

    const DiffractionExperiment x = TestExperiment(spindle.Normalize());
    const DiffractionGeometry truth = OffsetBy(x.GetDiffractionGeometry(), 3.0f, -2.5f);
    const Sweep sweep = MakeSweep(x, truth, PAIRS);

    SpindleEstimate fitted;
    const auto estimate = FindBeamCenterFromSpotSymmetry(x, sweep.angle_deg, sweep.spots, &fitted);

    REQUIRE(estimate.has_value());
    // The azimuth is measured, not assumed, and the vote is taller for it.
    CHECK(fitted.azimuth_rad == Catch::Approx(azimuth).margin(2e-4));
    CHECK(fitted.vote_excess > fitted.vote_excess_nominal);
    CHECK(estimate->beam_x_pxl == Catch::Approx(truth.GetBeamX_pxl()).margin(0.5));
    CHECK(estimate->beam_y_pxl == Catch::Approx(truth.GetBeamY_pxl()).margin(0.5));

    // And this is the failure it exists to remove: without it the same sweep is answered, with a
    // sigma that says nothing is wrong, further away. How much further depends on what else is in
    // the estimator - once the second crossing is guarded by its timing test the comb mis-pick that
    // made this several pixels is already gone, and what is left is the azimuth's own smearing of
    // the Friedel vote. So the claim is a ratio, which is what the correction actually owns.
    const auto uncorrected = FindBeamCenterFromSpotSymmetry(x, sweep.angle_deg, sweep.spots);
    REQUIRE(uncorrected.has_value());
    const float corrected_error = std::hypot(estimate->beam_x_pxl - truth.GetBeamX_pxl(),
                                             estimate->beam_y_pxl - truth.GetBeamY_pxl());
    CHECK(std::hypot(uncorrected->beam_x_pxl - truth.GetBeamX_pxl(),
                     uncorrected->beam_y_pxl - truth.GetBeamY_pxl()) > 3.0f * corrected_error);
}