// 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 "../image_analysis/geom_refinement/RingsFromProfile.h"
#include "../image_analysis/geom_refinement/AssignSpotsToRings.h"
#include "../image_analysis/geom_refinement/PowderAutoSeed.h"
#include "../image_analysis/geom_refinement/PowderCalibration.h"
#include "../common/Definitions.h"
#include "../common/JFJochMath.h"

namespace {

constexpr UnitCell LAB6{LAB6_CELL_A, LAB6_CELL_A, LAB6_CELL_A, 90.0f, 90.0f, 90.0f};
const std::vector<float> LAB6_RINGS = CalculateXtalRings(LAB6);

// A (q x azimuth) powder profile as the azimuthal integration would build it: the rings sit where
// geom_true puts them, but every pixel is binned with geom_assumed - which is the whole point, since a
// wrong assumed geometry is what makes a ring's apparent q wander with azimuth.
std::vector<float> SynthesiseProfile(const AzimuthalIntegrationMapping &mapping,
                                     const DiffractionGeometry &geom_assumed,
                                     const DiffractionGeometry &geom_true) {
    const auto &settings = mapping.Settings();
    const int32_t q_bins = mapping.GetQBinCount();
    const int32_t azim_bins = mapping.GetAzimuthalBinCount();
    std::vector<float> profile(static_cast<size_t>(q_bins) * azim_bins, 100.0f);   // flat background

    for (const float q_ring : LAB6_RINGS) {
        const float d = static_cast<float>(2.0 * PI) / q_ring;
        if (d <= geom_true.GetWavelength_A() / 2.0f)
            continue;
        for (int t = 0; t < 3600; ++t) {
            const float phi_true = static_cast<float>(2.0 * PI * t / 3600.0);
            const auto [px, py] = geom_true.ResPhiToPxl(d, phi_true);
            if (!std::isfinite(px) || !std::isfinite(py))
                continue;
            const float q_obs = geom_assumed.PxlToQ(px, py);
            float phi_deg = geom_assumed.Phi_rad(px, py) * 180.0f / static_cast<float>(PI);
            if (phi_deg < 0.0f)
                phi_deg += 360.0f;
            const uint16_t bin = settings.GetBin(q_obs, phi_deg);
            if (bin == UINT16_MAX)
                continue;
            // Lay a narrow peak over the neighbouring q bins of this azimuthal row.
            const int q_bin = bin % q_bins, phi_bin = bin / q_bins;
            for (int k = -3; k <= 3; ++k) {
                const int b = q_bin + k;
                if (b < 0 || b >= q_bins)
                    continue;
                profile[static_cast<size_t>(phi_bin) * q_bins + b] +=
                        2000.0f * std::exp(-0.5f * static_cast<float>(k * k) / (1.2f * 1.2f));
            }
        }
    }
    return profile;
}

}   // namespace

// The measurement this is for: a powder ring is a conic centred on the beam, so a wrong beam centre
// makes its apparent radius oscillate once per turn. Recovering the centre from that needs neither the
// calibrant's lattice constant nor the detector distance - only that the ring be round.
TEST_CASE("RingsFromProfile_RecoversBeamCenter", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    auto azint = x.GetAzimuthalIntegrationSettings();
    azint.AzimuthalBinCount(32);
    x.ImportAzimuthalIntegrationSettings(azint);

    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    const DiffractionGeometry geom_assumed = x.GetDiffractionGeometry();
    DiffractionGeometry geom_true = geom_assumed;
    geom_true.BeamX_pxl(geom_assumed.GetBeamX_pxl() + 6.0f)
             .BeamY_pxl(geom_assumed.GetBeamY_pxl() - 4.0f);

    const auto profile = SynthesiseProfile(mapping, geom_assumed, geom_true);
    const auto rings = RingsFromAzimuthalProfile(profile, mapping, geom_assumed, LAB6_RINGS);

    // Several rings, sampled all the way round: without azimuthal coverage there is no centre to find.
    REQUIRE(rings.size() > 64);

    RingOptimizer optimizer(geom_assumed);
    const auto fitted = optimizer.Run(rings);

    CHECK(fitted.GetBeamX_pxl() == Catch::Approx(geom_true.GetBeamX_pxl()).margin(0.5));
    CHECK(fitted.GetBeamY_pxl() == Catch::Approx(geom_true.GetBeamY_pxl()).margin(0.5));
    // The starting point was wrong by 6 and 4 pixels, so a fit that did nothing would fail the above -
    // but check explicitly that it moved toward the truth rather than merely landing near it.
    CHECK(std::abs(fitted.GetBeamX_pxl() - geom_true.GetBeamX_pxl())
          < std::abs(geom_assumed.GetBeamX_pxl() - geom_true.GetBeamX_pxl()));
}

