// 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 <algorithm>
#include <cmath>
#include <cstdint>
#include <cstdlib>
#include <limits>
#include <random>

#include "../image_analysis/geom_refinement/BeamCenterFFT.h"
#include "../image_analysis/geom_refinement/BeamCenterFFTCPU.h"
#include "../image_analysis/geom_refinement/BeamCenterFFTEngine.h"

namespace {

// An isotropic scattered background about (cx, cy): a decaying continuum with a solvent-ring bump,
// the mean of ~60 frames' worth of counting noise on top, NAN outside the detector's usable area.
// A horizontal dead stripe stands in for a module gap.
std::vector<float> Synthesise(int64_t w, int64_t h, float cx, float cy, float shadow_sector) {
    std::vector<float> mean(static_cast<size_t>(w * h), NAN);
    std::mt19937 rng(20260912);
    std::normal_distribution<float> gauss(0.0f, 1.0f);
    for (int64_t y = 0; y < h; y++) {
        if (y >= 300 && y < 320)
            continue;   // module gap
        for (int64_t x = 0; x < w; x++) {
            const float r = std::hypot(static_cast<float>(x) - cx, static_cast<float>(y) - cy);
            const float t = (r - 260.0f) / 30.0f;
            float value = 120.0f * std::exp(-r / 220.0f) + 45.0f * std::exp(-0.5f * t * t) + 5.0f;
            const float phi = std::atan2(static_cast<float>(y) - cy, static_cast<float>(x) - cx);
            if (phi > 0.0f && phi < 1.0f)
                value *= shadow_sector;
            mean[static_cast<size_t>(y * w + x)] = value + gauss(rng) * std::sqrt(value / 60.0f);
        }
    }
    return mean;
}

} // namespace

TEST_CASE("BeamCenterFFT_SyntheticBackground", "[BeamCenter]") {
    // The centre is nowhere near the image middle - the score must not care.
    const int64_t w = 900, h = 700;
    const float cx = 297.25f, cy = 411.5f;

    SECTION("clean isotropic background") {
        const auto mean = Synthesise(w, h, cx, cy, 1.0f);
        const auto r = BeamCenterFFTScore(w, h, mean);

        REQUIRE(!r.point.empty());
        REQUIRE(!r.line_x.empty());
        REQUIRE(!r.line_y.empty());
        // Whole-detector capture: the strongest point-symmetry centre is the true one.
        CHECK(std::hypot(r.point[0].beam_x_pxl - cx, r.point[0].beam_y_pxl - cy) < 2.0f);
        CHECK(std::abs(r.line_x[0].beam_x_pxl - cx) < 2.0f);
        CHECK(std::abs(r.line_y[0].beam_y_pxl - cy) < 2.0f);
        // A clean unimodal surface separates its argmax from the runner-up by a wide margin.
        CHECK(r.margin_point > 0.05f);
    }

    SECTION("a shadowed sextant keeps the true centre in the shortlist") {
        const auto mean = Synthesise(w, h, cx, cy, 0.5f);
        const auto r = BeamCenterFFTScore(w, h, mean);

        REQUIRE(!r.point.empty());
        float best = 1e9f;
        for (const auto &c : r.point)
            best = std::min(best, std::hypot(c.beam_x_pxl - cx, c.beam_y_pxl - cy));
        // The argmax is pulled by the asymmetry (measured ~9 px here - the same shadow bias the
        // background estimator carries); the shortlist must still hold the answer within the
        // 12 px a re-indexing verifier accepts.
        CHECK(best < 12.0f);
    }
}

