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

// Equivalence of the stored (stacked) and the assembled layout of a panel detector.
//
// One synthetic detector, recorded twice: as an assembled image (panels where they are on the plane, gaps
// masked) and as the stacked array a panel detector stores (panels one under the other, in four
// orientations, no gaps), with the DetectorModules that place the stacked panels on the plane. Both carry
// the same value in every physical pixel, so anything computed on the plane must come out the same.
//
// [PanelLayout] cases check what is already layout-correct. The [PanelLayoutTodo] cases are hidden: they
// state the equivalence each later stage has to deliver, and fail until it does. A stage that converts a
// consumer retags its case [PanelLayout]:
//   stage 1 - dense maps: PanelLayout_AzintMapping, PanelLayout_ResolutionMask
//   stage 2 - sparse/Bragg: PanelLayout_BraggIntegration
//   stage 3 - per-panel spot finding: PanelLayout_SpotsNeverCrossPanels; the raster algorithms
//             (ShadowFinder, BeamCenterFFT, BeamCenterFromBackground, DiagnosticOutput) add their case here,
//             comparing the result on the stacked layout with the result on the assembled image.

#include <catch2/catch_all.hpp>

#include <cmath>
#include <map>
#include <vector>

#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/CUDAWrapper.h"
#include "../common/BraggIntegrationSettings.h"
#include "../common/DetectorModules.h"
#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/PixelMask.h"
#include "../common/Reflection.h"
#include "../image_analysis/bragg_integration/BraggIntegrationEngineCPU.h"
#ifdef JFJOCH_USE_CUDA
#include "../image_analysis/bragg_integration/BraggIntegrationEngineGPU.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
#endif
#include "../image_analysis/image_preprocessing/ImagePreprocessorBuffer.h"
#include "../image_analysis/spot_finding/ImageSpotFinderCPU.h"
#include "../image_analysis/spot_finding/SpotUtils.h"

namespace {

// One panel as stored, and the plane the four of them are placed on.
constexpr int64_t PW = 96, PH = 64, NPANELS = 4;
constexpr int64_t PLANE_W = 220, PLANE_H = 200;

// Panel p is stored at rows [p * PH, (p + 1) * PH) and placed 2 x 2 on the plane with gaps of 8 and 16 px.
DetectorModules StackedPanels() {
    auto panel = [](int64_t p, double ox, double oy, double fx, double fy, double sx, double sy) {
        DetectorModule m;
        m.x0 = 0;
        m.y0 = p * PH;
        m.width = PW;
        m.height = PH;
        m.ox = ox; m.oy = oy;
        m.fx = fx; m.fy = fy;
        m.sx = sx; m.sy = sy;
        return m;
    };
    return DetectorModules({
        panel(0, 10, 10, 1, 0, 0, 1),        // as stored:     x  10..105, y  10..73
        panel(1, 209, 73, -1, 0, 0, -1),     // turned 180:    x 114..209, y  10..73
        panel(2, 73, 90, 0, 1, -1, 0),       // turned 90:     x  10..73,  y  90..185
        panel(3, 120, 185, 0, -1, 1, 0)      // turned -90:    x 120..183, y  90..185
    });
}

struct Spot {
    double x, y;   // on the plane
};

// Spots well inside their panels: more than the spot finder's 15 px window from every panel edge.
const std::vector<Spot> INTERIOR_SPOTS = {{50.3, 40.7}, {80.5, 45.3}, {160.6, 30.2}, {40.4, 140.8}, {150.2, 130.5}};

DiffractionExperiment Experiment(int64_t width, int64_t height, const DetectorModules &modules) {
    DiffractionExperiment x(DetDECTRIS(width, height, "PanelLayout", ""));
    // Beam in the gap cross; at 30 mm the plane reaches about 2 A in its corners.
    x.DetectorDistance_mm(30.0f).IncidentEnergy_keV(WVL_1A_IN_KEV).BeamX_pxl(110.0f).BeamY_pxl(82.0f);
    x.ImportBraggIntegrationSettings(BraggIntegrationSettings());
    x.Modules(modules);
    return x;
}

// The same detector, assembled and stacked.
struct Scene {
    DetectorModules modules = StackedPanels();
    DiffractionExperiment assembled = Experiment(PLANE_W, PLANE_H, {});
    DiffractionExperiment stacked = Experiment(PW, NPANELS * PH, modules);
    ImagePreprocessorBuffer assembled_image{PLANE_W * PLANE_H};
    ImagePreprocessorBuffer stacked_image{PW * NPANELS * PH};
    PixelMask assembled_mask;
    PixelMask stacked_mask{PW, NPANELS * PH};

