// 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 <cstring>
#include <optional>
#include <vector>

#include "../common/DetectorSetup.h"
#include "../common/DiffractionExperiment.h"
#include "../common/JFJochMessages.h"
#include "../common/PixelMask.h"
#include "../image_analysis/beam_stop/ShadowFinder.h"

namespace {
    // Odd and square, so the beam sits on a pixel and a cross-shaped scene is exactly 4-fold
    // symmetric; deliberately not a multiple of 64, so the column-blocked passes meet a short
    // final block.
    constexpr int W = 257, H = 257, C = 128;
    constexpr int NFRAMES = 12;
    constexpr int32_t BACKGROUND = 2;      // integer and noise-free, so every mean is exact
    constexpr int STOP_R = 22, ARM_HALF = 5;

    constexpr size_t I(int x, int y) { return static_cast<size_t>(y) * W + x; }

    DiffractionExperiment TestExperiment() {
        DiffractionExperiment x(DetDECTRIS(W, H, "Test detector", ""));
        x.IncidentEnergy_keV(WVL_1A_IN_KEV).DetectorDistance_mm(150.0f);
        x.BeamX_pxl(static_cast<float>(C)).BeamY_pxl(static_cast<float>(C));
        return x;
    }

    // Flat background, an opaque disk on the beam, and an arm running off it to the edge - a beam
    // stop. `cross` gives it four arms instead of one, making the scene invariant under a quarter
    // turn. `reflection` puts a cluster bright enough to count as a reflection inside the disk.
    std::vector<int32_t> Scene(bool cross, bool reflection) {
        std::vector<int32_t> f(static_cast<size_t>(W) * H, BACKGROUND);
        for (int y = 0; y < H; y++) {
            for (int x = 0; x < W; x++) {
                const int dx = x - C, dy = y - C;
                bool blocked = dx * dx + dy * dy <= STOP_R * STOP_R;
                blocked = blocked || (cross ? (std::abs(dy) <= ARM_HALF || std::abs(dx) <= ARM_HALF)
                                            : (std::abs(dy) <= ARM_HALF && dx >= 0));
                if (blocked)
                    f[I(x, y)] = 0;
            }
        }
        if (reflection) {
            for (int y = C - 4; y <= C; y++)
                for (int x = C - 16; x <= C - 12; x++)
                    f[I(x, y)] = 100;
        }
        return f;
    }

    // CompressedImage does not own its pixels, so the frames have to outlive the calls.
    void Feed(ShadowFinder &finder, std::vector<std::vector<int32_t>> &frames, bool cross,
              bool reflection_on_first) {
        std::vector<uint8_t> buffer;
        for (int f = 0; f < NFRAMES; f++) {
            frames.push_back(Scene(cross, reflection_on_first && (f == 0)));
            DataMessage msg{};
            msg.image = CompressedImage(frames.back(), W, H);
            finder.AddImage(msg, buffer);
        }
    }
}

// The scene is a beam stop: an opaque disk on the beam with an arm running off it. What comes back
// has to be the stop and nothing else - the corners of a detector are not shadowed - and a
// reflection recorded through the penumbra is given back rather than masked.
TEST_CASE("ShadowFinder_FindsAnInjectedBeamStop", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    std::vector<std::vector<int32_t>> frames;
    Feed(finder, frames, /*cross=*/false, /*reflection_on_first=*/true);
    REQUIRE(finder.GetFrameCount() == NFRAMES);

    const auto mask = finder.GetMask();
    REQUIRE(mask.size() == static_cast<size_t>(W) * H);

    CHECK(mask[I(C, C)] == 1);                        // the stop itself
    CHECK(mask[I(C + STOP_R - 3, C)] == 1);
    CHECK(mask[I(W - 3, C)] == 1);                    // the arm, followed to the edge
    CHECK(mask[I(W - 3, C + 4 * ARM_HALF)] == 0);     // and nothing beside it
    CHECK(mask[I(0, 0)] == 0);
    CHECK(mask[I(W - 1, 0)] == 0);
    CHECK(mask[I(0, H - 1)] == 0);
    CHECK(mask[I(W - 1, H - 1)] == 0);
    CHECK(mask[I(C - 14, C - 2)] == 0);               // a recorded reflection is given back

    // The mean projection is what the mask is computed from: exact here, because the scene is
    // integer and noise-free.
    const auto projection = finder.GetMeanProjection();
    REQUIRE(projection.size() == mask.size());
    CHECK(projection[I(0, 0)] == Catch::Approx(BACKGROUND));
    CHECK(projection[I(C, C)] == Catch::Approx(0.0));

    // Pinned from the serial implementation. A rewrite of the dilation, the hole fill or the ring
    // median that moves the mask by one pixel fails here, rather than in a merging statistic
    // several stages downstream.
    CHECK(std::count(mask.begin(), mask.end(), 1u) == 2612);
}