namespace {

// A value keyed on the distance from (cx, cy) in each coordinate, so the image is EXACTLY its own
// mirror about that centre under point inversion AND under either line mirror - the three symmetries
// the three surfaces measure. The preparation the score applies - clipping, median subtraction,
// clamping, the mean shift - is pointwise and monotone, so it preserves all of them, and the masked
// Pearson at the true centre is then exactly 1. Every other centre must score below it, because the
// score is a correlation coefficient and cannot exceed 1.
float SymmetricValue(int64_t x, int64_t y, float cx, float cy) {
    const int64_t dx = std::llabs(std::lround(2.0f * (static_cast<float>(x) - cx)));
    const int64_t dy = std::llabs(std::lround(2.0f * (static_cast<float>(y) - cy)));
    uint64_t k = static_cast<uint64_t>(dx) * 0x9e3779b97f4a7c15ULL
                 + static_cast<uint64_t>(dy) * 0xc2b2ae3d27d4eb4fULL;
    k ^= k >> 31;
    k *= 0xbf58476d1ce4e5b9ULL;
    k ^= k >> 29;
    return 100.0f + static_cast<float>(k % 4096u) * 0.25f;
}

// The symmetric pattern in a rectangle at (x0, y0), NAN everywhere else, so the same pattern can be
// placed anywhere on the detector. The true centre is (x0 + cx, y0 + cy).
std::vector<float> SymmetricPatch(int64_t w, int64_t h, int64_t x0, int64_t y0, int64_t rw,
                                  int64_t rh, float cx, float cy) {
    std::vector<float> mean(static_cast<size_t>(w * h), NAN);
    for (int64_t y = 0; y < rh; y++)
        for (int64_t x = 0; x < rw; x++)
            mean[static_cast<size_t>((y0 + y) * w + x0 + x)] = SymmetricValue(x, y, cx, cy);
    return mean;
}

// The four convolutions straight from their definitions, for an image small enough to afford it:
// C(t) = sum_x a(x) a(t - x), S = a (*) m, Q = a^2 (*) m, D = m (*) m.
BeamCenterConvSurfaces2D BruteForcePoint(const std::vector<float> &a, const std::vector<float> &m,
                                         int64_t h, int64_t w) {
    BeamCenterConvSurfaces2D out;
    const size_t n = static_cast<size_t>(4 * h * w);
    out.C.assign(n, 0.0f);
    out.S.assign(n, 0.0f);
    out.Q.assign(n, 0.0f);
    out.D.assign(n, 0.0f);
    for (int64_t ty = 0; ty < 2 * h; ty++)
        for (int64_t tx = 0; tx < 2 * w; tx++) {
            double c = 0, s = 0, q = 0, d = 0;
            for (int64_t y = std::max<int64_t>(0, ty - h + 1); y <= std::min(h - 1, ty); y++)
                for (int64_t x = std::max<int64_t>(0, tx - w + 1); x <= std::min(w - 1, tx); x++) {
                    const size_t i = static_cast<size_t>(y * w + x);
                    const size_t j = static_cast<size_t>((ty - y) * w + tx - x);
                    c += static_cast<double>(a[i]) * a[j];
                    s += static_cast<double>(a[i]) * m[j];
                    q += static_cast<double>(a[i]) * a[i] * m[j];
                    d += static_cast<double>(m[i]) * m[j];
                }
            const size_t t = static_cast<size_t>(ty * 2 * w + tx);
            out.C[t] = static_cast<float>(c);
            out.S[t] = static_cast<float>(s);
            out.Q[t] = static_cast<float>(q);
            out.D[t] = static_cast<float>(d);
        }
    return out;
}

// The same for the 1D mirror of the row coordinate: each column autoconvolved, summed over columns.
BeamCenterConvSurfaces1D BruteForceLine(const std::vector<float> &a, const std::vector<float> &m,
                                        int64_t h, int64_t w) {
    BeamCenterConvSurfaces1D out;
    out.C.assign(static_cast<size_t>(2 * h), 0.0);
    out.S.assign(static_cast<size_t>(2 * h), 0.0);
    out.Q.assign(static_cast<size_t>(2 * h), 0.0);
    out.D.assign(static_cast<size_t>(2 * h), 0.0);
    for (int64_t t = 0; t < 2 * h; t++)
        for (int64_t x = 0; x < w; x++)
            for (int64_t y = std::max<int64_t>(0, t - h + 1); y <= std::min(h - 1, t); y++) {
                const size_t i = static_cast<size_t>(y * w + x);
                const size_t j = static_cast<size_t>((t - y) * w + x);
                out.C[static_cast<size_t>(t)] += static_cast<double>(a[i]) * a[j];
                out.S[static_cast<size_t>(t)] += static_cast<double>(a[i]) * m[j];
                out.Q[static_cast<size_t>(t)] += static_cast<double>(a[i]) * a[i] * m[j];
                out.D[static_cast<size_t>(t)] += static_cast<double>(m[i]) * m[j];
            }
    return out;
}

// Largest relative difference of two surfaces, scaled by the largest term of the reference - a
// convolution surface spans orders of magnitude and its small entries are differences of large
// ones, so an elementwise relative test measures roundoff and nothing else.
template <class T> double MaxScaledDifference(const std::vector<T> &got, const std::vector<T> &ref) {
    double scale = 0.0, worst = 0.0;
    for (T v : ref)
        scale = std::max(scale, std::abs(static_cast<double>(v)));
    for (size_t i = 0; i < ref.size(); i++)
        worst = std::max(worst, std::abs(static_cast<double>(got[i]) - static_cast<double>(ref[i])));
    return scale > 0.0 ? worst / scale : worst;
}

} // namespace

