Jungfraujoch/tests/ShadowFinderTest.cpp

// 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 <climits>
#include <cmath>
#include <cstring>
#include <numbers>
#include <optional>
#include <thread>
#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"
#include "../image_analysis/beam_stop/ShadowFinderInternal.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. Nothing here lets part of the beam through, so the transmitting
    // shape may redraw the stop's edge but must not add to the mask: the count is the opaque one.
    CHECK(std::count(mask.begin(), mask.end(), 1u) == 2612);
}

// A holder arm that lets part of the beam through, on a background bright enough that every pixel
// reaches the reflection guard's count, crossing a module gap wider than the bridge. Each of the
// three hid the arm on its own: the partly transmitting stretch never reaches SHADOW_RATIO, the gap
// cut what did into pieces too small to be believed, and a count threshold read the attenuated
// background under the arm as recorded reflections and gave every pixel of it back.
TEST_CASE("ShadowFinder_FindsAnArmThatLetsPartOfTheBeamThrough", "[ShadowFinder]") {
    constexpr int32_t BRIGHT = 50;               // even the arm reaches MIN_REFLECTION
    constexpr int ARM_HALF_WIDE = 15, GAP_X0 = 200, GAP_X1 = 216, OPAQUE_FROM_X = 232;

    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, BRIGHT);
        auto &frame = frames.back();
        for (int y = 0; y < H; y++)
            for (int xi = 0; xi < W; xi++) {
                const int dx = xi - C, dy = y - C;
                if (dx * dx + dy * dy <= STOP_R * STOP_R)
                    frame[I(xi, y)] = 0;
                else if (dx >= 0 && std::abs(dy) <= ARM_HALF_WIDE)
                    frame[I(xi, y)] = xi >= OPAQUE_FROM_X ? 0 : BRIGHT * 6 / 10;
                if (xi >= GAP_X0 && xi <= GAP_X1)
                    frame[I(xi, y)] = INT32_MIN;  // no data
            }
        DataMessage msg{};
        msg.image = CompressedImage(frame, W, H);
        finder.AddImage(msg, buffer);
    }

    const auto mask = finder.GetMask();
    CHECK(mask[I(C, C)] == 1);                            // the stop
    CHECK(mask[I(GAP_X0 - 10, C)] == ShadowFinder::TRANSMITTING);  // the arm where it transmits, before the gap
    CHECK(mask[I(GAP_X1 + 10, C + ARM_HALF_WIDE - 3)] == ShadowFinder::TRANSMITTING);  // ... and after it
    CHECK(mask[I(W - 3, C)] != 0);                        // where it is opaque
    CHECK(mask[I(GAP_X1 + 10, C + ARM_HALF_WIDE + 25)] == 0);  // and nothing beside it
    CHECK(mask[I(0, 0)] == 0);
    CHECK(mask[I(C - 60, C)] == 0);
}

// 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 add their frames concurrently, each starting at a different band of the projection, so
// which worker added which frame, and in what order, must not reach the answer - including the
// maximum, which one frame alone holds when the reflection is on a single frame.
TEST_CASE("ShadowFinder_ConcurrentWorkersDoNotChangeTheProjection", "[ShadowFinder]") {
    const DiffractionExperiment x = TestExperiment();
    const PixelMask pixel_mask(x);

    ShadowFinder serial(x, pixel_mask);
    ShadowFinder concurrent(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=*/f == 0));
        DataMessage msg{};
        msg.image = CompressedImage(frames.back(), W, H);
        serial.AddImage(msg, buffer);
    }
    std::vector<std::thread> workers;
    for (int t = 0; t < 4; t++)
        workers.emplace_back([&, t] {
            std::vector<uint8_t> worker_buffer;
            for (int f = NFRAMES - 1 - t; f >= 0; f -= 4) {
                DataMessage msg{};
                msg.image = CompressedImage(frames[f], W, H);
                concurrent.AddImage(msg, worker_buffer);
            }
        });
    for (auto &w : workers) w.join();

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

    const auto a = serial.GetMeanProjection();
    const auto b = concurrent.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 only that frame's maximum sees it. If it were lost, the mask
    // would swallow the reflection instead of giving it back.
    CHECK(serial.GetMask(1) == concurrent.GetMask(1));
    CHECK(concurrent.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);
}

