// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include <catch2/catch_test_macros.hpp>
#include <catch2/matchers/catch_matchers_floating_point.hpp>

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

using namespace grid_scan;

TEST_CASE("GridScanComposite_Background", "[grid_scan]") {
    // Air only: a flat level with counting noise on it. Without a floor on the range this would
    // be stretched to fill the whole grey scale and look like a loop.
    std::vector<float> air;
    for (int i = 0; i < 100; i++)
        air.push_back(5.0f + 0.02f * ((i % 7) - 3));

    const auto flat = NormaliseBackground(air);
    for (float v : flat)
        CHECK(v < 0.1f);

    // A loop in the middle of the same air.
    std::vector<float> loop = air;
    for (int i = 40; i < 60; i++)
        loop[i] = 13.0f;

    const auto mapped = NormaliseBackground(loop);
    CHECK(mapped[0] < 0.1f);
    CHECK(mapped[50] > 0.9f);
}

TEST_CASE("GridScanComposite_BackgroundKeepsGaps", "[grid_scan]") {
    const std::vector<float> bkg{5.0f, std::numeric_limits<float>::quiet_NaN(), 13.0f};
    const auto out = NormaliseBackground(bkg);
    REQUIRE(out.size() == 3);
    CHECK(std::isnan(out[1]));
    CHECK(std::isfinite(out[0]));
    CHECK(std::isfinite(out[2]));
}

TEST_CASE("GridScanComposite_Color", "[grid_scan]") {
    // Nothing anywhere: the pale plate.
    const rgb air = CompositeColor(0.0f, 0.0f, 0.0f);
    CHECK(air.r == kPlate.r);
    CHECK(air.g == kPlate.g);
    CHECK(air.b == kPlate.b);

    // Material but no diffraction: the loop grey.
    const rgb loop = CompositeColor(1.0f, 0.0f, 0.0f);
    CHECK(loop.r == kLoop.r);
    CHECK(loop.b == kLoop.b);

    // A saturated score paints its own colour whatever the background says.
    const rgb protein = CompositeColor(1.0f, 1.0f, 0.0f);
    CHECK(protein.r == kProtein.r);
    CHECK(protein.g == kProtein.g);
    CHECK(protein.b == kProtein.b);

    const rgb ice = CompositeColor(0.0f, 0.0f, 1.0f);
    CHECK(ice.r == kIce.r);
    CHECK(ice.b == kIce.b);

    // Equal parts of both sit half way between the two, not on either.
    const rgb both = CompositeColor(1.0f, 0.5f, 0.5f);
    CHECK(both.r == static_cast<uint8_t>(std::lround(0.5f * (kProtein.r + kIce.r))));
    CHECK(both.b == static_cast<uint8_t>(std::lround(0.5f * (kProtein.b + kIce.b))));

    // Protein three times ice: still mostly orange, but on the way to cyan.
    const rgb mostly_protein = CompositeColor(1.0f, 0.75f, 0.25f);
    CHECK(mostly_protein.b > protein.b);
    CHECK(mostly_protein.b < both.b);

    // A weak score is a weak tint over the base, not a full-strength colour.
    const rgb faint = CompositeColor(1.0f, 0.1f, 0.0f);
    CHECK(faint.r > loop.r);
    CHECK(faint.b < loop.b);
    CHECK(faint.b > kProtein.b);
}

