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

#include <tuple>

#include <catch2/catch_all.hpp>
#include <nlohmann/json.hpp>

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

// A 5 x 4 raster with one 2 x 2 patch of diffracting cells in it - the smallest scan that has a
// crystal to report and cells around it that are not one.
namespace {
    constexpr int64_t GRID_NX = 5;
    constexpr int64_t GRID_NY = 4;

    GridScanSettings TestGrid() {
        GridScanSettings grid(GRID_NX, 20.0f, 16.0f, /*snake=*/false, /*vertical=*/false);
        grid.ImageNum(GRID_NX * GRID_NY);
        return grid;
    }

    ScanResult TestScan() {
        ScanResult scan;
        for (int64_t i = 0; i < GRID_NX * GRID_NY; i++) {
            const bool crystal = (i == 6 || i == 7 || i == 11 || i == 12);
            ScanResultElem elem;
            elem.number = i;
            elem.protein_score = crystal ? 0.9f : 0.1f;
            elem.ice_score = 0.05f;
            if (crystal)
                elem.res = 2.0f;
            scan.images.push_back(elem);
        }
        return scan;
    }

    RasterSettings TestSettings() {
        RasterSettings settings;
        settings.beam_size_x_um = 30.0f;
        settings.beam_size_y_um = 10.0f;
        settings.beam_size_source = "COMMAND_LINE";
        return settings;
    }
}

TEST_CASE("RasterReport_Render", "[Diagnostics]") {
    const GridScanSettings grid = TestGrid();
    const ScanResult scan = TestScan();
    const RasterSettings settings = TestSettings();
    const GridScanResult crystals = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                    settings.beam_size_y_um, settings.analysis);
    REQUIRE(crystals.crystals.size() == 1);

    const std::string report = RenderRasterReport("test_master.h5", grid, scan, crystals, settings);

    // The keys are the interface a battery greps; pin the spellings and the values that say what
    // the run saw.
    CHECK(report.find("RASTER_REPORT_VERSION= 3\n") != std::string::npos);
    CHECK(report.find("INPUT_FILE= test_master.h5\n") != std::string::npos);
    CHECK(report.find("GRID_SIZE= 5 4\n") != std::string::npos);
    CHECK(report.find("GRID_STEP_UM= 20.00 16.00\n") != std::string::npos);
    CHECK(report.find("GRID_N_ELEM= 20\n") != std::string::npos);
    CHECK(report.find("BEAM_SIZE_UM= 30.00 10.00\n") != std::string::npos);
    CHECK(report.find("BEAM_SIZE_SOURCE= COMMAND_LINE\n") != std::string::npos);
    CHECK(report.find("IMAGES_SCORED= 20\n") != std::string::npos);
    CHECK(report.find("IMAGES_ABOVE_THRESHOLD= 4\n") != std::string::npos);
    CHECK(report.find("ICE_ABOVE_THRESHOLD= 0\n") != std::string::npos);
    CHECK(report.find("GROW_SCORE_THRESHOLD= 0.35\n") != std::string::npos);
    CHECK(report.find("MIN_BLOB_CELLS= 3\n") != std::string::npos);
    CHECK(report.find("DECISIVE_PROTEIN_SCORE= 0.60\n") != std::string::npos);
    // Unset means no cap; "0" would read as "report no crystals".
    CHECK(report.find("MAX_CRYSTALS= none\n") != std::string::npos);
    CHECK(report.find("CRYSTAL_COUNT= 1\n") != std::string::npos);
    CHECK(report.find("END OF REPORT") != std::string::npos);
}

TEST_CASE("RasterReport_Json", "[Diagnostics]") {
    const GridScanSettings grid = TestGrid();
    const ScanResult scan = TestScan();
    const RasterSettings settings = TestSettings();
    const GridScanResult crystals = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                    settings.beam_size_y_um, settings.analysis);

    const nlohmann::json j = nlohmann::json::parse(
            RenderRasterJson("test_master.h5", grid, scan, crystals, settings));

    CHECK(j["grid"]["size_x"] == GRID_NX);
    CHECK(j["grid"]["size_y"] == GRID_NY);
    CHECK(j["scores"]["images_above_threshold"] == 4);
    CHECK(j["settings"]["beam_size_x_um"].get<float>() == Catch::Approx(30.0));
    CHECK(j["settings"]["beam_size_y_um"].get<float>() == Catch::Approx(10.0));
    REQUIRE(j["crystal_count"] == 1);
    REQUIRE(j["crystals"].size() == 1);

    // The crystal the JSON carries is the one the report's table carries - same numbers, typed.
    const auto &c = j["crystals"][0];
    CHECK(c["n_images"] == 4);
    CHECK(c["image_number"] == crystals.crystals[0].image_number);
    CHECK(c["res_A"].get<float>() == Catch::Approx(2.0));
    CHECK(c["score"].get<float>() == Catch::Approx(0.9).margin(1e-5));
    CHECK(c["peak_score"].get<float>() == Catch::Approx(0.9).margin(1e-5));
}