// A detector on a swung 2theta arm. The beam no longer lands on the geometry's centre (the point of
// normal incidence) and the rings of constant 2 theta are no circle about it, nor about any pixel;
// the solid angle of a pixel, which falls away from the normal, also varies along each of them. A
// background that depends on 2 theta alone, seen through that solid angle, has nothing in the
// beam - and circles drawn about the centre cut across its fall-off and read a large part of the
// detector as shadow.
TEST_CASE("ShadowFinder_ATiltedDetectorIsNotAShadow", "[ShadowFinder]") {
    // Short enough, and tilted enough, that the beam lands 47 px from the centre.
    constexpr float DISTANCE_MM = 20.0f;
    constexpr float TILT_RAD = 0.174533f;   // 10 degrees

    DiffractionExperiment x = TestExperiment();
    x.DetectorDistance_mm(DISTANCE_MM).PoniRot1_rad(TILT_RAD);
    const DiffractionGeometry geom = x.GetDiffractionGeometry();

    // Below MIN_REFLECTION, so no pixel is exempt as a recorded reflection.
    std::vector<int32_t> frame(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int xi = 0; xi < W; xi++) {
            const float two_theta = geom.TwoTheta_rad(static_cast<float>(xi), static_cast<float>(y));
            const float solid_angle = geom.CalcAzIntSolidAngleCorr(static_cast<float>(xi), static_cast<float>(y));
            frame[I(xi, y)] = static_cast<int32_t>(std::lround(24.0f * solid_angle * std::exp(-two_theta / 0.15f)));
        }

    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();
    CHECK(std::count_if(mask.begin(), mask.end(), [](uint32_t v) { return v != 0; }) == 0);
}

// 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);
}

// A beam stop with four arms on a background polarized in the plane along the image's y, at a short
// distance, in an experiment that states the usual horizontal plane. Divided by the factor of the
// stated plane the background's lobes read as shadow, and how much of them did followed the third
// decimal of the stated degree. The plane is taken from the background, so the stop and its arms come
// back, the lobes do not, and the mask is the same for any stated degree and either stated plane.
TEST_CASE("ShadowFinder_ThePolarizationPlaneIsTheOneTheBackgroundShows", "[ShadowFinder]") {
    constexpr float DISTANCE_MM = 8.0f;
    constexpr double TRUE_POLARIZATION = -0.998;   // the plane along the image's y
    constexpr int LEVEL = 24;

    const float pixel_mm = TestExperiment().GetPixelSize_mm();
    const auto stop = Scene(/*cross=*/true, false);
    std::vector<int32_t> frame(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int xi = 0; xi < W; xi++) {
            const double dx = xi - 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 - TRUE_POLARIZATION * cos_2phi * (1.0 - cos_2theta_sq));
            frame[I(xi, y)] = stop[I(xi, y)] == 0 ? 0 : static_cast<int32_t>(std::lround(LEVEL * factor));
        }

    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);
        }
        return finder.GetMask();
    };

    const auto mask = mask_for(0.99f);
    CHECK(mask[I(C, C)] == ShadowFinder::SHADOW);              // the stop
    CHECK(mask[I(W - 3, C)] == ShadowFinder::SHADOW);          // and its arms, along x
    CHECK(mask[I(C, 3)] == ShadowFinder::SHADOW);              // and along y, inside the suppressed lobe
    CHECK(mask[I(C + 20, 3)] == 0);                            // but not the lobe itself
    CHECK(mask[I(C - 20, H - 4)] == 0);
    CHECK(mask_for(0.998f) == mask);
    CHECK(mask_for(-0.998f) == mask);
}

// The plane is decided from the sector medians of each band: the cos 2phi harmonic has to follow the
// profile a polarized source gives it, with a one-sided shadow over a third of the ring left out by
// comparing opposite sectors. A background without modulation decides nothing.
TEST_CASE("ShadowFinder_PolarizationPlane", "[ShadowFinder]") {
    using namespace shadow_finder;
    constexpr int N_BANDS = 8;
    std::vector<float> amplitude(N_BANDS);
    for (int band = 0; band < N_BANDS; band++)
        amplitude[band] = 0.05f * static_cast<float>(band + 1);
    auto medians = [&](double sign, bool shadow) {
        std::vector<double> m(static_cast<size_t>(N_BANDS) * HARMONIC_SECTORS);
        for (int band = 0; band < N_BANDS; band++)
            for (int k = 0; k < HARMONIC_SECTORS; k++) {
                const double phi = (k + 0.5) * 2.0 * std::numbers::pi / HARMONIC_SECTORS - std::numbers::pi;
                double v = 1.0 + sign * amplitude[band] * std::cos(2.0 * phi);
                if (shadow && k >= 8 && k < 16)
                    v *= 0.7;                                   // a third of the ring around phi = 0
                m[static_cast<size_t>(band) * HARMONIC_SECTORS + k] = v;
            }
        return m;
    };
    CHECK(PolarizationPlane(medians(-1.0, false), amplitude) == 1.0f);   // dim along x: horizontal
    CHECK(PolarizationPlane(medians(+1.0, false), amplitude) == -1.0f);  // dim along y: vertical
    CHECK(PolarizationPlane(medians(+1.0, true), amplitude) == -1.0f);   // the shadow along x does not hide it
    CHECK(PolarizationPlane(medians(0.0, false), amplitude) == 0.0f);    // no modulation: no decision
}