TEST_CASE("GridScanComposite_FrameGeometry", "[grid_scan]") {
    // Steps differ, as they do in a real raster: 20 um across, 16 um down.
    constexpr float step_x = 20.0f;
    constexpr float step_y = 16.0f;

    // Along the grid x axis: 100 um is 5 cells across, 10 um is 0.625 cells down.
    const auto along_x = CrystalFrameCorners(4.0f, 3.0f, 100.0f, 10.0f, 0.0f, step_x, step_y);
    CHECK_THAT(along_x[0].x, Catch::Matchers::WithinAbs(4.5 + 2.5, 1e-4));
    CHECK_THAT(along_x[0].y, Catch::Matchers::WithinAbs(3.5 + 0.3125, 1e-4));
    CHECK_THAT(along_x[2].x, Catch::Matchers::WithinAbs(4.5 - 2.5, 1e-4));

    // The same needle at 90 degrees: the long axis is now along y, and because the y step is the
    // finer one it covers more cells than it did across - 100/16 rather than 100/20.
    const auto along_y = CrystalFrameCorners(4.0f, 3.0f, 100.0f, 10.0f, 90.0f, step_x, step_y);
    float min_x = along_y[0].x, max_x = along_y[0].x, min_y = along_y[0].y, max_y = along_y[0].y;
    for (const auto &c : along_y) {
        min_x = std::min(min_x, c.x);
        max_x = std::max(max_x, c.x);
        min_y = std::min(min_y, c.y);
        max_y = std::max(max_y, c.y);
    }
    CHECK_THAT(max_x - min_x, Catch::Matchers::WithinAbs(10.0 / step_x, 1e-3));
    CHECK_THAT(max_y - min_y, Catch::Matchers::WithinAbs(100.0 / step_y, 1e-3));

    // At 45 degrees the frame is still a 100 x 10 needle - its axis-aligned bounding box is the
    // ~78 x 78 blob the oriented frame exists to avoid drawing.
    const auto tilted = CrystalFrameCorners(4.0f, 3.0f, 100.0f, 10.0f, 45.0f, step_x, step_x);
    const float long_side = std::hypot(tilted[0].x - tilted[1].x, tilted[0].y - tilted[1].y);
    const float short_side = std::hypot(tilted[1].x - tilted[2].x, tilted[1].y - tilted[2].y);
    CHECK_THAT(long_side, Catch::Matchers::WithinAbs(100.0 / step_x, 1e-3));
    CHECK_THAT(short_side, Catch::Matchers::WithinAbs(10.0 / step_x, 1e-3));
    float bbox = 0;
    for (const auto &c : tilted)
        bbox = std::max(bbox, std::abs(c.x - 4.5f));
    CHECK(bbox > 1.9f);   // the bounding box is far wider than the needle itself
}

TEST_CASE("GridScanComposite_RoundBlob", "[grid_scan]") {
    CHECK(IsRoundBlob(30.0f, 30.0f));
    CHECK(IsRoundBlob(33.0f, 30.0f));
    CHECK(!IsRoundBlob(100.0f, 10.0f));
    CHECK(!IsRoundBlob(40.0f, 30.0f));
}

namespace {
    // A plain raster: positive steps, no snake, so image number == grid index.
    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;
    }

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

// ---------------------------------------------------------------------------
// The producer measures angle_deg from second moments in micrometres; the viewer
// draws it counter-clockwise in the grid index frame. If those two conventions
// disagree every needle is drawn mirrored, and NEITHER side's own tests notice: the
// producer checks an angle it computed itself, the viewer checks a frame from an
// angle it was handed, and the corpus needle sits at 45 degrees, which is invariant
// under exactly that flip. This is the test that spans the seam, so it uses a
// deliberately asymmetric angle on an anisotropic step.
TEST_CASE("GridScanComposite_FrameFollowsTheNeedleAcrossTheSeam", "[grid_scan]") {
    // A staircase running lower-left to upper-right: nx and ny increase together.
    const std::vector<std::pair<int64_t, int64_t>> needle{{1, 1}, {2, 2}, {3, 2}, {4, 3}, {5, 4}};
    auto grid = MakeGrid(8, 8, 20.0f, 16.0f);
    const auto result = AnalyzeGridScan(MakeScan(8, 8, needle, 0.9f), grid, 0.0f, 0.0f);
    REQUIRE(result.crystals.size() == 1);
    const auto &c = result.crystals.front();
    REQUIRE(c.major_um > c.minor_um);          // it is a needle, not a blob

    const auto corners = CrystalFrameCorners(c.nx, c.ny, c.major_um, c.minor_um, c.angle_deg,
                                             grid.GetGridStepX_um(), grid.GetGridStepY_um());

    // The two corners furthest apart span the frame's long axis. In cell coordinates that axis
    // must rise as it runs right, exactly like the cells it was measured from - a mirrored
    // convention would give it the opposite slope while leaving its length untouched.
    int a = 0, b = 0;
    float longest = -1.0f;
    for (int i = 0; i < 4; i++)
        for (int j = i + 1; j < 4; j++) {
            const float d = std::hypot(corners[i].x - corners[j].x, corners[i].y - corners[j].y);
            if (d > longest) { longest = d; a = i; b = j; }
        }
    const float run  = corners[b].x - corners[a].x;
    const float rise = corners[b].y - corners[a].y;
    CHECK(run * rise > 0.0f);
}