// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

#include <cmath>

#include "../image_analysis/bragg_integration/BraggStencil.h"

namespace {

BraggStencilParams Params(float k_sigma, float bw_sigma = 0.002f) {
    BraggStencilParams p;
    p.beam_x = 400.0f;
    p.beam_y = 400.0f;
    p.r2 = 6.0f;
    p.r3 = 10.0f;
    p.bw_sigma = bw_sigma;
    p.k_sigma = k_sigma;
    p.max_grow = 2.0f * p.r3;
    return p;
}

}   // namespace

// The whole change rests on this: with no elongation asked for, the three squared distances the
// integrator tests against must be the SAME BITS as the plain circular distance used before, so
// that every pixel is classified exactly as it was, not merely nearly.
TEST_CASE("BraggStencil_ZeroElongationIsExactlyCircular", "[Integration]") {
    const BraggStencilParams p = Params(0.0f);   // a bandwidth, but k_sigma = 0

    for (float py = 0.0f; py < 800.0f; py += 37.0f)
        for (float px = 0.0f; px < 800.0f; px += 41.0f) {
            const BraggStencil s = MakeBraggStencil(px, py, p);
            REQUIRE(s.q_in == 0.0f);
            REQUIRE(s.q_out == 0.0f);
            for (int dy = -12; dy <= 12; ++dy)
                for (int dx = -12; dx <= 12; ++dx) {
                    const auto d = BraggStencilDistances(s, static_cast<float>(dx), static_cast<float>(dy));
                    const float circular = static_cast<float>(dx) * dx + static_cast<float>(dy) * dy;
                    REQUIRE(d.signal == circular);
                    REQUIRE(d.inner == circular);
                    REQUIRE(d.outer == circular);
                }
        }
}

// A monochromatic beam has no streak, so nothing is elongated whatever k_sigma says - which is what
// makes the feature inert on every monochromatic dataset rather than merely small.
TEST_CASE("BraggStencil_MonochromaticIsInert", "[Integration]") {
    const BraggStencilParams p = Params(4.0f, 0.0f);
    for (float py = 0.0f; py < 800.0f; py += 53.0f)
        for (float px = 0.0f; px < 800.0f; px += 59.0f) {
            const BraggStencil s = MakeBraggStencil(px, py, p);
            REQUIRE(s.grow == 0.0f);
            REQUIRE(s.q_in == 0.0f);
            REQUIRE(s.q_out == 0.0f);
        }
}

// The elongated region really is the ellipse it claims: radial semi-axis r + grow, tangential r.
TEST_CASE("BraggStencil_ElongatedSemiAxes", "[Integration]") {
    const BraggStencilParams p = Params(3.0f);

    for (float py = 120.0f; py < 800.0f; py += 91.0f)
        for (float px = 120.0f; px < 800.0f; px += 97.0f) {
            const BraggStencil s = MakeBraggStencil(px, py, p);
            const float grow = s.grow;
            REQUIRE(grow > 0.0f);
            REQUIRE(grow <= p.max_grow);
            REQUIRE(s.grow == BraggStencilGrow_px(s.r0, p));   // the kernel table indexes on this

            // On the radial axis the inner boundary sits at r2 + grow, the outer at r3 + grow.
            const auto rad_in = BraggStencilDistances(s, (p.r2 + grow) * s.ux, (p.r2 + grow) * s.uy);
            const auto rad_out = BraggStencilDistances(s, (p.r3 + grow) * s.ux, (p.r3 + grow) * s.uy);
            CHECK_THAT(rad_in.inner, Catch::Matchers::WithinRel(p.r2 * p.r2, 1e-4f));
            CHECK_THAT(rad_out.outer, Catch::Matchers::WithinRel(p.r3 * p.r3, 1e-4f));

            // Across it, at the untouched tangential half-widths r2 and r3. Testing on the exact
            // tangential axis would be a tautology - rad is 0 there, so q never enters - so the
            // point that matters is that the SAME offset is inside the region radially and outside
            // it tangentially. That is the anisotropy, and it fails if q is built from the wrong
            // radius or from a constant.
            const float probe = p.r2 + 0.5f * grow;
            const auto radial_probe = BraggStencilDistances(s, probe * s.ux, probe * s.uy);
            const auto tangent_probe = BraggStencilDistances(s, -probe * s.uy, probe * s.ux);
            CHECK(radial_probe.inner < p.r2 * p.r2);      // still signal, the ring starts further out
            CHECK(tangent_probe.inner > p.r2 * p.r2);     // already background across the streak
            const auto tan_in = BraggStencilDistances(s, -p.r2 * s.uy, p.r2 * s.ux);
            const auto tan_out = BraggStencilDistances(s, -p.r3 * s.uy, p.r3 * s.ux);
            CHECK_THAT(tan_in.inner, Catch::Matchers::WithinRel(p.r2 * p.r2, 1e-4f));
            CHECK_THAT(tan_out.outer, Catch::Matchers::WithinRel(p.r3 * p.r3, 1e-4f));
        }
}