TEST_CASE("RasterReport_NoCrystal", "[Diagnostics]") {
    // A raster over an empty loop: the report still has to say so, with a table that is empty rather
    // than absent - a battery that greps CRYSTAL_COUNT must find it on every run.
    const GridScanSettings grid = TestGrid();
    ScanResult scan = TestScan();
    for (auto &elem : scan.images)
        elem.protein_score = 0.1f;

    const RasterSettings settings = TestSettings();
    const GridScanResult crystals = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                    settings.beam_size_y_um, settings.analysis);
    CHECK(crystals.crystals.empty());

    const std::string report = RenderRasterReport("test_master.h5", grid, scan, crystals, settings);
    CHECK(report.find("CRYSTAL_COUNT= 0\n") != std::string::npos);
    CHECK(report.find("IMAGES_ABOVE_THRESHOLD= 0\n") != std::string::npos);

    const nlohmann::json j = nlohmann::json::parse(
            RenderRasterJson("test_master.h5", grid, scan, crystals, settings));
    CHECK(j["crystal_count"] == 0);
    CHECK(j["crystals"].empty());
}

// The blob search returns every patch it finds, so a raster with two crystals in it has to come back
// with two rows in score order - and a caller that only wants the best few has to be able to say so.
namespace {
    // A wider grid, so several crystals fit in it with empty cells between them.
    constexpr int64_t WIDE_NX = 12;
    constexpr int64_t WIDE_NY = 5;

    GridScanSettings WideGrid() {
        GridScanSettings grid(WIDE_NX, 20.0f, 16.0f, /*snake=*/false, /*vertical=*/false);
        grid.ImageNum(WIDE_NX * WIDE_NY);
        return grid;
    }

    // Each patch is a 2 x 2 block whose top-left cell is (x, y), all of its cells at score.
    ScanResult WideScan(const std::vector<std::tuple<int64_t, int64_t, float>> &patches) {
        ScanResult scan;
        for (int64_t i = 0; i < WIDE_NX * WIDE_NY; i++) {
            ScanResultElem elem;
            elem.number = i;
            elem.protein_score = 0.1f;
            scan.images.push_back(elem);
        }
        for (const auto &[x, y, score] : patches) {
            for (int64_t dy = 0; dy < 2; dy++) {
                for (int64_t dx = 0; dx < 2; dx++)
                    scan.images[(y + dy) * WIDE_NX + x + dx].protein_score = score;
            }
        }
        return scan;
    }
}

TEST_CASE("RasterReport_TwoCrystals", "[Diagnostics]") {
    const GridScanSettings grid = WideGrid();
    const ScanResult scan = WideScan({{1, 1, 0.9f}, {8, 2, 0.7f}});
    RasterSettings settings = TestSettings();

    const GridScanResult crystals = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                    settings.beam_size_y_um, settings.analysis);
    REQUIRE(crystals.crystals.size() == 2);
    // Best first, and the two are far enough apart that neither took the other's cells.
    CHECK(crystals.crystals[0].score == Catch::Approx(0.9).margin(1e-5));
    CHECK(crystals.crystals[1].score == Catch::Approx(0.7).margin(1e-5));
    CHECK(crystals.crystals[0].n_images == 4);
    CHECK(crystals.crystals[1].n_images == 4);
    CHECK(crystals.crystals[0].nx < crystals.crystals[1].nx);

    const std::string report = RenderRasterReport("test_master.h5", grid, scan, crystals, settings);
    CHECK(report.find("CRYSTAL_COUNT= 2\n") != std::string::npos);

    const nlohmann::json j = nlohmann::json::parse(
            RenderRasterJson("test_master.h5", grid, scan, crystals, settings));
    REQUIRE(j["crystals"].size() == 2);
    CHECK(j["crystals"][0]["score"].get<float>() > j["crystals"][1]["score"].get<float>());
}