// The identity the whole component rests on: the centrosymmetry score about c is the
// autoconvolution at 2c, so one transform set scores every centre. On an exactly centrosymmetric
// image the score at the true centre is exactly 1 and nowhere else can reach it - at integer and
// at half-pixel centres alike, because the surface is sampled on the half-pixel grid.
TEST_CASE("BeamCenterFFT_RecoversAnExactlySymmetricCentre", "[BeamCenter]") {
    const int64_t w = 220, h = 180;

    const float cx = GENERATE(83.0f, 83.5f);
    const float cy = GENERATE(61.0f, 61.5f);

    const auto mean = SymmetricPatch(w, h, 0, 0, w, h, cx, cy);
    const auto r = BeamCenterFFTScore(w, h, mean);

    REQUIRE(!r.point.empty());
    CHECK(r.point[0].beam_x_pxl == cx);
    CHECK(r.point[0].beam_y_pxl == cy);
    CHECK(r.point[0].score == Catch::Approx(1.0f).margin(1e-4));
    // The two line mirrors see the same symmetry one coordinate at a time.
    REQUIRE(!r.line_x.empty());
    REQUIRE(!r.line_y.empty());
    CHECK(r.line_x[0].beam_x_pxl == cx);
    CHECK(r.line_y[0].beam_y_pxl == cy);
}

// Move the image and the answer moves with it, exactly. The score is a convolution, so this is a
// property and not a tolerance: an off-by-one in the padding, the crop or the t = 2c indexing shows
// up here as a fixed offset, which is the class of error a synthetic image at one position cannot
// see.
TEST_CASE("BeamCenterFFT_IsShiftEquivariant", "[BeamCenter]") {
    const int64_t w = 260, h = 200;
    const int64_t rw = 150, rh = 120;
    const float cx = 71.5f, cy = 55.0f;

    const auto at_origin = BeamCenterFFTScore(w, h, SymmetricPatch(w, h, 0, 0, rw, rh, cx, cy));
    REQUIRE(!at_origin.point.empty());

    const int64_t tx = 37, ty = 43;
    const auto shifted = BeamCenterFFTScore(w, h, SymmetricPatch(w, h, tx, ty, rw, rh, cx, cy));
    REQUIRE(!shifted.point.empty());

    CHECK(shifted.point[0].beam_x_pxl == at_origin.point[0].beam_x_pxl + static_cast<float>(tx));
    CHECK(shifted.point[0].beam_y_pxl == at_origin.point[0].beam_y_pxl + static_cast<float>(ty));
    CHECK(shifted.line_x[0].beam_x_pxl == at_origin.line_x[0].beam_x_pxl + static_cast<float>(tx));
    CHECK(shifted.line_y[0].beam_y_pxl == at_origin.line_y[0].beam_y_pxl + static_cast<float>(ty));
}

