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

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