// The same round trip with the detector tilted. A tilt and a centre error BOTH show up as cos(phi);
// what separates them is that the tilt's amplitude grows as the ring radius squared, so it takes
// several rings to tell them apart. This mainly guards the conventions: RingOptimizer open-codes its
// rotation instead of going through DiffractionGeometry, and this holds the two against each other.
// Fewer ring points than the centred case is expected - a tilt this size carries part of some rings
// out of the extractor's search window, which is centred on where the ring is EXPECTED to be.
TEST_CASE("RingsFromProfile_RecoversTilt", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    auto azint = x.GetAzimuthalIntegrationSettings();
    azint.AzimuthalBinCount(64);
    x.ImportAzimuthalIntegrationSettings(azint);

    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    const DiffractionGeometry geom_assumed = x.GetDiffractionGeometry();
    DiffractionGeometry geom_true = geom_assumed;
    geom_true.PoniRot1_rad(0.02f).PoniRot2_rad(-0.015f);

    const auto profile = SynthesiseProfile(mapping, geom_assumed, geom_true);
    const auto rings = RingsFromAzimuthalProfile(profile, mapping, geom_assumed, LAB6_RINGS);
    REQUIRE(rings.size() > 60);

    RingOptimizer optimizer(geom_assumed);
    const auto fitted = optimizer.Run(rings);

    CHECK(fitted.GetPoniRot1_rad() == Catch::Approx(0.02).margin(0.004));
    CHECK(fitted.GetPoniRot2_rad() == Catch::Approx(-0.015).margin(0.004));
}

// One azimuthal bin is a plain radial profile: the ring has been averaged over every direction, so
// nothing is left to say where its centre is. Refuse rather than return points that cannot constrain it.
TEST_CASE("RingsFromProfile_NeedsAzimuthalBins", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);

    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);
    REQUIRE(mapping.GetAzimuthalBinCount() == 1);

    const std::vector<float> profile(static_cast<size_t>(mapping.GetQBinCount()), 1000.0f);
    CHECK(RingsFromAzimuthalProfile(profile, mapping, x.GetDiffractionGeometry(), LAB6_RINGS).empty());
}

// A profile with no rings in it must yield no ring points: the peak has to stand clear of the scatter
// of the background either side of it, or every azimuthal sector would contribute its largest noise
// excursion as though it were a measurement.
TEST_CASE("RingsFromProfile_FlatProfileGivesNothing", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    auto azint = x.GetAzimuthalIntegrationSettings();
    azint.AzimuthalBinCount(32);
    x.ImportAzimuthalIntegrationSettings(azint);

    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    const std::vector<float> profile(
            static_cast<size_t>(mapping.GetQBinCount()) * mapping.GetAzimuthalBinCount(), 100.0f);
    CHECK(RingsFromAzimuthalProfile(profile, mapping, x.GetDiffractionGeometry(), LAB6_RINGS).empty());
}

