// 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 "../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());
}