// The per-pixel passes are split across threads, so where the split falls must not be visible in the
// answer. The x pass and the y pass of the dilation and of the pooled sum are written differently -
// one a plain scan, the other blocked by column - so an x/y asymmetry is the plausible regression.
TEST_CASE("ShadowFinder_MaskDoesNotDependOnTheThreadCount", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    std::vector<std::vector<int32_t>> frames;
    Feed(finder, frames, /*cross=*/false, /*reflection_on_first=*/true);

    const auto one = finder.GetMask(1);
    CHECK(finder.GetMask(3) == one);
    CHECK(finder.GetMask(8) == one);
}

// Workers accumulate into shards of their own and the shards are summed when the projection is read,
// so which worker saw which frame must not reach the answer - including the maximum, which only one
// shard holds when the reflection is on a single frame.
TEST_CASE("ShadowFinder_ShardingDoesNotChangeTheProjection", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);

    ShadowFinder serial(x, pixel_mask);
    ShadowFinder sharded(x, pixel_mask);
    sharded.SetShardCount(4);

    std::vector<std::vector<int32_t>> frames;
    std::vector<uint8_t> buffer;
    for (int f = 0; f < NFRAMES; f++) {
        frames.push_back(Scene(/*cross=*/false, /*reflection=*/f == 0));
        DataMessage msg{};
        msg.image = CompressedImage(frames.back(), W, H);
        serial.AddImage(msg, buffer, 0);
        sharded.AddImage(msg, buffer, static_cast<size_t>(f) % 4);
    }

    CHECK(serial.GetFrameCount() == sharded.GetFrameCount());

    const auto a = serial.GetMeanProjection();
    const auto b = sharded.GetMeanProjection();
    REQUIRE(a.size() == b.size());
    // NAN marks a pixel nothing counted, and NAN != NAN, so compare the bits rather than the values.
    CHECK(memcmp(a.data(), b.data(), a.size() * sizeof(float)) == 0);

    // The reflection is on one frame, so its maximum lives in a single shard. If the fold lost it,
    // the mask would swallow the reflection instead of giving it back.
    CHECK(serial.GetMask(1) == sharded.GetMask(1));
    CHECK(sharded.GetMask(1)[I(C - 14, C - 2)] == 0);
}

// Four opaque arms and a centred disk: the scene is invariant under a quarter turn, so the mask must
// be too, whatever the thread count.
TEST_CASE("ShadowFinder_ASymmetricSceneGivesASymmetricMask", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    std::vector<std::vector<int32_t>> frames;
    Feed(finder, frames, /*cross=*/true, /*reflection_on_first=*/false);

    const auto mask = finder.GetMask(8);
    for (int y = 0; y < H; y++) {
        for (int xi = 0; xi < W; xi++) {
            REQUIRE(mask[I(xi, y)] == mask[I(y, xi)]);            // transpose
            REQUIRE(mask[I(xi, y)] == mask[I(W - 1 - y, xi)]);    // quarter turn
        }
    }
}

// Hardware that shadows the detector need not touch the direct beam: a pin or a loop begins some way
// out in radius, with lit detector between it and the stop. The scene here is that - a beam stop
// with its arm, and a separate opaque patch far from both - and the patch has to come back as
// shadow. Anchoring the search at the beam centre, as an earlier version did, returned the stop and
// discarded the patch however deep it was.
TEST_CASE("ShadowFinder_FindsAShadowThatDoesNotTouchTheBeam", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    // Well clear of the stop, and wide enough that the ring it sits on still has lit pixels to be
    // compared against - as a real pin shadow does, covering part of a ring and not all of it.
    constexpr int PATCH_X0 = 175, PATCH_X1 = 245, PATCH_Y0 = 30, PATCH_Y1 = 90;

    std::vector<std::vector<int32_t>> frames;
    std::vector<uint8_t> buffer;
    for (int f = 0; f < NFRAMES; f++) {
        frames.push_back(Scene(/*cross=*/false, /*reflection=*/false));
        for (int y = PATCH_Y0; y <= PATCH_Y1; y++)
            for (int xi = PATCH_X0; xi <= PATCH_X1; xi++)
                frames.back()[I(xi, y)] = 0;
        DataMessage msg{};
        msg.image = CompressedImage(frames.back(), W, H);
        finder.AddImage(msg, buffer);
    }

    const auto mask = finder.GetMask();
    CHECK(mask[I((PATCH_X0 + PATCH_X1) / 2, (PATCH_Y0 + PATCH_Y1) / 2)] == 1);
    CHECK(mask[I(PATCH_X0 + 5, PATCH_Y0 + 5)] == 1);
    CHECK(mask[I(C, C)] == 1);                  // the stop is still found
    CHECK(mask[I(PATCH_X0 - 40, PATCH_Y0)] == 0);   // and the lit detector between them is kept
    CHECK(mask[I(0, H - 1)] == 0);
}

