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

#include "../common/GridScanSettings.h"
#include "../common/ScanResult.h"
#include "../image_analysis/grid_scan_analysis/AnalyzeGridScan.h"

namespace {
    // A plain raster with positive steps and no snake, so image number == grid index and a test
    // can name cells by (column, row) directly.
    GridScanSettings MakeGrid(int64_t nx, int64_t ny, float step_x_um, float step_y_um) {
        GridScanSettings grid(nx, step_x_um, step_y_um, false, false);
        grid.ImageNum(nx * ny);
        return grid;
    }

    // Every cell of the raster is collected; the ones named in hits get a protein score, the
    // rest get nothing.
    ScanResult MakeScan(int64_t nx, int64_t ny,
                        const std::vector<std::pair<int64_t, int64_t>> &hits, float score) {
        ScanResult scan;
        for (int64_t i = 0; i < nx * ny; i++) {
            ScanResultElem elem;
            elem.number = i;
            scan.images.push_back(elem);
        }
        for (const auto &[x, y]: hits)
            scan.images[y * nx + x].protein_score = score;
        return scan;
    }
}

TEST_CASE("AnalyzeGridScan finds nothing in an empty scan", "[AnalyzeGridScan]") {
    auto grid = MakeGrid(6, 6, 10.0f, 10.0f);

    SECTION("no image scored at all") {
        auto result = AnalyzeGridScan(MakeScan(6, 6, {}, 0.0f), grid, 5.0f, 5.0f);
        CHECK(result.crystals.empty());
    }

    SECTION("every image scored, none above threshold") {
        auto scan = MakeScan(6, 6, {}, 0.0f);
        for (auto &elem: scan.images)
            elem.protein_score = 0.1f;
        auto result = AnalyzeGridScan(scan, grid, 5.0f, 5.0f);
        CHECK(result.crystals.empty());
    }

    SECTION("a two-cell patch is below the minimum blob size and its score is not decisive") {
        // Above the admission threshold, so the cells are in a patch at all, but under the score at
        // which one cell would be evidence enough on its own. RasterReportTest covers the other
        // side of that rule - a patch this small kept because its best cell IS decisive.
        auto scan = MakeScan(6, 6, {{2, 2}, {3, 2}}, 0.55f);
        auto result = AnalyzeGridScan(scan, grid, 5.0f, 5.0f);
        CHECK(result.crystals.empty());
    }
}

TEST_CASE("AnalyzeGridScan measures a round blob", "[AnalyzeGridScan]") {
// A plus sign centred on cell (3,3): symmetric, so the centre is the centre cell exactly.
    auto grid = MakeGrid(7, 7, 10.0f, 10.0f);
    auto scan = MakeScan(7, 7, {{3, 2}, {2, 3}, {3, 3}, {4, 3}, {3, 4}}, 0.8f);

    auto result = AnalyzeGridScan(scan, grid, 5.0f, 5.0f);
    REQUIRE(result.crystals.size() == 1);
    const auto &c = result.crystals[0];

    CHECK(c.nx == Catch::Approx(3.0f));
    CHECK(c.ny == Catch::Approx(3.0f));
    CHECK(c.x_um == Catch::Approx(30.0f));
    CHECK(c.y_um == Catch::Approx(30.0f));
    CHECK(c.image_number == 3 * 7 + 3);
    CHECK(c.n_images == 5);

// Score is the mean over the blob, not the peak, and every cell here carries the same 0.8
    CHECK(c.score == Catch::Approx(0.8f));
    CHECK(c.ice_score == Catch::Approx(0.0f));
    CHECK(std::isnan(c.res_A));    // no resolution was measured anywhere

// Round: the two axes come out equal, whatever arbitrary angle the numerics picked
    CHECK(c.major_um == Catch::Approx(c.minor_um).margin(1e-3));
    CHECK(c.angle_deg >= 0.0f);
    CHECK(c.angle_deg < 180.0f);

// The beam is not removed - it is reported so a consumer can remove it itself
    CHECK(result.beam_size_x_um == Catch::Approx(5.0f));
    CHECK(result.beam_size_y_um == Catch::Approx(5.0f));
}