// The other side of that correction: a polarization factor that is not the beam's. Here the
// background is unpolarized and the experiment states 0.99, so after the division every outer ring is
// bright along x and dim along y - a smooth lobe below PENUMBRA_RATIO across the top and bottom of the
// detector with nothing in the beam. On its own the lobe is never deep, but a small absorbing speck
// inside it is, and the transmitting-hardware step joined the two into one piece "dim everywhere and
// deep in places" and masked the whole lobe. The step compares a pixel with its ring's own second
// harmonic, so the lobe is explained away and the speck, far below a shadow's size, stays unmasked.
TEST_CASE("ShadowFinder_AWrongPolarizationFactorIsNotAnArm", "[ShadowFinder]") {
    constexpr float DISTANCE_MM = 6.0f;
    constexpr int LEVEL = 24;
    constexpr int SPECK_X = C, SPECK_Y = 20, SPECK_R = 11;

    DiffractionExperiment x = TestExperiment();
    x.DetectorDistance_mm(DISTANCE_MM).PolarizationFactor(0.99f);
    const float pixel_mm = x.GetPixelSize_mm();

    std::vector<int32_t> frame(static_cast<size_t>(W) * H);
    for (int y = 0; y < H; y++)
        for (int xi = 0; xi < W; xi++) {
            const double rho = std::hypot(xi - C, y - C) * pixel_mm;
            const double cos_2theta_sq = DISTANCE_MM * DISTANCE_MM / (DISTANCE_MM * DISTANCE_MM + rho * rho);
            frame[I(xi, y)] = static_cast<int32_t>(std::lround(LEVEL * 0.5 * (1.0 + cos_2theta_sq)));
            const int sx = xi - SPECK_X, sy = y - SPECK_Y;
            if (sx * sx + sy * sy <= SPECK_R * SPECK_R)
                frame[I(xi, y)] /= 10;
        }

    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();
    CHECK(mask[I(C, 5)] == 0);                   // the dim lobe along y
    CHECK(mask[I(SPECK_X, SPECK_Y)] == 0);       // the speck, too small to be a shadow
    CHECK(std::count(mask.begin(), mask.end(), 1u) == 0);
}

// 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);
}

#ifdef JFJOCH_USE_CUDA
#include "../common/CUDAWrapper.h"
#include "../compression/JFJochCompressor.h"

namespace {
    // The mask and the mean projection of the same frames, once from the host projection (the frames
    // handed over uncompressed) and once from the device's (handed over as bitshuffle+LZ4, which is
    // what sends them to the GPU).
    struct HostAndDevice {
        std::vector<uint32_t> host_mask, device_mask;
        std::vector<float> host_mean, device_mean;
    };

    HostAndDevice MaskBothWays(const DiffractionExperiment &x, const std::vector<std::vector<int32_t>> &frames) {
        const PixelMask pixel_mask(x);
        HostAndDevice out;
        std::vector<uint8_t> buffer;
        {
            ShadowFinder finder(x, pixel_mask);
            for (const auto &frame : frames) {
                DataMessage msg{};
                msg.image = CompressedImage(frame, W, H);
                finder.AddImage(msg, buffer);
            }
            out.host_mask = finder.GetMask();
            out.host_mean = finder.GetMeanProjection();
        }
        {
            ShadowFinder finder(x, pixel_mask);
            JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
            std::vector<std::vector<uint8_t>> compressed;
            for (const auto &frame : frames) {
                compressed.push_back(compressor.Compress(frame));
                DataMessage msg{};
                msg.image = CompressedImage(compressed.back().data(), compressed.back().size(), W, H,
                                            CompressedImageMode::Int32, CompressionAlgorithm::BSHUF_LZ4);
                finder.AddImage(msg, buffer);
            }
            out.device_mask = finder.GetMask();
            out.device_mean = finder.GetMeanProjection();
        }
        return out;
    }
}

