// 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 "../common/AnalysisSettings.h"
#include "../common/DiffractionExperiment.h"
#include "../frame_serialize/CBORStream2Serializer.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../writer/FileWriter.h"
#include "../reader/JFJochHDF5Reader.h"

TEST_CASE("AnalysisMode_Default", "[AnalysisMode]") {
    // MXStills, not None: a deployment whose configuration predates the mode must keep analysing.
    CHECK(AnalysisSettings().GetMode() == AnalysisMode::MXStills);
    CHECK(DiffractionExperiment().GetAnalysisMode() == AnalysisMode::MXStills);
}

TEST_CASE("AnalysisMode_Names", "[AnalysisMode]") {
    for (auto mode : {AnalysisMode::None, AnalysisMode::MXRotation, AnalysisMode::MXStills,
                      AnalysisMode::Azint, AnalysisMode::Grid, AnalysisMode::PowderCalibration}) {
        auto back = AnalysisModeFromName(AnalysisModeName(mode));
        REQUIRE(back.has_value());
        CHECK(*back == mode);
    }
    CHECK_FALSE(AnalysisModeFromName("").has_value());
    CHECK_FALSE(AnalysisModeFromName("mx").has_value());
}

TEST_CASE("AnalysisMode_Stages", "[AnalysisMode]") {
    // The table itself. Bragg integration must never be on where indexing is off - nothing is
    // predicted without a lattice.
    for (auto mode : {AnalysisMode::None, AnalysisMode::MXRotation, AnalysisMode::MXStills,
                      AnalysisMode::Azint, AnalysisMode::Grid, AnalysisMode::PowderCalibration}) {
        const auto s = AnalysisModeStages(mode);
        CHECK((!s.bragg_integration || s.indexing));
        CHECK((!s.scale_merge || s.bragg_integration));
    }

    CHECK(AnalysisModeStages(AnalysisMode::MXStills).indexing);
    CHECK(AnalysisModeStages(AnalysisMode::MXRotation).indexing);

    const auto none = AnalysisModeStages(AnalysisMode::None);
    CHECK_FALSE(none.spot_finding);
    CHECK_FALSE(none.azimuthal_integration);

    const auto azint = AnalysisModeStages(AnalysisMode::Azint);
    CHECK_FALSE(azint.spot_finding);
    CHECK_FALSE(azint.indexing);
    CHECK_FALSE(azint.scoring);
    CHECK(azint.azimuthal_integration);

    // Calibration keeps the spot finder: --calibration spots fits the pooled spots.
    const auto calib = AnalysisModeStages(AnalysisMode::PowderCalibration);
    CHECK(calib.spot_finding);
    CHECK_FALSE(calib.indexing);

    // Grid indexing is a setting, so the table only carries its default.
    const auto grid = AnalysisModeStages(AnalysisMode::Grid);
    CHECK(grid.spot_finding);
    CHECK(grid.scoring);
    CHECK(grid.indexing == GridScanAnalysisSettings().IsIndexing());
}

TEST_CASE("AnalysisMode_GridIndexingIsASetting", "[AnalysisMode]") {
    // The one stage the mode does not fix. The configured value must reach the gates, which read
    // GetAnalysisStages and not the table.
    DiffractionExperiment experiment;
    experiment.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::Grid));
    CHECK(experiment.GetAnalysisStages().indexing);          // the default is on

    experiment.ImportGridScanAnalysisSettings(GridScanAnalysisSettings().Indexing(false));
    CHECK_FALSE(experiment.GetAnalysisStages().indexing);

    experiment.ImportGridScanAnalysisSettings(GridScanAnalysisSettings().Indexing(true));
    CHECK(experiment.GetAnalysisStages().indexing);

    // ...and only under Grid. No other mode's indexing answer moves with it.
    experiment.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::Azint));
    experiment.ImportGridScanAnalysisSettings(GridScanAnalysisSettings().Indexing(true));
    CHECK_FALSE(experiment.GetAnalysisStages().indexing);
}

TEST_CASE("AnalysisMode_GridScanSettingsDefaults", "[AnalysisMode]") {
    const GridScanAnalysisSettings s;
    // A single cell is admitted only well above the threshold that admits a cell into a blob;
    // otherwise the decisive rule would just be a lower general threshold.
    CHECK(s.GetDecisiveSingleCellScore() > s.GetProteinScoreThreshold());
    CHECK(s.GetMinBlobCells() >= 1);
    // No cap by default: a crystal found and then dropped cannot be recovered by the caller.
    CHECK_FALSE(s.GetMaxCrystals().has_value());

    GridScanAnalysisSettings set;
    set.ProteinScoreThreshold(0.7f).MinBlobCells(5).DecisiveSingleCellScore(0.95f)
       .MaxCrystals(3).Indexing(false);
    CHECK(set.GetProteinScoreThreshold() == Catch::Approx(0.7));
    CHECK(set.GetMinBlobCells() == 5);
    CHECK(set.GetDecisiveSingleCellScore() == Catch::Approx(0.95));
    CHECK(set.GetMaxCrystals() == 3);
    CHECK_FALSE(set.IsIndexing());
}