TEST_CASE("AnalyzeGridScan keeps a corner-touching diagonal needle whole", "[AnalyzeGridScan]") {
// A needle lying at 45 degrees across a raster whose step is coarser than the beam is sampled as
// cells that meet only at their corners. 4-connectivity would cut this into four blobs of one
// cell each and the minimum-size rule would then discard the crystal entirely; 8-connectivity
// keeps it as the one oriented needle it is.
    const std::vector<std::pair<int64_t, int64_t>> diagonal = {{1, 1}, {2, 2}, {3, 3}, {4, 4}};

    SECTION("square 20 x 20 um steps") {
        auto grid = MakeGrid(7, 7, 20.0f, 20.0f);
        auto result = AnalyzeGridScan(MakeScan(7, 7, diagonal, 0.9f), grid, 5.0f, 5.0f);
        REQUIRE(result.crystals.size() == 1);
        const auto &c = result.crystals[0];

        CHECK(c.n_images == 4);
        CHECK(c.angle_deg == Catch::Approx(45.0f).margin(0.01));
        CHECK(c.nx == Catch::Approx(2.5f).margin(1e-4));
        CHECK(c.ny == Catch::Approx(2.5f).margin(1e-4));
        CHECK(c.major_um == Catch::Approx(113.137f).margin(0.01));
        CHECK(c.minor_um == Catch::Approx(28.284f).margin(0.01));
    }

    SECTION("anisotropic 20 x 16 um steps") {
        auto grid = MakeGrid(7, 7, 20.0f, 16.0f);
        auto result = AnalyzeGridScan(MakeScan(7, 7, diagonal, 0.9f), grid, 5.0f, 5.0f);
        REQUIRE(result.crystals.size() == 1);
        const auto &c = result.crystals[0];

        CHECK(c.n_images == 4);
// atan(16/20) - the same cells, read on a raster with a shorter y step
        CHECK(c.angle_deg == Catch::Approx(38.6598f).margin(0.01));
        CHECK(c.x_um == Catch::Approx(50.0f).margin(1e-3));
        CHECK(c.y_um == Catch::Approx(40.0f).margin(1e-3));
        CHECK(c.major_um == Catch::Approx(102.450f).margin(0.01));
        CHECK(c.minor_um == Catch::Approx(24.988f).margin(0.01));
    }
}

TEST_CASE("AnalyzeGridScan takes the needle angle from micrometres, not cells", "[AnalyzeGridScan]") {
// The SAME five cells - a staircase running up and to the right - read on two rasters that
// differ only in their y step. A purely diagonal line would not do: it is not 4-connected, and
// a real needle is sampled as a staircase anyway, the beam being wider than one step.
    const std::vector<std::pair<int64_t, int64_t>> staircase = {{1, 1}, {2, 1}, {2, 2}, {3, 2}, {3, 3}};

    SECTION("square 20 x 20 um steps") {
        auto grid = MakeGrid(7, 7, 20.0f, 20.0f);
        auto result = AnalyzeGridScan(MakeScan(7, 7, staircase, 0.9f), grid, 5.0f, 5.0f);
        REQUIRE(result.crystals.size() == 1);
        const auto &c = result.crystals[0];

// Symmetric about the diagonal, so the axis is at exactly 45 degrees
        CHECK(c.angle_deg == Catch::Approx(45.0f).margin(0.01));
        CHECK(c.nx == Catch::Approx(2.2f).margin(1e-4));
        CHECK(c.ny == Catch::Approx(1.8f).margin(1e-4));
        CHECK(c.x_um == Catch::Approx(44.0f).margin(1e-3));
        CHECK(c.y_um == Catch::Approx(36.0f).margin(1e-3));
        CHECK(c.major_um == Catch::Approx(84.853f).margin(0.01));
        CHECK(c.minor_um == Catch::Approx(42.426f).margin(0.01));
    }

    SECTION("anisotropic 20 x 16 um steps") {
        auto grid = MakeGrid(7, 7, 20.0f, 16.0f);
        auto result = AnalyzeGridScan(MakeScan(7, 7, staircase, 0.9f), grid, 5.0f, 5.0f);
        REQUIRE(result.crystals.size() == 1);
        const auto &c = result.crystals[0];

// The cells are identical, so second moments taken in CELL units would still say 45 degrees.
// In micrometres the shorter y step tilts the axis down to 37.01 - eight degrees away, which
// is what makes this assertion worth making.
        CHECK(c.angle_deg == Catch::Approx(37.010f).margin(0.01));
        CHECK(c.angle_deg < 44.0f);

// The centre sits a fifth of a cell off a grid node in both directions
        CHECK(c.nx == Catch::Approx(2.2f).margin(1e-4));
        CHECK(c.ny == Catch::Approx(1.8f).margin(1e-4));
        CHECK(c.x_um == Catch::Approx(44.0f).margin(1e-3));
        CHECK(c.y_um == Catch::Approx(28.8f).margin(1e-3));

        CHECK(c.major_um == Catch::Approx(76.806f).margin(0.01));
        CHECK(c.minor_um == Catch::Approx(38.329f).margin(0.01));
        CHECK(c.major_um > c.minor_um);
    }
}