// The bounding boxes the engines scan must contain the regions they classify - a box one pixel too
// small silently drops background pixels on one side of every reflection, which no parity test
// between two engines making the same mistake would catch.
TEST_CASE("BraggStencil_BoundingBoxContainsRegion", "[Integration]") {
    const BraggStencilParams p = Params(3.0f);
    const float r2_sq = p.r2 * p.r2, r3_sq = p.r3 * p.r3;

    for (float py = 0.0f; py < 800.0f; py += 53.0f)
        for (float px = 0.0f; px < 800.0f; px += 59.0f) {
            const BraggStencil s = MakeBraggStencil(px, py, p);
            const int span = static_cast<int>(std::ceil(p.r3 + p.max_grow)) + 4;
            for (int dy = -span; dy <= span; ++dy)
                for (int dx = -span; dx <= span; ++dx) {
                    const auto d = BraggStencilDistances(s, static_cast<float>(dx), static_cast<float>(dy));
                    const float ax = std::fabs(static_cast<float>(dx)), ay = std::fabs(static_cast<float>(dy));
                    // No slack: the offsets are integers from an exactly centred stencil, so the
                    // extents bound them outright. A tolerance of a pixel here would accept a box
                    // one pixel too small, which is the error this exists to catch.
                    if (d.inner < r2_sq) {
                        INFO("inner region outside its box at " << dx << "," << dy);
                        REQUIRE(ax <= s.ex_in);
                        REQUIRE(ay <= s.ey_in);
                    }
                    if (d.inner >= r2_sq && d.outer < r3_sq) {
                        INFO("ring outside its box at " << dx << "," << dy);
                        REQUIRE(ax <= s.ex_out);
                        REQUIRE(ay <= s.ey_out);
                    }
                }
        }
}

// The growth is capped, so a mis-declared bandwidth cannot run away with the bounding box.
TEST_CASE("BraggStencil_GrowthIsCapped", "[Integration]") {
    BraggStencilParams p = Params(3.0f, 0.5f);   // an absurdly declared bandwidth
    for (float r0 = 0.0f; r0 < 4000.0f; r0 += 17.0f)
        REQUIRE(BraggStencilGrow_px(r0, p) <= p.max_grow);

    const BraggStencil s = MakeBraggStencil(4000.0f, 4000.0f, p);
    REQUIRE(s.ex_out <= p.r3 + p.max_grow + 1e-3f);
    REQUIRE(s.ey_out <= p.r3 + p.max_grow + 1e-3f);
}

// The kernel table is indexed by the growth rounded to whole pixels, so the table has to have a row
// for every index any reflection on the detector can produce. An off-by-one here is an out-of-range
// read of k_diff - on the GPU, a device-side one.
TEST_CASE("BraggStencil_KernelIndexInRange", "[Integration][portable]") {
    for (const float k : {0.0f, 0.4f, 1.0f, 2.5f, 3.0f, 6.0f}) {
        const BraggStencilParams p = Params(k);
        const float r_max = std::hypot(800.0f - p.beam_x, 800.0f - p.beam_y);
        const int n_kern = static_cast<int>(std::lround(BraggStencilGrow_px(r_max, p))) + 1;
        REQUIRE(n_kern >= 1);
        for (float py = 0.0f; py <= 800.0f; py += 13.0f)
            for (float px = 0.0f; px <= 800.0f; px += 17.0f) {
                const BraggStencil s = MakeBraggStencil(px, py, p);
                const int idx = BraggStencilKernelIndex(s, n_kern);
                INFO("k " << k << " at " << px << "," << py << " grow " << s.grow);
                REQUIRE(idx >= 0);
                REQUIRE(idx < n_kern);
                // The clamp must never be what saves it: the table is sized so the row exists.
                REQUIRE(static_cast<int>(std::lround(s.grow)) == idx);
            }
    }
}

