// 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 "../common/CUDAWrapper.h"

#ifdef JFJOCH_USE_CUDA

#include <cmath>
#include <cstdint>
#include <limits>
#include <random>
#include <vector>

#include "../image_analysis/geom_refinement/BeamCenterFFTCPU.h"
#include "../image_analysis/geom_refinement/BeamCenterFFTGPU.h"

namespace {

// An isotropic scattered background about (cx, cy) with a solvent ring, a module gap and a darker
// sextant - the same shape of image the corpus measurement runs on, at a size a test can afford.
std::vector<float> SynthesiseBackground(int64_t w, int64_t h, float cx, float cy) {
    std::vector<float> mean(static_cast<size_t>(w * h), NAN);
    std::mt19937 rng(20260913);
    std::normal_distribution<float> gauss(0.0f, 1.0f);
    for (int64_t y = 0; y < h; y++) {
        if (y >= 500 && y < 520)
            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 - 330.0f) / 40.0f;
            float value = 130.0f * std::exp(-r / 260.0f) + 50.0f * std::exp(-0.5f * t * t) + 6.0f;
            const float phi = std::atan2(static_cast<float>(y) - cy, static_cast<float>(x) - cx);
            if (phi > 0.0f && phi < 1.0f)
                value *= 0.8f;
            mean[static_cast<size_t>(y * w + x)] = value + gauss(rng) * std::sqrt(value / 60.0f);
        }
    }
    return mean;
}

void CompareShortlist(const std::vector<BeamCenterFFTCandidate> &gpu,
                      const std::vector<BeamCenterFFTCandidate> &cpu, const char *what) {
    INFO(what);
    REQUIRE(gpu.size() == cpu.size());
    for (size_t i = 0; i < cpu.size(); i++) {
        INFO("candidate " << i);
        // The positions are what a caller consumes, and they are a grid index: they must be the
        // same index, not a nearby one.
        CHECK(((gpu[i].beam_x_pxl == cpu[i].beam_x_pxl)
               || (std::isnan(gpu[i].beam_x_pxl) && std::isnan(cpu[i].beam_x_pxl))));
        CHECK(((gpu[i].beam_y_pxl == cpu[i].beam_y_pxl)
               || (std::isnan(gpu[i].beam_y_pxl) && std::isnan(cpu[i].beam_y_pxl))));
        CHECK(gpu[i].score == Catch::Approx(cpu[i].score).epsilon(1e-5));
    }
}

// Largest difference of two scored surfaces. Entries that are -infinity in both (a centre the gates
// refuse) are equal; one refused on one side only is counted as a gate mismatch.
//
// `pair_count` and `floor` exclude the centres sitting on the pair-count gate itself. That is not
// special pleading, it is what the score is: at the gate a centre's mirror overlap is a quarter of
// the largest, the masked Pearson there is a cancellation over the fewest pixels it is ever
// computed over, and a 1e-7 difference in the convolutions lands at 1e-4 in r - measured, the worst
// difference between the two paths on this image is at pair fraction 0.250021. Those centres are
// discarded by both paths anyway; what has to agree is the surface the shortlist searches.
struct SurfaceDifference {
    double worst = 0.0;
    size_t gate_mismatches = 0;
};

SurfaceDifference Compare(const std::vector<float> &got, const std::vector<float> &ref,
                          const std::vector<float> *pair_count = nullptr, float floor = 0.0f) {
    SurfaceDifference out;
    for (size_t i = 0; i < ref.size(); i++) {
        if (pair_count && (*pair_count)[i] < floor)
            continue;
        const bool a = std::isfinite(got[i]), b = std::isfinite(ref[i]);
        if (a != b)
            out.gate_mismatches++;
        else if (a)
            out.worst = std::max(out.worst, std::abs(static_cast<double>(got[i]) - ref[i]));
    }
    return out;
}

// The same for a convolution surface, which is not a correlation coefficient: it spans orders of
// magnitude and its small entries are differences of large ones, so the difference is scaled by the
// largest term of the reference rather than taken elementwise.
double ScaledDifference(const std::vector<float> &got, const std::vector<float> &ref) {
    double scale = 0.0, worst = 0.0;
    for (float 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]) - ref[i]));
    return scale > 0.0 ? worst / scale : worst;
}

} // namespace