    // Assembled pixel at the plane position of a stacked pixel (the placements are whole-pixel).
    static size_t AssembledIndex(double px, double py) {
        return static_cast<size_t>(std::lround(py)) * PLANE_W + static_cast<size_t>(std::lround(px));
    }

    explicit Scene(const std::vector<Spot> &spots) {
        auto value = [&](double px, double py) {
            double v = 2.0;
            for (const auto &s: spots) {
                const double dx = px - s.x, dy = py - s.y;
                v += 800.0 * std::exp(-(dx * dx + dy * dy) / (2.0 * 1.3 * 1.3));
            }
            return static_cast<int32_t>(std::lround(v));
        };
        // Gaps are what the preprocessor makes of a masked pixel.
        for (size_t i = 0; i < assembled_image.size(); i++)
            assembled_image[i] = INT32_MIN;
        std::vector<uint32_t> gaps(PLANE_W * PLANE_H, 1u << PixelMask::ModuleGapPixelBit);
        modules.ForEachPixel(PW, NPANELS * PH, [&](int64_t, size_t i, double px, double py) {
            stacked_image[i] = value(px, py);
            assembled_image[AssembledIndex(px, py)] = value(px, py);
            gaps[AssembledIndex(px, py)] = 0;
        });
        assembled_mask = PixelMask(gaps);
    }
};

std::vector<SpotToSave> FindSpots(const DiffractionExperiment &x, const ImagePreprocessorBuffer &image,
                                  float snr = 0.0f) {
    SpotFindingSettings settings;
    settings.signal_to_noise_threshold = snr;
    settings.photon_count_threshold = 20;
    settings.min_pix_per_spot = 1;
    settings.high_res_gap_Q_recipA = std::nullopt;
    ImageSpotFinderCPU finder(static_cast<int32_t>(x.GetXPixelsNum()), static_cast<int32_t>(x.GetYPixelsNum()));
    finder.SetResolutionMask(std::vector<bool>(x.GetPixelsNum(), false));
    DataMessage msg{};
    msg.number = 0;
    SpotAnalyze(x, settings, finder.Run(image, settings), msg);
    return msg.spots;
}

// Every spot of `a` has one in `b` at the same plane position, and the counts agree.
void CheckSameSpots(const std::vector<SpotToSave> &a, const std::vector<SpotToSave> &b) {
    REQUIRE(a.size() == b.size());
    for (const auto &s: a) {
        size_t matched = 0;
        for (const auto &t: b)
            if (std::hypot(s.x - t.x, s.y - t.y) < 1e-3) {
                CHECK(s.intensity == t.intensity);
                CHECK(s.d_A == Catch::Approx(t.d_A));
                matched++;
            }
        CHECK(matched == 1);
    }
}

// Integrate predictions placed on the plane: on the assembled image directly, on the stacked one through
// the panel table - PlaneToPanel in, ImageToPlane of the same panel out. Keyed by h.
std::map<int, Reflection> Integrate(const DiffractionExperiment &x, const PixelMask &mask,
                                    const ImagePreprocessorBuffer &image, const std::vector<Spot> &spots,
                                    bool gpu = false) {
    const auto &modules = x.GetModules();
    const auto geom = x.GetDiffractionGeometry();
    std::vector<Reflection> predicted;
    std::map<int, int64_t> panel_of;
    for (size_t i = 0; i < spots.size(); i++) {
        const auto p = modules.PlaneToPanel(spots[i].x, spots[i].y);
        REQUIRE(p.has_value());
        Reflection r{};
        r.h = static_cast<int>(i) + 1; r.k = 0; r.l = 1;
        r.predicted_x = static_cast<float>(p->x);
        r.predicted_y = static_cast<float>(p->y);
        r.d = geom.PxlToRes(static_cast<float>(spots[i].x), static_cast<float>(spots[i].y));
        r.prescaling_corr = 1.0f;
        r.partiality = 1.0f;
        predicted.push_back(r);
        panel_of[r.h] = p->panel;
    }
    std::vector<Reflection> integrated;
#ifdef JFJOCH_USE_CUDA
    if (gpu) {
        auto stream = std::make_shared<CudaStream>();
        ImagePreprocessorBufferGPU gpu_image(image.size());
        for (size_t i = 0; i < image.size(); i++)
            gpu_image[i] = image[i];
        REQUIRE(cudaMemcpyAsync(gpu_image.getGPUBuffer(), gpu_image.getBuffer().data(), image.size() * sizeof(int32_t),
                                cudaMemcpyHostToDevice, *stream) == cudaSuccess);
        BraggIntegrationEngineGPU engine(x, stream, mask);
        integrated = engine.Run(gpu_image, predicted, predicted.size(), 0);
    }
#endif
    if (!gpu) {
        BraggIntegrationEngineCPU engine(x, mask);
        integrated = engine.Run(image, predicted, predicted.size(), 0);
    }
    std::map<int, Reflection> out;
    for (auto r: integrated) {
        const int64_t panel = panel_of.at(r.h);
        const auto [px, py] = modules.ImageToPlane(panel, r.predicted_x, r.predicted_y);
        r.predicted_x = static_cast<float>(px);
        r.predicted_y = static_cast<float>(py);
        const auto [ox, oy] = modules.ImageToPlane(panel, r.observed_x, r.observed_y);
        r.observed_x = static_cast<float>(ox);
        r.observed_y = static_cast<float>(oy);
        out[r.h] = r;
    }
    return out;
}

}  // namespace

