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