TEST_CASE("AnalyzeGridScan snaps a centre that falls outside the blob", "[AnalyzeGridScan]") {
// An L: a vertical arm at column 0 (rows 0..5) and a horizontal arm along row 5 (columns 1..3).
// The centroid of that is (0.667, 3.333), which rounds to cell (1,3) - a cell that is NOT in
// the blob and was never collected as part of this crystal.
    auto grid = MakeGrid(6, 6, 10.0f, 10.0f);
    auto scan = MakeScan(6, 6, {{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5},
                                {1, 5}, {2, 5}, {3, 5}}, 0.7f);

    auto result = AnalyzeGridScan(scan, grid, 5.0f, 5.0f);
    REQUIRE(result.crystals.size() == 1);
    const auto &c = result.crystals[0];

    CHECK(c.nx == Catch::Approx(6.0f / 9.0f).margin(1e-4));
    CHECK(c.ny == Catch::Approx(30.0f / 9.0f).margin(1e-4));
    CHECK(c.n_images == 9);

// The reported image is the nearest cell that is actually in the blob, (0,3), not (1,3)
    CHECK(c.image_number == 3 * 6 + 0);
}

TEST_CASE("AnalyzeGridScan centres on resolution rank, not resolution value", "[AnalyzeGridScan]") {
// A horizontal bar of five cells. The best resolution sits at one end, so it pulls the centre
// off the geometric middle - but by exactly the same amount whether that end reports a
// plausible 2.0 A or an absurd 0.8 A, because only the rank is used.
    auto grid = MakeGrid(7, 3, 10.0f, 10.0f);
    const std::vector<std::pair<int64_t, int64_t>> bar = {{1, 1}, {2, 1}, {3, 1}, {4, 1}, {5, 1}};

    auto with_best_at_left = [&](float best_res) {
        auto scan = MakeScan(7, 3, bar, 0.6f);
        const float res[5] = {best_res, 3.0f, 3.5f, 4.0f, 4.5f};
        for (int i = 0; i < 5; i++)
            scan.images[1 * 7 + 1 + i].res = res[i];
        return AnalyzeGridScan(scan, grid, 5.0f, 5.0f).crystals.at(0);
    };

    const auto plausible = with_best_at_left(2.0f);
    const auto artefact = with_best_at_left(0.8f);

// Pulled towards the good end, so left of the geometric centre at 3.0
    CHECK(plausible.nx < 3.0f);
// ...and the absurd number moves it not one bit
    CHECK(artefact.nx == Catch::Approx(plausible.nx));
    CHECK(artefact.ny == Catch::Approx(plausible.ny));

// res_A is the 25th percentile of the five, i.e. the second best (3.0), never the best
    CHECK(plausible.res_A == Catch::Approx(3.0f));
    CHECK(artefact.res_A == Catch::Approx(3.0f));
}