// The panel table itself: every stored pixel is visited once and lands on a free plane pixel, plane ->
// panel inverts it, a gap is on no panel, and the plane bounds hold every panel.
TEST_CASE("PanelLayout_PanelTable", "[PanelLayout]") {
    const DetectorModules modules = StackedPanels();
    const PlaneBox box = modules.PlaneBounds(PW, NPANELS * PH);
    CHECK(box.x0 == 9.5);
    CHECK(box.y0 == 9.5);
    CHECK(box.x1 == 209.5);
    CHECK(box.y1 == 185.5);

    std::vector<int> visits(PW * NPANELS * PH, 0), on_plane(PLANE_W * PLANE_H, 0);
    modules.ForEachPixel(PW, NPANELS * PH, [&](int64_t panel, size_t i, double px, double py) {
        visits[i]++;
        on_plane[Scene::AssembledIndex(px, py)]++;
        CHECK(panel == static_cast<int64_t>(i / (PW * PH)));
        const auto p = modules.PlaneToPanel(px, py);
        REQUIRE(p.has_value());
        CHECK(p->panel == panel);
        CHECK(static_cast<size_t>(std::lround(p->y)) * PW + static_cast<size_t>(std::lround(p->x)) == i);
        const auto [qx, qy] = modules.PlaneToImage(panel, px, py);
        CHECK(qx == Catch::Approx(p->x));
        CHECK(qy == Catch::Approx(p->y));
    });
    for (int v: visits)
        CHECK(v == 1);
    for (int v: on_plane)
        CHECK(v <= 1);

    CHECK_FALSE(modules.PlaneToPanel(110.0, 40.0).has_value());   // gap between panels 0 and 1
    CHECK_FALSE(modules.PlaneToPanel(110.0, 82.0).has_value());   // gap cross, where the beam is

    // Empty table: the identity, the whole image one panel -1.
    const DetectorModules none;
    const auto p = none.PlaneToPanel(12.5, -3.0);
    REQUIRE(p.has_value());
    CHECK(p->panel == -1);
    CHECK(p->x == 12.5);
    CHECK(p->y == -3.0);
    CHECK(none.ImageToPlane(-1, 4.0, 5.0) == std::pair{4.0, 5.0});
    const PlaneBox image = none.PlaneBounds(PW, PH);
    CHECK(image.x0 == 0.0);
    CHECK(image.x1 == static_cast<double>(PW));
    size_t n = 0;
    none.ForEachPixel(PW, PH, [&](int64_t panel, size_t i, double px, double py) {
        CHECK(panel == -1);
        CHECK(i == n++);
        CHECK(px == static_cast<double>(i % PW));
        CHECK(py == static_cast<double>(i / PW));
    });
    CHECK(n == static_cast<size_t>(PW * PH));
}