TEST_CASE("AnalysisMode_PrecedenceOverSpotFindingSwitch", "[AnalysisMode]") {
    DiffractionExperiment experiment;
    DatasetSettings dataset;
    dataset.SpotFindingEnable(true).MaxSpotCount(500);
    experiment.ImportDatasetSettings(dataset);

    CHECK(experiment.IsSpotFindingEnabled());

    // A mode that analyses no spots wins over the deprecated per-dataset switch, and takes the spot
    // budget with it.
    experiment.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::Azint));
    CHECK_FALSE(experiment.IsSpotFindingEnabled());
    CHECK(experiment.GetMaxSpotCount() == 0);

    // Under a mode that does find spots, the deprecated switch is still able to turn it off.
    experiment.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::MXStills));
    CHECK(experiment.IsSpotFindingEnabled());
    dataset.SpotFindingEnable(false);
    experiment.ImportDatasetSettings(dataset);
    CHECK_FALSE(experiment.IsSpotFindingEnabled());
}

TEST_CASE("AnalysisMode_PowderCalibrationForcesCPUAzInt", "[AnalysisMode]") {
    // The FPGA integration core holds 2048 bins in total, so a sectored profile cannot be built there.
    DiffractionExperiment experiment;
    REQUIRE_FALSE(experiment.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow());

    experiment.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::PowderCalibration));
    CHECK(experiment.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow());
    CHECK(experiment.GetAzimuthalIntegrationSettings().GetAzimuthalBinCount()
          == CALIBRATION_AZIM_BINS_DEFAULT);

    // An explicit sector count stands; the mode only supplies one where none is usable.
    DiffractionExperiment explicit_bins;
    explicit_bins.ImportAzimuthalIntegrationSettings(
            AzimuthalIntegrationSettings().AzimuthalBinCount(64));
    explicit_bins.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::PowderCalibration));
    CHECK(explicit_bins.GetAzimuthalIntegrationSettings().GetAzimuthalBinCount() == 64);

    // The coupling is applied from both imports, so the order the two are set in does not matter.
    DiffractionExperiment settings_first;
    settings_first.ImportCalibrationSettings(CalibrationSettings().Calibrant("lab6"));
    CHECK_FALSE(settings_first.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow());
    settings_first.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::PowderCalibration));
    CHECK(settings_first.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow());

    DiffractionExperiment mode_first;
    mode_first.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::PowderCalibration));
    mode_first.ImportCalibrationSettings(CalibrationSettings().Calibrant("lab6"));
    CHECK(mode_first.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow());
    CHECK(mode_first.GetCalibrationSettings().GetCalibrant() == "lab6");

    // A mode that is not a calibration is left alone.
    DiffractionExperiment stills;
    stills.ImportCalibrationSettings(CalibrationSettings().Method(CalibrationMethod::Spots));
    CHECK_FALSE(stills.GetAzimuthalIntegrationSettings().IsForceCPUinFPGAWorkflow());
    CHECK(stills.GetCalibrationSettings().GetMethod() == CalibrationMethod::Spots);
}

TEST_CASE("AnalysisMode_CBORStartRoundTrip", "[AnalysisMode][CBOR]") {
    std::vector<uint8_t> buffer(1024 * 1024);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());

    DiffractionExperiment experiment;
    experiment.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::Grid));

    StartMessage message{};
    experiment.FillMessage(message);
    REQUIRE(message.analysis_mode.has_value());
    REQUIRE_NOTHROW(serializer.SerializeSequenceStart(message));

    auto deserialized = CBORStream2Deserialize(buffer.data(), serializer.GetBufferSize());
    REQUIRE(deserialized);
    REQUIRE(deserialized->start_message);
    REQUIRE(deserialized->start_message->analysis_mode.has_value());
    CHECK(*deserialized->start_message->analysis_mode == AnalysisMode::Grid);
}

TEST_CASE("AnalysisMode_CBORStartAbsentMode", "[AnalysisMode][CBOR]") {
    // A stream written before the mode existed says nothing, and is read back as saying nothing -
    // naming a mode for it would be an invention rather than provenance.
    std::vector<uint8_t> buffer(1024 * 1024);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());

    StartMessage message{};
    REQUIRE_NOTHROW(serializer.SerializeSequenceStart(message));

    auto deserialized = CBORStream2Deserialize(buffer.data(), serializer.GetBufferSize());
    REQUIRE(deserialized);
    REQUIRE(deserialized->start_message);
    CHECK_FALSE(deserialized->start_message->analysis_mode.has_value());
}