// The GPU engine exists to make the capture affordable, not to answer differently. It no longer
// shares the masked Pearson and the greedy search with the CPU path - both run on the device, where
// the surfaces already are - so the two paths are compared at every level they still have in
// common, not only at the answer:
//
//   1. the four convolutions, which both engines produce (`PointSurfaces`);
//   2. the scored surface, the device kernel against `BeamCenterPointScore`, which is the reference
//      the CPU path runs - `PointScoreSurface` exists for this and for nothing else;
//   3. the shortlist, the device search against `BeamCenterShortlist2D` on the same surface;
//   4. the whole score, shortlist positions and margins alike.
//
// What this can no longer do is what the first version of it did: prove that a difference could
// only have come from the transform. A difference in (2) or (3) is now a difference between two
// implementations of the same formula, which is why (2) and (3) are compared elementwise and
// exactly rather than through their effect on the answer.
TEST_CASE("BeamCenterFFTGPU_MatchesTheCPUEngine", "[BeamCenter][BeamCenterFFTGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_MatchesTheCPUEngine");
        return;
    }

    const int64_t w = 1100, h = 900;
    const float cx = 397.25f, cy = 511.5f;
    const auto mean = SynthesiseBackground(w, h, cx, cy);

    BeamCenterFFTCPU cpu_engine;
    BeamCenterFFTGPU gpu_engine;
    const BeamCenterFFTSettings settings;

    // The score works on the prepared image, which is internal; the engines below are fed the
    // projection through the same entry point, so all four comparisons see the same preparation.
    const auto cpu = BeamCenterFFTScore(w, h, mean, settings, cpu_engine);
    const auto gpu = BeamCenterFFTScore(w, h, mean, settings, gpu_engine);

    SECTION("the four convolutions") {
        // Fed the raw projection rather than the prepared image: what matters here is that two
        // transforms of the same arrays agree, not what the arrays are.
        std::vector<float> a(mean), m(mean.size(), 1.0f);
        for (size_t i = 0; i < a.size(); i++)
            if (!std::isfinite(a[i])) {
                a[i] = 0.0f;
                m[i] = 0.0f;
            }
        const auto c = cpu_engine.PointSurfaces(a, m, h, w);
        const auto g = gpu_engine.PointSurfaces(a, m, h, w);
        CHECK(ScaledDifference(g.C, c.C) < 1e-6);
        CHECK(ScaledDifference(g.S, c.S) < 1e-6);
        CHECK(ScaledDifference(g.Q, c.Q) < 1e-6);
        CHECK(ScaledDifference(g.D, c.D) < 1e-6);

        // And the two combinations of them, through the one reference implementation: with the same
        // four surfaces in, the formula is the formula.
        const double variance = 1.0;   // only scales the no-evidence floor
        const auto rc = BeamCenterPointScore(c, settings.min_pair_fraction, variance);
        const auto rg = BeamCenterPointScore(g, settings.min_pair_fraction, variance);
        float d_max = 0.0f;
        for (float v : c.D)
            d_max = std::max(d_max, v);
        // Where the shortlist looks. Everywhere else the comparison measures the conditioning of
        // the formula rather than either implementation of it: r is a ratio of two cancellations,
        // and on a weak centre - a small mirror overlap, a denominator near nothing - the
        // transform's 1e-6 lands at 1e-4 in r whichever path computed it (measured: 1.5e-4 at worst
        // over the gated surface, 2.2e-4 if the centres sitting exactly on the pair-count gate are
        // included). None of that reaches a peak, which is the only part of the surface anything
        // reads.
        float r_max = -std::numeric_limits<float>::infinity();
        for (float v : rc)
            if (std::isfinite(v))
                r_max = std::max(r_max, v);
        double peak_worst = 0.0;
        for (size_t i = 0; i < rc.size(); i++)
            if (std::isfinite(rc[i]) && rc[i] > 0.5f * r_max && std::isfinite(rg[i]))
                peak_worst = std::max(peak_worst,
                                      std::abs(static_cast<double>(rg[i]) - rc[i]));
        CHECK(peak_worst < 1e-5);
        // And the gate agrees: over four million centres, the two paths disagree about whether a
        // centre is scored at all only for a handful sitting exactly on the threshold.
        const auto diff = Compare(rg, rc, &c.D, 0.26f * d_max);
        CHECK(diff.gate_mismatches == 0);
        CHECK(Compare(rg, rc).gate_mismatches < 10);

        // The device kernel against that reference, on the device's own surfaces - so this measures
        // the combination and not the transform.
        const auto device_score = gpu_engine.PointScoreSurface(a, m, h, w,
                                                               settings.min_pair_fraction, variance);
        const auto kernel = Compare(device_score, rg);
        CHECK(kernel.gate_mismatches == 0);
        CHECK(kernel.worst < 1e-6);   // the same formula on the same numbers, so this is exact-ish

        // The device search against the host reference, on one surface. The device path is
        // deterministic, so `PointShortlist` sees exactly the surface `PointScoreSurface` returned
        // and this comparison is exact.
        std::vector<float> surface = device_score;
        const auto host_peaks = BeamCenterShortlist2D(surface, 2 * h, 2 * w,
                                                      settings.nms_radius_pxl,
                                                      settings.candidates_point);
        BeamCenterFFTSettings scored_settings = settings;
        const auto device_peaks = gpu_engine.PointShortlist(a, m, h, w, scored_settings, variance);
        REQUIRE(device_peaks.size() == host_peaks.size());
        for (size_t i = 0; i < host_peaks.size(); i++) {
            INFO("peak " << i);
            CHECK(device_peaks[i].beam_x_pxl == host_peaks[i].beam_x_pxl);
            CHECK(device_peaks[i].beam_y_pxl == host_peaks[i].beam_y_pxl);
            CHECK(device_peaks[i].score == host_peaks[i].score);
        }
    }

    SECTION("the shortlists and the margins") {
        REQUIRE(!cpu.point.empty());
        CompareShortlist(gpu.point, cpu.point, "point surface");
        CompareShortlist(gpu.line_x, cpu.line_x, "x line mirror");
        CompareShortlist(gpu.line_y, cpu.line_y, "y line mirror");

        // The margin is the surface's self-diagnostic, and it is a difference of two scores that sit
        // within a fraction of a percent of each other - the quantity most exposed to a change of
        // transform.
        CHECK(gpu.margin_point == Catch::Approx(cpu.margin_point).margin(1e-5));
        CHECK(gpu.margin_line_x == Catch::Approx(cpu.margin_line_x).margin(1e-5));
        CHECK(gpu.margin_line_y == Catch::Approx(cpu.margin_line_y).margin(1e-5));
    }
}