// Spot centroids: found in the image, mapped to the plane by SpotAnalyze. A count threshold on interior
// spots, so the finder's neighbourhood never reaches another panel.
TEST_CASE("PanelLayout_SpotPlaneCoordinates", "[PanelLayout]") {
    const Scene scene(INTERIOR_SPOTS);
    const auto a = FindSpots(scene.assembled, scene.assembled_image);
    const auto s = FindSpots(scene.stacked, scene.stacked_image);
    CHECK(a.size() == INTERIOR_SPOTS.size());
    CheckSameSpots(a, s);
}

// Predictions go plane -> panel -> image and come back to where they were on the plane; the integrated
// spot is observed where it is on the plane in both layouts.
TEST_CASE("PanelLayout_IntegrationRoundTrip", "[PanelLayout]") {
    const Scene scene(INTERIOR_SPOTS);
    const auto a = Integrate(scene.assembled, scene.assembled_mask, scene.assembled_image, INTERIOR_SPOTS);
    const auto s = Integrate(scene.stacked, scene.stacked_mask, scene.stacked_image, INTERIOR_SPOTS);
    REQUIRE(a.size() == INTERIOR_SPOTS.size());
    REQUIRE(s.size() == INTERIOR_SPOTS.size());
    for (const auto &[h, r]: s) {
        const Spot &spot = INTERIOR_SPOTS[h - 1];
        CHECK(r.predicted_x == Catch::Approx(spot.x).margin(1e-3));
        CHECK(r.predicted_y == Catch::Approx(spot.y).margin(1e-3));
        CHECK(r.observed_x == Catch::Approx(a.at(h).observed_x).margin(1e-3));
        CHECK(r.observed_y == Catch::Approx(a.at(h).observed_y).margin(1e-3));
    }
}

// TODO stage 1: the azimuthal-integration map built per stored pixel through the panel table - same bin,
// resolution and correction as the assembled pixel it lands on.
TEST_CASE("PanelLayout_AzintMapping", "[.][PanelLayoutTodo]") {
    const Scene scene(INTERIOR_SPOTS);
    const AzimuthalIntegrationMapping a(scene.assembled, scene.assembled_mask);
    const AzimuthalIntegrationMapping s(scene.stacked, scene.stacked_mask);
    size_t bin = 0, resolution = 0, correction = 0;   // pixels that differ
    scene.modules.ForEachPixel(PW, NPANELS * PH, [&](int64_t, size_t i, double px, double py) {
        const size_t j = Scene::AssembledIndex(px, py);
        bin += s.GetPixelToBin()[i] != a.GetPixelToBin()[j];
        resolution += std::fabs(s.Resolution()[i] - a.Resolution()[j]) > 1e-5f * a.Resolution()[j];
        correction += std::fabs(s.Corrections()[i] - a.Corrections()[j]) > 1e-5f * a.Corrections()[j];
    });
    CHECK(bin == 0);
    CHECK(resolution == 0);
    CHECK(correction == 0);
}

// TODO stage 1: the spot finders' resolution mask, per stored pixel through the panel table.
TEST_CASE("PanelLayout_ResolutionMask", "[.][PanelLayoutTodo]") {
    const Scene scene(INTERIOR_SPOTS);
    const auto a = AzimuthalIntegrationMapping(scene.assembled, scene.assembled_mask).ResolutionMaskBits(3.0f, 8.0f);
    const auto s = AzimuthalIntegrationMapping(scene.stacked, scene.stacked_mask).ResolutionMaskBits(3.0f, 8.0f);
    auto bit = [](const std::vector<uint32_t> &v, size_t i) { return (v[i / 32] >> (i % 32)) & 1u; };
    size_t differ = 0;
    scene.modules.ForEachPixel(PW, NPANELS * PH, [&](int64_t, size_t i, double px, double py) {
        differ += bit(*s, i) != bit(*a, Scene::AssembledIndex(px, py));
    });
    CHECK(differ == 0);
}