// The four convolutions the score is built from, against their definitions. The masked Pearson is a
// cancellation of large near-equal terms and the shortlist margins ride on differences of ~0.3 %,
// so what matters is not that the transform is approximately right but by how much it is wrong.
TEST_CASE("BeamCenterFFT_ConvolutionsAgreeWithBruteForce", "[BeamCenter]") {
    const int64_t w = 23, h = 17;
    std::vector<float> a(static_cast<size_t>(w * h)), m(static_cast<size_t>(w * h));
    std::mt19937 rng(20260913);
    std::uniform_real_distribution<float> uniform(-1.0f, 1.0f);
    for (size_t i = 0; i < a.size(); i++) {
        // A quarter of the pixels masked, in a pattern with no symmetry of its own.
        m[i] = (i * 7919u) % 4u == 0u ? 0.0f : 1.0f;
        a[i] = m[i] * (30.0f + 10.0f * uniform(rng));
    }

    BeamCenterFFTCPU engine;
    const auto got = engine.PointSurfaces(a, m, h, w);
    const auto ref = BruteForcePoint(a, m, h, w);
    CHECK(MaxScaledDifference(got.C, ref.C) < 1e-6);
    CHECK(MaxScaledDifference(got.S, ref.S) < 1e-6);
    CHECK(MaxScaledDifference(got.Q, ref.Q) < 1e-6);
    CHECK(MaxScaledDifference(got.D, ref.D) < 1e-6);

    // And the score built from them, through the reference implementation both engines are held
    // to - compared where the shortlist looks. r is a ratio of two cancellations, so on a weak
    // centre the transform's 1e-6 lands at 1e-4 in r; at a peak it does not, and a peak is the only
    // part of the surface anything reads. The margin here is what a float32 transform costs against
    // an exact double reference on an image this small, where an overlap is a few hundred pixels.
    const auto score = BeamCenterPointScore(got, 0.25f, 1.0);
    const auto score_ref = BeamCenterPointScore(ref, 0.25f, 1.0);
    float r_max = -std::numeric_limits<float>::infinity();
    for (float v : score_ref)
        if (std::isfinite(v))
            r_max = std::max(r_max, v);
    for (size_t i = 0; i < score.size(); i++) {
        if (!std::isfinite(score_ref[i]) || score_ref[i] <= 0.5f * r_max)
            continue;
        INFO("surface element " << i);
        REQUIRE(std::isfinite(score[i]));
        REQUIRE(score[i] == Catch::Approx(score_ref[i]).margin(5e-5));
    }

    const auto got_line = engine.LineSurfaces(a, m, h, w, BeamCenterMirror::Rows);
    const auto ref_line = BruteForceLine(a, m, h, w);
    CHECK(MaxScaledDifference(got_line.C, ref_line.C) < 1e-6);
    CHECK(MaxScaledDifference(got_line.S, ref_line.S) < 1e-6);
    CHECK(MaxScaledDifference(got_line.Q, ref_line.Q) < 1e-6);
    CHECK(MaxScaledDifference(got_line.D, ref_line.D) < 1e-6);
}

// A mirror overlap with no variance in it - a region the preparation has flattened to a constant -
// makes the score 0/0, and in floating point that mints an r of several hundred out of pure
// roundoff. The score is a correlation coefficient: no candidate may ever come back above 1.
TEST_CASE("BeamCenterFFT_ANoVarianceOverlapScoresNothing", "[BeamCenter]") {
    const int64_t w = 400, h = 320;
    const int64_t pw = 240, ph = 200;
    const float cx = 0.5f * static_cast<float>(pw - 1), cy = 0.5f * static_cast<float>(ph - 1);

    // A flat detector-wide background - which the preparation clamps to a constant - with the
    // textured, exactly symmetric patch filling one corner of it. Every centre whose overlap lies
    // inside the flat part carries no evidence at all.
    std::vector<float> mean(static_cast<size_t>(w * h), 100.0f);
    for (int64_t y = 0; y < ph; y++)
        for (int64_t x = 0; x < pw; x++)
            mean[static_cast<size_t>(y * w + x)] = SymmetricValue(x, y, cx, cy);

    const auto r = BeamCenterFFTScore(w, h, mean);
    REQUIRE(!r.point.empty());
    for (const auto &c : r.point)
        CHECK(c.score <= 1.0f + 1e-3f);
    CHECK(r.point[0].beam_x_pxl == cx);
    CHECK(r.point[0].beam_y_pxl == cy);
}

// The transform length. A smooth length is what makes the whole-detector transform affordable: the
// naive next power of two for a 16 Mpixel detector costs 3x the time and, on cuFFT, 2.5 GB of plan
// memory against a few hundred MB.
TEST_CASE("BeamCenterFFT_PadsToASmoothLength", "[BeamCenter]") {
    CHECK(BeamCenterFFTPadSize(1) == 1);
    CHECK(BeamCenterFFTPadSize(1024) == 1024);
    CHECK(BeamCenterFFTPadSize(1025) == 1029);   // 3 * 7^3
    CHECK(BeamCenterFFTPadSize(2027) == 2048);
    CHECK(BeamCenterFFTPadSize(4931) == 5000);   // 2^3 * 5^4

    for (int64_t n = 2000; n < 20000; n += 7) {
        const int64_t pad = BeamCenterFFTPadSize(n);
        REQUIRE(pad >= n);
        int64_t v = pad;
        for (int p : {2, 3, 5, 7})
            while (v % p == 0)
                v /= p;
        REQUIRE(v == 1);
        // And it is always close: the padding never costs a whole extra octave the way a
        // power-of-two length can.
        REQUIRE(pad < n + n / 20);
    }
}