// The mask is made on the GPU wherever the projection is there. On scenes where no pixel sits within a
// rounding of a threshold the two must agree to the pixel: every step but the polarization factor,
// the Poisson test's logarithm and the arm search's azimuth is exact on both, and these scenes are
// built so that none of the three decides anything at an edge.
TEST_CASE("ShadowFinder_DeviceMaskMatchesHost", "[ShadowFinder]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");

    SECTION("a beam stop with a reflection behind it") {
        std::vector<std::vector<int32_t>> frames;
        for (int f = 0; f < NFRAMES; f++)
            frames.push_back(Scene(false, f == 0));
        const auto r = MaskBothWays(TestExperiment(), frames);
        CHECK(std::count(r.host_mask.begin(), r.host_mask.end(), 1u) == 2612);
        CHECK(r.device_mask == r.host_mask);
        CHECK(std::memcmp(r.device_mean.data(), r.host_mean.data(), r.host_mean.size() * sizeof(float)) == 0);
    }

    SECTION("an arm that lets part of the beam through, across a module gap") {
        constexpr int32_t BRIGHT = 50;
        constexpr int ARM_HALF_WIDE = 15, GAP_X0 = 200, GAP_X1 = 216, OPAQUE_FROM_X = 232;
        std::vector<std::vector<int32_t>> frames;
        for (int f = 0; f < NFRAMES; f++) {
            frames.emplace_back(static_cast<size_t>(W) * H, BRIGHT);
            auto &frame = frames.back();
            for (int y = 0; y < H; y++)
                for (int xi = 0; xi < W; xi++) {
                    const int dx = xi - C, dy = y - C;
                    if (dx * dx + dy * dy <= STOP_R * STOP_R)
                        frame[I(xi, y)] = 0;
                    else if (dx >= 0 && std::abs(dy) <= ARM_HALF_WIDE)
                        frame[I(xi, y)] = xi >= OPAQUE_FROM_X ? 0 : BRIGHT * 6 / 10;
                    if (xi >= GAP_X0 && xi <= GAP_X1)
                        frame[I(xi, y)] = INT32_MIN;
                }
        }
        const auto r = MaskBothWays(TestExperiment(), frames);
        CHECK(std::count(r.host_mask.begin(), r.host_mask.end(), ShadowFinder::TRANSMITTING) > 0);
        CHECK(r.device_mask == r.host_mask);
    }
}

// The pre-scan keeps one accumulator per card and folds them together when the mask is read. Two
// accumulators, each given half of the frames, folded into one must hold exactly what one given every
// frame holds - including a pixel whose only valid values are negative and fell all to the second half,
// which the maximum must take although the first half never counted the pixel.
TEST_CASE("ShadowAccumulatorGPU_FoldedCardsMatchOneCard", "[ShadowFinder]") {
    if (get_gpu_count() == 0)
        SKIP("no GPU");
    constexpr size_t N = static_cast<size_t>(W) * H;
    std::vector<std::vector<int32_t>> frames;
    uint32_t state = 12345;
    const auto next = [&] { state = state * 1664525u + 1013904223u; return state >> 8; };
    for (int f = 0; f < 2 * NFRAMES; f++) {
        frames.emplace_back(N);
        for (auto &v : frames.back())
            v = next() % 50 == 0 ? INT32_MIN : static_cast<int32_t>(next() % 1100) - 100;
        frames.back()[I(3, 3)] = f % 2 == 0 ? INT32_MIN : -7 - f;
        frames.back()[I(5, 5)] = INT32_MIN;
    }
    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);
    std::vector<std::vector<uint8_t>> compressed;
    for (const auto &frame : frames)
        compressed.push_back(compressor.Compress(frame));
    const auto image = [&](size_t f) {
        return CompressedImage(compressed[f].data(), compressed[f].size(), W, H, CompressedImageMode::Int32,
                               CompressionAlgorithm::BSHUF_LZ4);
    };

    ShadowAccumulatorGPU one(N), first(N), second(N);
    for (size_t f = 0; f < frames.size(); f++) {
        one.Add(image(f));
        (f % 2 == 0 ? first : second).Add(image(f));
    }
    first.Fold(second);

    std::vector<int64_t> max_one, sum_one, max_folded, sum_folded;
    std::vector<uint32_t> count_one, count_folded;
    one.Download(max_one, sum_one, count_one);
    first.Download(max_folded, sum_folded, count_folded);
    CHECK(first.GetFrameCount() == one.GetFrameCount());
    CHECK(max_folded[I(3, 3)] == -8);
    CHECK(count_folded[I(5, 5)] == 0);
    CHECK(max_folded == max_one);
    CHECK(sum_folded == sum_one);
    CHECK(count_folded == count_one);
}
#endif