// A flat, unobstructed scene has no shadow in it. The per-pixel test no longer has to reach the
// beam centre, so what keeps it from calling a wandering background a shadow is size alone - and
// that has to hold on a scene where nothing is blocked at all.
TEST_CASE("ShadowFinder_FindsNothingOnAnUnobstructedScene", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    std::vector<std::vector<int32_t>> frames;
    std::vector<uint8_t> buffer;
    for (int f = 0; f < NFRAMES; f++) {
        frames.emplace_back(static_cast<size_t>(W) * H, BACKGROUND);
        DataMessage msg{};
        msg.image = CompressedImage(frames.back(), W, H);
        finder.AddImage(msg, buffer);
    }

    const auto mask = finder.GetMask();
    CHECK(std::count(mask.begin(), mask.end(), 1u) == 0);
}

// The rings are drawn about a beam centre, and a centre in a file can be a long way from the truth.
// On a background that falls with radius, rings drawn about the wrong point cut across that
// fall-off, and pixels that are simply further out than the ring's median read as shadow: a large
// part of the detector comes back masked with nothing blocking it. The caller therefore names the
// centre it has measured, and that is the one the comparison has to use.
TEST_CASE("ShadowFinder_TheRingsFollowTheCentreTheCallerNames", "[ShadowFinder]") {
    // Far enough out that the ring about the file's centre spans a wide range of true radii.
    constexpr int OFFSET = 100;
    // Opaque, well clear of the beam, and larger than the smallest region the test may return.
    constexpr int PATCH_X0 = 175, PATCH_X1 = 245, PATCH_Y0 = 30, PATCH_Y1 = 90;
    constexpr int PATCH_PX = (PATCH_X1 - PATCH_X0 + 1) * (PATCH_Y1 - PATCH_Y0 + 1);

    DiffractionExperiment x = TestExperiment();
    x.BeamX_pxl(static_cast<float>(C + OFFSET));   // what the file claims
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    // A background falling with the true radius, kept under MIN_REFLECTION so no pixel is exempt.
    std::vector<std::vector<int32_t>> frames;
    std::vector<uint8_t> buffer;
    for (int f = 0; f < NFRAMES; f++) {
        frames.emplace_back(static_cast<size_t>(W) * H, 0);
        for (int y = 0; y < H; y++)
            for (int xi = 0; xi < W; xi++) {
                const float r = std::hypot(static_cast<float>(xi - C), static_cast<float>(y - C));
                const bool blocked = xi >= PATCH_X0 && xi <= PATCH_X1 && y >= PATCH_Y0 && y <= PATCH_Y1;
                frames.back()[I(xi, y)] = blocked ? 0 : std::lround(24.0f * std::exp(-r / 80.0f));
            }
        DataMessage msg{};
        msg.image = CompressedImage(frames.back(), W, H);
        finder.AddImage(msg, buffer);
    }

    const auto at_file = finder.GetMask();
    finder.BeamCenter(static_cast<float>(C), static_cast<float>(C));
    const auto at_measured = finder.GetMask();

    // The patch is shadow either way - it is opaque, and no centre makes it look lit.
    CHECK(at_file[I((PATCH_X0 + PATCH_X1) / 2, (PATCH_Y0 + PATCH_Y1) / 2)] == 1);
    CHECK(at_measured[I((PATCH_X0 + PATCH_X1) / 2, (PATCH_Y0 + PATCH_Y1) / 2)] == 1);

    // About the true centre the rings are flat and only the patch and its penumbra come back;
    // about the file's they cut across the fall-off and a large part of the detector does.
    const auto masked_at_file = std::count(at_file.begin(), at_file.end(), 1u);
    const auto masked_at_measured = std::count(at_measured.begin(), at_measured.end(), 1u);
    CHECK(masked_at_measured < 2 * PATCH_PX);
    CHECK(masked_at_file > 3 * PATCH_PX);
}