namespace {

// A footprint growing linearly from sigma 1 px at the beam to `edge` px at 800 px, radial and tangential
// alike unless told otherwise.
BraggStencilParams FootprintParams(float edge_rad, float edge_tan) {
    BraggStencilParams p = Params(0.0f, 0.0f);
    p.r1 = 4.0f;
    p.fp_n = 8;
    p.fp_bin_px = 100.0f;
    for (int i = 0; i < p.fp_n; ++i) {
        const float t = (i + 0.5f) / p.fp_n;
        p.fp_sigma_rad[i] = 1.0f + t * (edge_rad - 1.0f);
        p.fp_sigma_tan[i] = 1.0f + t * (edge_tan - 1.0f);
    }
    return p;
}

}   // namespace

// Where the footprint fits the r1 disk (3 sigma <= r1) the stencil is the circular one, bit for bit:
// compact spots integrate exactly as without a footprint.
TEST_CASE("BraggStencil_FootprintInsideDiskIsInert", "[Integration]") {
    const BraggStencilParams p = FootprintParams(1.3f, 1.3f);   // 3 sigma < 4 everywhere
    const BraggStencilParams none = Params(0.0f, 0.0f);
    for (float py = 0.0f; py < 800.0f; py += 37.0f)
        for (float px = 0.0f; px < 800.0f; px += 41.0f) {
            const BraggStencil s = MakeBraggStencil(px, py, p);
            const BraggStencil n = MakeBraggStencil(px, py, none);
            REQUIRE(s.fp_s2r == 0.0f);
            REQUIRE(s.grow == 0.0f);
            REQUIRE(s.grow_tan == 0.0f);
            for (int dy = -12; dy <= 12; ++dy)
                for (int dx = -12; dx <= 12; ++dx) {
                    const auto d = BraggStencilDistances(s, static_cast<float>(dx), static_cast<float>(dy));
                    const auto e = BraggStencilDistances(n, static_cast<float>(dx), static_cast<float>(dy));
                    REQUIRE(d.inner == e.inner);
                    REQUIRE(d.outer == e.outer);
                }
        }
}

// A spot wider than the disk pushes the ring to 3 sigma along each axis separately: a pixel at 3 sigma
// along the radius (or across it) is no longer background, one just beyond the grown ring's inner edge
// is.
TEST_CASE("BraggStencil_FootprintGrowsRingAlongEachAxis", "[Integration]") {
    const BraggStencilParams p = FootprintParams(3.0f, 5.0f);
    const float px = 400.0f + 700.0f, py = 400.0f;   // on +x, 700 px out: radial = x, tangential = y
    const BraggStencil s = MakeBraggStencil(px, py, p);
    float sr, st;
    BraggFootprintAt(700.0f, p, sr, st);
    REQUIRE(s.fp_s2r == sr * sr);
    REQUIRE(s.fp_s2t == st * st);
    REQUIRE_THAT(s.grow, Catch::Matchers::WithinAbs(3.0f * sr - p.r2, 1e-4));
    REQUIRE_THAT(s.grow_tan, Catch::Matchers::WithinAbs(3.0f * st - p.r2, 1e-4));
    const float ar = p.r2 + s.grow, at = p.r2 + s.grow_tan;
    // Just inside the inner ellipse along each axis: not background.
    REQUIRE(BraggStencilDistances(s, 0.98f * ar, 0.0f).inner < p.r2 * p.r2);
    REQUIRE(BraggStencilDistances(s, 0.0f, 0.98f * at).inner < p.r2 * p.r2);
    // Just outside: background ring.
    REQUIRE(BraggStencilDistances(s, 1.02f * ar, 0.0f).inner >= p.r2 * p.r2);
    REQUIRE(BraggStencilDistances(s, 0.0f, 1.02f * at).inner >= p.r2 * p.r2);
    // The bounding box holds the outer ellipse.
    REQUIRE(s.ex_out >= p.r3 + s.grow - 1e-3f);
    REQUIRE(s.ey_out >= p.r3 + s.grow_tan - 1e-3f);
    REQUIRE(BraggStencilMaxGrow_px(1000.0f, p) >= std::max(s.grow, s.grow_tan));
}