TEST_CASE("AnalyzeGridScan separates two crystals and sorts them by score", "[AnalyzeGridScan]") {
// Two blobs with a clear gap between them; the weaker one is collected first. Labelling is
// 8-connected, so the gap has to be more than one cell in EVERY direction, diagonals included -
// three empty columns here.
    auto grid = MakeGrid(9, 5, 10.0f, 10.0f);
    auto scan = MakeScan(9, 5, {}, 0.0f);

    for (auto [x, y]: std::vector<std::pair<int64_t, int64_t>>{{1, 1}, {1, 2}, {2, 1}, {2, 2}})
        scan.images[y * 9 + x].protein_score = 0.6f;
    for (auto [x, y]: std::vector<std::pair<int64_t, int64_t>>{{6, 2}, {7, 2}, {6, 3}})
        scan.images[y * 9 + x].protein_score = 0.95f;

    scan.images[2 * 9 + 6].ice_score = 0.4f;

    auto result = AnalyzeGridScan(scan, grid, 5.0f, 5.0f);
    REQUIRE(result.crystals.size() == 2);

// Element 0 is the one to collect
    CHECK(result.crystals[0].score == Catch::Approx(0.95f));
    CHECK(result.crystals[0].n_images == 3);
    CHECK(result.crystals[0].nx > 4.0f);
    CHECK(result.crystals[0].ice_score == Catch::Approx(0.4f / 3.0f));

    CHECK(result.crystals[1].score == Catch::Approx(0.6f));
    CHECK(result.crystals[1].n_images == 4);
    CHECK(result.crystals[1].nx < 4.0f);
    CHECK(result.crystals[1].ice_score == Catch::Approx(0.0f));
}

TEST_CASE("AnalyzeGridScan reports major_um along angle_deg", "[AnalyzeGridScan]") {
// A hook: two cells down, a step across, one more down. On a 20 x 10 um raster the axis of the
// larger second MOMENT runs at 58.3 degrees, but the blob's measured extent along that axis
// (46.5 um) is SHORTER than across it (49.8 um) - the moment is dominated by two clumps at the
// ends, the extent is not. A consumer draws a major by minor frame rotated by angle_deg, so the
// two facts are kept by turning the frame a quarter turn, not by discarding one of them.
    auto grid = MakeGrid(5, 5, 20.0f, 10.0f);
    auto scan = MakeScan(5, 5, {{0, 0}, {0, 1}, {1, 1}, {1, 2}, {0, 3}}, 0.8f);

    auto result = AnalyzeGridScan(scan, grid, 5.0f, 5.0f);
    REQUIRE(result.crystals.size() == 1);
    const auto &c = result.crystals[0];

    CHECK(c.n_images == 5);
    CHECK(c.major_um == Catch::Approx(49.798f).margin(0.01));
    CHECK(c.minor_um == Catch::Approx(46.549f).margin(0.01));
// 58.28 + 90, folded back into [0,180)
    CHECK(c.angle_deg == Catch::Approx(148.283f).margin(0.01));
}

TEST_CASE("AnalyzeGridScan keeps major_um >= minor_um for every blob", "[AnalyzeGridScan]") {
// The invariant a consumer relies on, over every shape the other cases exercise.
    const std::vector<std::pair<int64_t, int64_t>> shapes[] = {
            {{3, 2}, {2, 3}, {3, 3}, {4, 3}, {3, 4}},                     // round
            {{1, 1}, {2, 2}, {3, 3}, {4, 4}},                             // diagonal needle
            {{1, 1}, {2, 1}, {2, 2}, {3, 2}, {3, 3}},                     // staircase needle
            {{0, 0}, {0, 1}, {1, 1}, {1, 2}, {0, 3}},                     // hook
            {{0, 0}, {0, 1}, {0, 2}, {0, 3}, {0, 4}, {0, 5},
             {1, 5}, {2, 5}, {3, 5}},                                     // L
            {{1, 1}, {2, 1}, {3, 1}, {4, 1}, {5, 1}},                     // bar
    };

    for (const auto &shape: shapes) {
        for (auto [sx, sy]: std::vector<std::pair<float, float>>{{10.0f, 10.0f},
                                                                 {20.0f, 16.0f},
                                                                 {20.0f, 10.0f},
                                                                 {10.0f, 20.0f}}) {
            auto grid = MakeGrid(7, 7, sx, sy);
            auto result = AnalyzeGridScan(MakeScan(7, 7, shape, 0.8f), grid, 5.0f, 5.0f);
            REQUIRE(result.crystals.size() == 1);
            const auto &c = result.crystals[0];

            CHECK(c.major_um >= c.minor_um);
            CHECK(c.angle_deg >= 0.0f);
            CHECK(c.angle_deg < 180.0f);
        }
    }
}