// The background is not flat around a ring: a polarized source suppresses it in its own plane, by
// a factor that reaches four at the 2 theta a short detector distance puts in a corner. That is
// several times the dip this class looks for, so with the modulation left in, half of every outer
// ring reads as shadow with nothing in the beam. The scene here is exactly that and nothing else -
// a background carrying the Kahn factor of a horizontally polarized source, no hardware anywhere -
// and it must come back empty when the experiment states its polarization, and does not when it
// says nothing, which is what the correction is for.
TEST_CASE("ShadowFinder_APolarizedBackgroundIsNotAShadow", "[ShadowFinder]") {
    // Short enough that this small detector reaches 2 theta = 70 degrees in its corner, where
    // the source suppresses its own plane to a fifth. A real geometry gets there with a short
    // crystal-to-detector distance and a large detector.
    constexpr float DISTANCE_MM = 5.0f;
    constexpr float POLARIZATION = 0.99f;
    // Counts per pixel per frame across the azimuth the source does not suppress. Below
    // MIN_REFLECTION, so no pixel of this scene is exempt as a recorded reflection.
    constexpr int LEVEL = 24;

    DiffractionExperiment reference = TestExperiment();
    const float pixel_mm = reference.GetPixelSize_mm();

    std::vector<int32_t> frame(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++) {
            const double dx = x - C, dy = y - C;
            const double rho = std::hypot(dx, dy) * pixel_mm;
            const double cos_2theta_sq = DISTANCE_MM * DISTANCE_MM
                                         / (DISTANCE_MM * DISTANCE_MM + rho * rho);
            const double rr = dx * dx + dy * dy;
            const double cos_2phi = rr > 0.0 ? (dx * dx - dy * dy) / rr : 0.0;
            const double factor = 0.5 * (1.0 + cos_2theta_sq
                                         - POLARIZATION * cos_2phi * (1.0 - cos_2theta_sq));
            // Normalised per ring by the azimuth the source does not suppress, so what is left
            // varies around a ring and not with radius - the ring median absorbs the rest.
            const double at_peak = 0.5 * (1.0 + cos_2theta_sq
                                          + POLARIZATION * (1.0 - cos_2theta_sq));
            frame[I(x, y)] = static_cast<int32_t>(std::lround(LEVEL * factor / at_peak));
        }

    auto mask_for = [&](const std::optional<float> &polarization) {
        DiffractionExperiment x = TestExperiment();
        x.DetectorDistance_mm(DISTANCE_MM).PolarizationFactor(polarization);
        const PixelMask pixel_mask(x);
        ShadowFinder finder(x, pixel_mask);
        std::vector<uint8_t> buffer;
        for (int f = 0; f < NFRAMES; f++) {
            DataMessage msg{};
            msg.image = CompressedImage(frame, W, H);
            finder.AddImage(msg, buffer);
        }
        const auto mask = finder.GetMask();
        return std::count(mask.begin(), mask.end(), 1u);
    };

    CHECK(mask_for(POLARIZATION) == 0);
    CHECK(mask_for(std::nullopt) > static_cast<long>(W) * H / 20);
}

// Inside the stop a whole ring is blocked, so its own median is blocked too and the per-pixel
// comparison has nothing to work with; the walk that finds those rings is what covers the disk.
// It decides by comparing a ring against what this detector's background typically is - and NOT
// against the brightest ring anywhere further out, which a sample whose background peaks well away
// from the beam turns into a beam stop the size of that peak. The scene here is that: an ordinary
// stop, and a strong ring at four times the background far outside it. What comes back has to be
// the stop, not a disk reaching the ring.
TEST_CASE("ShadowFinder_ABrightRingIsNotABeamStop", "[ShadowFinder]") {
    constexpr int RING_R = 100, RING_HALF = 2, RING_LEVEL = 8 * BACKGROUND;

    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);
    ShadowFinder finder(x, pixel_mask);

    std::vector<std::vector<int32_t>> frames;
    std::vector<uint8_t> buffer;
    for (int f = 0; f < NFRAMES; f++) {
        frames.push_back(Scene(/*cross=*/false, /*reflection=*/false));
        for (int y = 0; y < H; y++)
            for (int xi = 0; xi < W; xi++) {
                const double r = std::hypot(xi - C, y - C);
                if (std::abs(r - RING_R) <= RING_HALF)
                    frames.back()[I(xi, y)] = RING_LEVEL;
            }
        DataMessage msg{};
        msg.image = CompressedImage(frames.back(), W, H);
        finder.AddImage(msg, buffer);
    }

    const auto mask = finder.GetMask();
    CHECK(mask[I(C, C)] == 1);                   // the stop is still covered
    CHECK(mask[I(C, C - STOP_R - 20)] == 0);     // lit detector between the stop and the ring
    CHECK(mask[I(C, C - RING_R + 5)] == 0);
    CHECK(mask[I(C, C - RING_R)] == 0);          // and the ring itself
    // Anything much beyond the stop, its arm and their penumbra means the walk ran away.
    CHECK(std::count(mask.begin(), mask.end(), 1u) < 6000);
}