TEST_CASE("RasterReport_ThreeCrystalsAndTheCap", "[Diagnostics]") {
    const GridScanSettings grid = WideGrid();
    const ScanResult scan = WideScan({{0, 0, 0.95f}, {5, 2, 0.75f}, {10, 0, 0.85f}});
    RasterSettings settings = TestSettings();

    const GridScanResult all = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                               settings.beam_size_y_um, settings.analysis);
    REQUIRE(all.crystals.size() == 3);
    CHECK(all.crystals[0].score == Catch::Approx(0.95).margin(1e-5));
    CHECK(all.crystals[1].score == Catch::Approx(0.85).margin(1e-5));
    CHECK(all.crystals[2].score == Catch::Approx(0.75).margin(1e-5));

    // The cap keeps the strongest, not the first found: the 0.85 patch is the last one in grid order.
    settings.analysis.MaxCrystals(2);
    const GridScanResult capped = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                  settings.beam_size_y_um, settings.analysis);
    REQUIRE(capped.crystals.size() == 2);
    CHECK(capped.crystals[0].score == Catch::Approx(0.95).margin(1e-5));
    CHECK(capped.crystals[1].score == Catch::Approx(0.85).margin(1e-5));

    const std::string report = RenderRasterReport("test_master.h5", grid, scan, capped, settings);
    CHECK(report.find("MAX_CRYSTALS= 2\n") != std::string::npos);
    CHECK(report.find("CRYSTAL_COUNT= 2\n") != std::string::npos);
}

TEST_CASE("RasterReport_OneDecisiveCellIsACrystal", "[Diagnostics]") {
    // A patch below MIN_BLOB_CELLS is kept when its best cell is decisive and dropped when it is not.
    // Both patches here are one cell, so the only thing separating them is the evidence in them.
    const GridScanSettings grid = WideGrid();
    ScanResult scan = WideScan({});
    scan.images[1 * WIDE_NX + 2].protein_score = 0.95f;   // decisive
    scan.images[3 * WIDE_NX + 9].protein_score = 0.55f;   // above the admission threshold, not decisive

    RasterSettings settings = TestSettings();
    const GridScanResult crystals = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                    settings.beam_size_y_um, settings.analysis);
    REQUIRE(crystals.crystals.size() == 1);
    CHECK(crystals.crystals[0].n_images == 1);
    CHECK(crystals.crystals[0].peak_score == Catch::Approx(0.95).margin(1e-5));

    // Raise the bar past the one cell that passed and nothing is left; a patch that meets the cell
    // minimum is unaffected by the bar, which is what makes this an OR and not a second gate.
    settings.analysis.DecisiveSingleCellScore(0.99f);
    CHECK(AnalyzeGridScan(scan, grid, settings.beam_size_x_um, settings.beam_size_y_um,
                          settings.analysis).crystals.empty());

    const ScanResult big = WideScan({{1, 1, 0.55f}});
    CHECK(AnalyzeGridScan(big, grid, settings.beam_size_x_um, settings.beam_size_y_um,
                          settings.analysis).crystals.size() == 1);
}

TEST_CASE("RasterReport_HysteresisHealsASplit", "[Diagnostics]") {
    // The case this exists for, taken from the corpus: one crystal whose middle grid point fell two
    // thousandths under the seed threshold, which broke it into two crystals that were never two.
    const GridScanSettings grid = WideGrid();
    ScanResult scan = WideScan({});
    for (int64_t x = 2; x <= 8; x++)
        scan.images[2 * WIDE_NX + x].protein_score = 0.85f;
    scan.images[2 * WIDE_NX + 5].protein_score = 0.498f;   // above the grow level, under the seed

    RasterSettings settings = TestSettings();
    const GridScanResult healed = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                  settings.beam_size_y_um, settings.analysis);
    REQUIRE(healed.crystals.size() == 1);
    CHECK(healed.crystals[0].n_images == 7);   // the bridging cell is part of the crystal

    // Growing no further than the seed level is the behaviour hysteresis replaced, and it splits.
    settings.analysis.GrowScoreThreshold(settings.analysis.GetProteinScoreThreshold());
    const GridScanResult split = AnalyzeGridScan(scan, grid, settings.beam_size_x_um,
                                                 settings.beam_size_y_um, settings.analysis);
    CHECK(split.crystals.size() == 2);
}

TEST_CASE("RasterReport_GrowthCannotStartOnItsOwn", "[Diagnostics]") {
    // A patch that reaches the grow level and never the seed level is not a crystal however large it
    // is or however low the grow level is set - which is what makes the level safe to lower.
    const GridScanSettings grid = WideGrid();
    const ScanResult scan = WideScan({{1, 1, 0.45f}, {8, 2, 0.40f}});

    RasterSettings settings = TestSettings();
    CHECK(AnalyzeGridScan(scan, grid, settings.beam_size_x_um, settings.beam_size_y_um,
                          settings.analysis).crystals.empty());

    settings.analysis.GrowScoreThreshold(0.05f);
    CHECK(AnalyzeGridScan(scan, grid, settings.beam_size_x_um, settings.beam_size_y_um,
                          settings.analysis).crystals.empty());
}