// The same identity test the CPU engine carries, on the device: an exactly centrosymmetric image
// scores exactly 1 at its centre and the capture lands on that grid point. It is the cheapest
// check there is for an indexing or sign error that only the device path has.
TEST_CASE("BeamCenterFFTGPU_RecoversAnExactlySymmetricCentre", "[BeamCenter][BeamCenterFFTGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_RecoversAnExactlySymmetricCentre");
        return;
    }

    // Keyed on the canonical member of the +- pair, so the image is exactly its own mirror about
    // (cx, cy) - the same construction the CPU identity test uses.
    const int64_t w = 220, h = 180;
    const float cx = 83.5f, cy = 61.5f;
    std::vector<float> mean(static_cast<size_t>(w * h));
    for (int64_t y = 0; y < h; y++)
        for (int64_t x = 0; x < w; x++) {
            int64_t dx = std::lround(2.0f * (static_cast<float>(x) - cx));
            int64_t dy = std::lround(2.0f * (static_cast<float>(y) - cy));
            if (dy < 0 || (dy == 0 && dx < 0)) {
                dx = -dx;
                dy = -dy;
            }
            uint64_t k = static_cast<uint64_t>(dx) * 0x9e3779b97f4a7c15ULL
                         + static_cast<uint64_t>(dy) * 0xc2b2ae3d27d4eb4fULL;
            k ^= k >> 31;
            k *= 0xbf58476d1ce4e5b9ULL;
            k ^= k >> 29;
            mean[static_cast<size_t>(y * w + x)] = 100.0f + static_cast<float>(k % 4096u) * 0.25f;
        }

    BeamCenterFFTGPU engine;
    const auto r = BeamCenterFFTScore(w, h, mean, BeamCenterFFTSettings{}, engine);
    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 card is shared, so the engine has to be able to say that an image will not fit before it
// starts allocating - that is what sends the capture to the CPU instead of taking the run down.
TEST_CASE("BeamCenterFFTGPU_SizesItsOwnFootprint", "[BeamCenter][BeamCenterFFTGPU]") {
    if (get_gpu_count() == 0) {
        WARN("No CUDA GPU present. Skipping BeamCenterFFTGPU_SizesItsOwnFootprint");
        return;
    }

    // A 16 Mpixel detector is the largest this runs on, and it has to be a bounded cost.
    CHECK(BeamCenterFFTGPU::DeviceMemoryNeeded(4148, 4362) < (3ull << 30));
    // Bigger images need more.
    CHECK(BeamCenterFFTGPU::DeviceMemoryNeeded(4148, 4362)
          > BeamCenterFFTGPU::DeviceMemoryNeeded(1030, 1064));
    // Nothing on this card can hold a hundred-gigapixel image.
    CHECK(!BeamCenterFFTGPU::FitsInDeviceMemory(300000, 300000));
}

#endif