// The distance recovered from the rings alone, with the header deliberately wrong. This is the property
// the whole seed exists for: a calibration must not need to be told the distance, because the header is
// the number a calibration is run to check. Nothing here reads the assumed distance except to bin the
// profile - the answer comes from the ring radii, the wavelength and the pixel size.
TEST_CASE("PowderAutoSeed_RecoversDistanceFromAWrongHeader", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    auto azint = x.GetAzimuthalIntegrationSettings();
    azint.AzimuthalBinCount(32);
    x.ImportAzimuthalIntegrationSettings(azint);

    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    const DiffractionGeometry geom_assumed = x.GetDiffractionGeometry();
    const float true_distance = geom_assumed.GetDetectorDistance_mm();

    DiffractionGeometry geom_true = geom_assumed;
    const auto profile = SynthesiseProfile(mapping, geom_assumed, geom_true);

    // The rings the profile actually shows, found with no calibrant involved at all.
    const auto observed = RingRadiiFromProfile(profile, mapping, geom_assumed);
    REQUIRE(observed.size() >= 2);

    const auto [r_min, r_max] = ProfileRadiusRange_pxl(mapping, geom_assumed);
    const auto candidates = CandidateDistancesFromPowderRings(observed, LAB6_RINGS, geom_assumed,
                                                              r_min, r_max);
    REQUIRE(!candidates.empty());

    // The true distance is among the candidates. It need not be the FIRST: a powder pattern has real
    // distance aliases - for a cubic primitive standard the rings go as sqrt(N), so scaling by sqrt(2)
    // maps ring N onto ring 2N - which is exactly why the caller fits every candidate and lets the
    // residual choose rather than trusting the best score.
    const bool found = std::any_of(candidates.begin(), candidates.end(),
                                   [&](const DistanceCandidate &c) {
        return std::abs(c.distance_mm - true_distance) < 0.02f * true_distance;
    });
    CHECK(found);
}

// The seed measures radii, and a radius does not care what distance was assumed when the profile was
// binned: bin i holds the pixels at one particular radius whatever q that radius was called. So the
// ring radii recovered from a profile binned at half the true distance are the same radii - which is
// what lets the distance be measured before it is known.
TEST_CASE("PowderAutoSeed_RingRadiiDoNotDependOnTheAssumedDistance", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    auto azint = x.GetAzimuthalIntegrationSettings();
    azint.AzimuthalBinCount(32);
    x.ImportAzimuthalIntegrationSettings(azint);
    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    const DiffractionGeometry geom = x.GetDiffractionGeometry();
    const auto profile = SynthesiseProfile(mapping, geom, geom);
    const auto observed = RingRadiiFromProfile(profile, mapping, geom);
    REQUIRE(observed.size() >= 3);

    // Every ring the finder reports must sit on a real LaB6 ring of this geometry, to a pixel.
    for (const auto &o : observed) {
        float nearest = std::numeric_limits<float>::max();
        for (const float q : LAB6_RINGS) {
            const float d = static_cast<float>(2.0 * PI) / q;
            if (d <= geom.GetWavelength_A() / 2.0f)
                continue;   // past the Ewald limit - no such ring on any detector
            const auto [px, py] = geom.ResPhiToPxl(d, 0.0f);
            if (!std::isfinite(px) || !std::isfinite(py))
                continue;
            const float r = std::hypot(px - geom.GetBeamX_pxl(), py - geom.GetBeamY_pxl());
            nearest = std::min(nearest, std::abs(r - o.radius_pxl));
        }
        CHECK(nearest < 2.0f);
    }
}

// Where a ring APPEARS in a profile binned with one geometry, if the truth is another. The round trip
// has to be exact when the two agree, or a correctly-seeded run would move its own search windows off
// the rings it is looking for.
TEST_CASE("PowderAutoSeed_RingTrackRoundTripsWhenTheGeometryIsRight", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    const DiffractionGeometry geom = x.GetDiffractionGeometry();
    constexpr int32_t AZIM_BINS = 32;

    for (const float q : {1.0f, 2.0f, 3.0f}) {
        const auto track = ProfileRingTrack(q, geom, geom, AZIM_BINS);
        REQUIRE(track.size() == AZIM_BINS);
        for (const float t : track)
            if (std::isfinite(t))
                CHECK(t == Catch::Approx(q).epsilon(1e-3));
    }
}