// Bragg integration on the stacked layout measures what it measures on the assembled image - stencil
// direction, radial background and polarisation taken from the plane position, not the stored one.
// Also with what reads the radial direction hardest: polarisation, a ring elongated by a bandwidth, the
// radial background correction, the owner map, and the box sum. The interior spots moved off the exact
// half pixel: there the profile grid centres on the rounded position, and rounding a half picks the other
// neighbour once a panel is turned.
void CheckBraggIntegration(bool gpu) {
    const std::vector<Spot> spots = {{50.3, 40.7}, {80.4, 45.3}, {160.6, 30.2}, {40.4, 140.8}, {150.2, 130.6}};
    const Scene scene(spots);
    auto setup = [](DiffractionExperiment x, int variant) {
        BraggIntegrationSettings bis = x.GetBraggIntegrationSettings();
        x.PolarizationFactor(0.99f);
        if (variant == 1) {
            x.BandwidthFWHM(0.02f);
            bis.StencilKSigma(2.0f);
        }
        if (variant == 2)
            bis.BackgroundRadialCorrection(true);
        if (variant == 3)
            bis.Overlap(OverlapMode::Exclude);
        if (variant == 4)
            bis.Integrator(IntegratorMode::BoxSum);
        if (variant == 5)
            bis.Integrator(IntegratorMode::ProfileEmpirical);
        x.ImportBraggIntegrationSettings(bis);
        return x;
    };
    for (int variant = 0; variant < 6; variant++) {
        CAPTURE(variant);
        const auto a = Integrate(setup(scene.assembled, variant), scene.assembled_mask, scene.assembled_image, spots,
                                 gpu);
        const auto s = Integrate(setup(scene.stacked, variant), scene.stacked_mask, scene.stacked_image, spots, gpu);
        REQUIRE(s.size() == a.size());
        for (const auto &[h, r]: s) {
            CAPTURE(h);
            CHECK(r.I == Catch::Approx(a.at(h).I).epsilon(1e-5));
            CHECK(r.sigma == Catch::Approx(a.at(h).sigma).epsilon(1e-5));
            CHECK(r.bkg == Catch::Approx(a.at(h).bkg).epsilon(1e-5));
            CHECK(r.image_scale_corr == Catch::Approx(a.at(h).image_scale_corr).epsilon(1e-5));
        }
    }
}

TEST_CASE("PanelLayout_BraggIntegration", "[PanelLayout]") {
    CheckBraggIntegration(false);
}

// The same on the CUDA engine, stacked against assembled on the GPU.
TEST_CASE("PanelLayout_BraggIntegrationGPU", "[PanelLayout][GPU]") {
    if (get_gpu_count() == 0)
        SKIP("No GPU");
    CheckBraggIntegration(true);
}

// TODO stage 3: spot finding per panel. Two half spots on the facing edges of panels 0 and 1 as stored -
// adjacent rows of the stacked array, 139 px apart on the plane - stay two spots, and the local-background
// (SNR) finder agrees with the assembled image away from panel edges.
TEST_CASE("PanelLayout_SpotsNeverCrossPanels", "[.][PanelLayoutTodo]") {
    // Panel 0 stored row 63 and panel 1 stored row 0, both at stored column 30.
    const std::vector<Spot> edge = {{40.0, 73.0}, {179.0, 73.0}};
    const Scene scene(edge);
    CheckSameSpots(FindSpots(scene.assembled, scene.assembled_image),
                   FindSpots(scene.stacked, scene.stacked_image));

    const Scene interior(INTERIOR_SPOTS);
    CheckSameSpots(FindSpots(interior.assembled, interior.assembled_image, 3.0f),
                   FindSpots(interior.stacked, interior.stacked_image, 3.0f));
}