TEST_CASE("AnalysisMode_CBORGridCrystals", "[AnalysisMode][CBOR]") {
    std::vector<uint8_t> buffer(1024 * 1024);
    CBORStream2Serializer serializer(buffer.data(), buffer.size());

    EndMessage message{};
    message.max_image_number = 0;
    // Two crystals, not one: a raster can find any number and nothing may assume at most one.
    message.grid_crystals.push_back(GridScanCrystal{
        .nx = 3.5f, .ny = 7.25f, .x_um = -12.5f, .y_um = 40.0f, .image_number = 143,
        .major_um = 22.5f, .minor_um = 8.0f, .angle_deg = 179.5f, .score = 0.92f,
        .ice_score = 0.03f, .res_A = 1.85f, .n_images = 17});
    message.grid_crystals.push_back(GridScanCrystal{
        .nx = 11.0f, .ny = 2.0f, .x_um = 60.0f, .y_um = -30.0f, .image_number = 44,
        .major_um = 9.0f, .minor_um = 7.5f, .angle_deg = 0.5f, .score = 0.41f,
        .ice_score = 0.30f, .res_A = 3.2f, .n_images = 4});
    REQUIRE_NOTHROW(serializer.SerializeSequenceEnd(message));

    auto deserialized = CBORStream2Deserialize(buffer.data(), serializer.GetBufferSize());
    REQUIRE(deserialized);
    REQUIRE(deserialized->end_message);
    const auto &out = deserialized->end_message->grid_crystals;
    REQUIRE(out.size() == 2);
    CHECK(out[0].nx == Catch::Approx(3.5));
    CHECK(out[0].image_number == 143);
    CHECK(out[0].major_um == Catch::Approx(22.5));
    CHECK(out[0].angle_deg == Catch::Approx(179.5));
    CHECK(out[0].res_A == Catch::Approx(1.85));
    CHECK(out[0].n_images == 17);
    CHECK(out[1].score == Catch::Approx(0.41));
    CHECK(out[1].ice_score == Catch::Approx(0.30));
    CHECK(out[1].y_um == Catch::Approx(-30.0));
}

TEST_CASE("AnalysisMode_HDF5MasterRoundTrip", "[AnalysisMode][HDF5][Full]") {
    // The mode and the crystal list are dataset-wide metadata: written to the master, read back from
    // it, so a stored raster re-opens knowing what produced it and what it found.
    DiffractionExperiment x(DetJF(1));
    x.FilePrefix("test_analysis_mode").ImagesPerTrigger(1).OverwriteExistingFiles(true);
    x.BeamX_pxl(100).BeamY_pxl(200).DetectorDistance_mm(150)
            .IncidentEnergy_keV(WVL_1A_IN_KEV).PixelSigned(false).BitDepthImage(16)
            .FrameTime(std::chrono::microseconds(500), std::chrono::microseconds(10));
    x.ImportAnalysisSettings(AnalysisSettings().Mode(AnalysisMode::Grid));
    RegisterHDF5Filter();

    std::vector<uint16_t> image(x.GetPixelsNum(), 0);
    {
        StartMessage start_message;
        x.FillMessage(start_message);

        FileWriter file_set(start_message);

        DataMessage message{};
        message.image = CompressedImage(image, x.GetXPixelsNum(), x.GetYPixelsNum());
        message.number = 0;
        REQUIRE_NOTHROW(file_set.WriteHDF5(message));

        EndMessage end_message;
        end_message.max_image_number = 1;
        end_message.grid_crystals.push_back(GridScanCrystal{
            .nx = 3.5f, .ny = 7.25f, .x_um = -12.5f, .y_um = 40.0f, .image_number = 143,
            .major_um = 22.5f, .minor_um = 8.0f, .angle_deg = 179.5f, .score = 0.92f,
            .ice_score = 0.03f, .res_A = 1.85f, .n_images = 17});
        end_message.grid_crystals.push_back(GridScanCrystal{
            .nx = 11.0f, .ny = 2.0f, .x_um = 60.0f, .y_um = -30.0f, .image_number = 44,
            .major_um = 9.0f, .minor_um = 7.5f, .angle_deg = 0.5f, .score = 0.41f,
            .ice_score = 0.30f, .res_A = 3.2f, .n_images = 4});
        REQUIRE_NOTHROW(file_set.WriteHDF5(end_message));
        file_set.Finalize();
    }
    {
        JFJochHDF5Reader reader;
        reader.ReadFile("test_analysis_mode_master.h5");
        auto dataset = reader.GetDataset();
        REQUIRE(dataset->file_analysis_mode.has_value());
        CHECK(*dataset->file_analysis_mode == AnalysisMode::Grid);
        REQUIRE(dataset->grid_crystals.size() == 2);
        CHECK(dataset->grid_crystals[0].nx == Catch::Approx(3.5));
        CHECK(dataset->grid_crystals[0].image_number == 143);
        CHECK(dataset->grid_crystals[0].angle_deg == Catch::Approx(179.5));
        CHECK(dataset->grid_crystals[0].n_images == 17);
        CHECK(dataset->grid_crystals[1].score == Catch::Approx(0.41));
        CHECK(dataset->grid_crystals[1].res_A == Catch::Approx(3.2));
    }
    remove("test_analysis_mode_master.h5");
    remove("test_analysis_mode_data_000001.h5");
}