// ...and a beam centre that is wrong makes the ring wander in q ONCE PER TURN, which is the property
// the track exists to follow. A single window centred on one q cannot hold a ring that does this, which
// is why the extraction searches sector by sector when it has a seed to search from.
TEST_CASE("PowderAutoSeed_RingTrackFollowsAWrongBeamCentre", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    const DiffractionGeometry binned = x.GetDiffractionGeometry();
    DiffractionGeometry truth = binned;
    truth.BeamX_pxl(binned.GetBeamX_pxl() + 20.0f);
    constexpr int32_t AZIM_BINS = 32;

    const auto track = ProfileRingTrack(2.0f, truth, binned, AZIM_BINS);
    float lo = std::numeric_limits<float>::max(), hi = std::numeric_limits<float>::lowest();
    int finite = 0;
    for (const float t : track)
        if (std::isfinite(t)) { lo = std::min(lo, t); hi = std::max(hi, t); ++finite; }
    REQUIRE(finite > AZIM_BINS / 2);
    // It has to swing by appreciably more than nothing, or there would be no need to track it...
    CHECK(hi - lo > 0.01f);
    // ...and the swing has to bracket the ring's own q, since the centre error only moves it.
    CHECK(lo < 2.0f);
    CHECK(hi > 2.0f);
}

// The beam-centre offset read off the ring radii alone - no calibrant, no distance. This is the seed
// that lets a header centre further out than the extraction window be recovered at all.
TEST_CASE("PowderAutoSeed_RecoversTheBeamCentreOffset", "[DetGeomCalib]") {
    DiffractionExperiment x(DetJF4M());
    x.QSpacingForAzimInt_recipA(0.004).QRangeForAzimInt_recipA(0.5, 4.0);
    auto azint = x.GetAzimuthalIntegrationSettings();
    azint.AzimuthalBinCount(32);
    x.ImportAzimuthalIntegrationSettings(azint);

    // A beam centre on the detector, not at its corner where the fixture leaves it. The seed reads a
    // once-per-turn wobble, so it needs the rings to go round: with the beam in the corner only a
    // quarter of the azimuth carries any ring at all and the two components cannot separate.
    x.BeamX_pxl(1000.0f).BeamY_pxl(1050.0f);

    PixelMask pixel_mask(x);
    AzimuthalIntegrationMapping mapping(x, pixel_mask);

    const DiffractionGeometry geom_assumed = x.GetDiffractionGeometry();
    DiffractionGeometry geom_true = geom_assumed;
    // A small offset, because the rings this fixture draws are about a pixel wide and its q bins are
    // finer than a real run's. Past a few pixels those sharp rings split in the azimuthal average into
    // the two HORNS of the sinusoid they trace - density piles up where r(phi) turns round, at R+|d| and
    // R-|d| - and the radius finder then reports two rings where there is one. Real powder rings are
    // broad enough to smear that out, which is why the offset this recovers on a real LaB6 exposure is
    // four times the one it can be shown recovering here. The point of the test is the sign convention
    // and the magnitude, which are what a caller would get catastrophically wrong.
    geom_true.BeamX_pxl(geom_assumed.GetBeamX_pxl() + 4.0f)
             .BeamY_pxl(geom_assumed.GetBeamY_pxl() - 3.0f);

    const auto profile = SynthesiseProfile(mapping, geom_assumed, geom_true);
    const auto observed = RingRadiiFromProfile(profile, mapping, geom_assumed);
    REQUIRE(observed.size() >= 2);

    const auto offset = BeamCentreOffsetFromProfile(profile, mapping, geom_assumed, observed);
    REQUIRE(offset.has_value());
    CHECK(offset->first == Catch::Approx(4.0).margin(1.5));
    CHECK(offset->second == Catch::Approx(-3.0).margin(1.5));
}