TEST_CASE("AnalyzeGridScan follows the snake and the step signs", "[AnalyzeGridScan]") {
// The same crystal, described once on a plain raster and once on a snake raster with a
// negative slow step. Rearrange puts both on the same display grid, so both have to report
// the same position - only the image numbers differ.
    const int64_t nx = 6, ny = 6;
    GridScanSettings plain(nx, 10.0f, 10.0f, false, false);
    plain.ImageNum(nx * ny);
    GridScanSettings snake(nx, 10.0f, -10.0f, true, false);
    snake.ImageNum(nx * ny);

    const std::vector<std::pair<int64_t, int64_t>> blob = {{2, 1}, {3, 1}, {2, 2}, {3, 2}};

    auto scan_plain = MakeScan(nx, ny, blob, 0.8f);
    ScanResult scan_snake;
    for (int64_t i = 0; i < nx * ny; i++) {
        ScanResultElem elem;
        elem.number = i;
        elem.protein_score = scan_plain.images[snake.Rearrange(i)].protein_score;
        scan_snake.images.push_back(elem);
    }

    auto a = AnalyzeGridScan(scan_plain, plain, 5.0f, 5.0f).crystals.at(0);
    auto b = AnalyzeGridScan(scan_snake, snake, 5.0f, 5.0f).crystals.at(0);

    CHECK(a.nx == Catch::Approx(b.nx));
    CHECK(a.ny == Catch::Approx(b.ny));
    CHECK(a.x_um == Catch::Approx(b.x_um));
    CHECK(a.y_um == Catch::Approx(b.y_um));
    CHECK(snake.Rearrange(b.image_number) == plain.Rearrange(a.image_number));
}

TEST_CASE("AnalyzeGridScan falls back to the grid step for an unstated beam", "[AnalyzeGridScan]") {
    auto grid = MakeGrid(8, 8, 20.0f, 16.0f);
    const std::vector<std::pair<int64_t, int64_t>> blob{{2, 2}, {3, 2}, {2, 3}, {3, 3}};
    const auto scan = MakeScan(8, 8, blob, 0.9f);

    // Stated: echoed back untouched, whatever the step is.
    const auto stated = AnalyzeGridScan(scan, grid, 80.0f, 20.0f);
    CHECK(stated.beam_size_x_um == Catch::Approx(80.0));
    CHECK(stated.beam_size_y_um == Catch::Approx(20.0));

    // Unstated: the step stands in. Zero would tell a consumer the extents are exact when they
    // still contain a whole beam, and the two axes must not collapse to one value.
    const auto unstated = AnalyzeGridScan(scan, grid, 0.0f, 0.0f);
    CHECK(unstated.beam_size_x_um == Catch::Approx(20.0));
    CHECK(unstated.beam_size_y_um == Catch::Approx(16.0));

    // The extents themselves are measured and must not move with the beam.
    REQUIRE(stated.crystals.size() == 1);
    REQUIRE(unstated.crystals.size() == 1);
    CHECK(stated.crystals[0].major_um == Catch::Approx(unstated.crystals[0].major_um));
    CHECK(stated.crystals[0].minor_um == Catch::Approx(unstated.crystals[0].minor_um));
}