Jungfraujoch/tests/JFJochReceiverLiteTest.cpp

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

#include <catch2/catch_all.hpp>

#include <fstream>

#include "../writer/HDF5Objects.h"
#include "../receiver/JFJochReceiverService.h"
#include "../writer/StreamWriter.h"
#include "../image_pusher/HDF5FilePusher.h"
#include "../image_puller/TestImagePuller.h"

// PixelSigned(true) below is not incidental: the frames fed in come from compression_benchmark.h5,
// which stores signed int16, while a DECTRIS experiment declares unsigned by default. The receiver
// forwards images byte-for-byte, so the start message it re-emits has to describe them - otherwise
// the master file declares a type the data files do not have, which HDF5DataFile now refuses.
TEST_CASE("JFJochReceiverLite", "[JFJochReceiver]") {
    Logger logger("JFJochReceiverLite");

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetDECTRIS(2068, 2164, "Test", {}));

    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(2)
            .FilePrefix("crystal_test_lite").JungfrauConvPhotonCnt(false).PixelSigned(true).SetFileWriterFormat(
                    FileWriterFormat::NXmxVDS).OverwriteExistingFiles(true)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90});
    PixelMask pixel_mask(experiment);

    // Load example image
    HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5");
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    REQUIRE(file_space.GetDimensions()[2] == experiment.GetXPixelsNum());
    REQUIRE(file_space.GetDimensions()[1] == experiment.GetYPixelsNum());
    std::vector<int16_t> image_conv (file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

    std::vector<hsize_t> start = {4,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    HDF5FilePusher pusher;

    auto puller = std::make_shared<TestImagePuller>();

    StartMessage start_msg;
    experiment.FillMessage(start_msg);
    // As a DECTRIS detector sends it by default; the flags must reach the master file.
    start_msg.countrate_correction_enabled = true;
    start_msg.flatfield_enabled = true;
    start_msg.countrate_correction_lookup_table = {0, 1, 3, 7};
    start_msg.virtual_pixel_interpolation_enabled = true;
    // What the detector did to the pixels, not the local setting
    start_msg.pixel_mask_enabled = !experiment.IsApplyPixelMask();

    puller->Put(ImagePullerOutput{
        .cbor = std::make_shared<CBORStream2DeserializerOutput>(start_msg)
    });

    DataMessage data_msg;
    data_msg.image = CompressedImage(image_conv, file_space.GetDimensions()[2], file_space.GetDimensions()[1]);

    for (int i = 0; i < 5; i++) {
        data_msg.number = i;
        puller->Put(ImagePullerOutput{
                .cbor = std::make_shared<CBORStream2DeserializerOutput>(data_msg)
        });
    }

    EndMessage end_msg{};

    puller->Put(ImagePullerOutput{
            .cbor = std::make_shared<CBORStream2DeserializerOutput>(end_msg)
    });

    AcquisitionDeviceGroup group;
    JFJochReceiverService service(group, logger, pusher);

    service.NumThreads(nthreads);

    // No progress value at the start of measurement
    REQUIRE(!service.GetProgress().has_value());

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 1;
    settings.max_pix_per_spot = 200;
    settings.high_resolution_limit = 2.0;
    settings.low_resolution_limit = 50.0;
    service.SetSpotFindingSettings(settings);
    service.Indexing(experiment.GetIndexingSettings());

    service.Start(experiment, pixel_mask, nullptr, puller);
    auto receiver_out = service.Stop();

    CHECK(receiver_out.efficiency == 1.0);
    REQUIRE(receiver_out.status.indexing_rate);
    CHECK(receiver_out.status.indexing_rate.value() == 1.0);
    CHECK(receiver_out.status.images_sent == experiment.GetImageNum());
    CHECK(!receiver_out.status.cancelled);

    // No progress value at the end of the measurement
    REQUIRE(!service.GetProgress().has_value());

    HDF5ReadOnlyFile master("crystal_test_lite_master.h5");
    CHECK(master.GetBool("/entry/instrument/detector/countrate_correction_applied"));
    CHECK(master.GetBool("/entry/instrument/detector/flatfield_applied"));
    CHECK(master.GetBool("/entry/instrument/detector/virtual_pixel_interpolation_applied"));
    CHECK(master.GetBool("/entry/instrument/detector/pixel_mask_applied") == start_msg.pixel_mask_enabled);
    CHECK(master.ReadVector<uint32_t>("/entry/instrument/detector/countrate_correction_lookup_table")
          == std::vector<uint32_t>{0, 1, 3, 7});
}

TEST_CASE("JFJochReceiverLite_Cancel", "[JFJochReceiver]") {
    Logger logger("JFJochReceiverLite_Cancel");

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetDECTRIS(2068, 2164, "Test", {}));

    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(2)
            .FilePrefix("crystal_test_lite").JungfrauConvPhotonCnt(false).PixelSigned(true).SetFileWriterFormat(
                    FileWriterFormat::NXmxVDS).OverwriteExistingFiles(true)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90});
    PixelMask pixel_mask(experiment);

    HDF5FilePusher pusher;

    auto puller = std::make_shared<TestImagePuller>();

    AcquisitionDeviceGroup group;
    JFJochReceiverService service(group, logger, pusher);

    service.Indexing(experiment.GetIndexingSettings());
    service.NumThreads(nthreads);

    // No progress value at the start of measurement
    REQUIRE(!service.GetProgress().has_value());

    service.Start(experiment, pixel_mask, nullptr, puller);
    std::this_thread::sleep_for(std::chrono::seconds(10));
    service.Cancel(false);

    auto receiver_out = service.Stop();

    CHECK(receiver_out.efficiency == 0.0);
    CHECK(receiver_out.status.cancelled);
}

TEST_CASE("JFJochReceiverLite_DarkMask", "[JFJochReceiver]") {
    Logger logger("JFJochReceiverLite_DarkMask");

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetDECTRIS(2068, 2164, "Test", {}));

    DarkMaskSettings mask_settings;
    mask_settings.NumberOfFrames(10).MaxCounts(1).MaxFramesWithCounts(5);
    experiment.Mode(DetectorMode::DarkMask).ImportDarkMaskSettings(mask_settings);

    PixelMask pixel_mask(experiment);

    HDF5FilePusher pusher;

    auto puller = std::make_shared<TestImagePuller>();

    StartMessage start_msg;
    experiment.FillMessage(start_msg);

    puller->Put(ImagePullerOutput{
        .cbor = std::make_shared<CBORStream2DeserializerOutput>(start_msg)
    });


    std::vector<uint16_t> image(experiment.GetPixelsNum(), 0);
    image[1] = 1;
    image[2] = 2;
    image[3] = 3;

    std::vector<uint16_t> image_last(experiment.GetPixelsNum(), 0);
    image_last[1] = 0;
    image_last[2] = 0;
    image_last[3] = 0;
    image_last[68] = 10000;

    for (int i = 0; i < experiment.GetFrameNum(); i++) {
        DataMessage data_msg;
        if (i == experiment.GetFrameNum() - 1)
            data_msg.image = CompressedImage(image_last, experiment.GetXPixelsNum(), experiment.GetYPixelsNum());
        else
            data_msg.image = CompressedImage(image, experiment.GetXPixelsNum(), experiment.GetYPixelsNum());
        data_msg.number = i;
        puller->Put(ImagePullerOutput{
                .cbor = std::make_shared<CBORStream2DeserializerOutput>(data_msg)
        });
    }

    EndMessage end_msg{};

    puller->Put(ImagePullerOutput{
            .cbor = std::make_shared<CBORStream2DeserializerOutput>(end_msg)
    });

    AcquisitionDeviceGroup group;
    JFJochReceiverService service(group, logger, pusher);

    service.NumThreads(nthreads);

    // No progress value at the start of measurement
    REQUIRE(!service.GetProgress().has_value());

    service.Start(experiment, pixel_mask, nullptr, puller);
    auto receiver_out = service.Stop();

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.images_sent == 0);
    CHECK(receiver_out.dark_mask_result.size() == experiment.GetPixelsNum());
    CHECK(receiver_out.dark_mask_result[0] == 0);
    CHECK(receiver_out.dark_mask_result[1] == 0);
    CHECK(receiver_out.dark_mask_result[2] != 0);
    CHECK(receiver_out.dark_mask_result[3] != 0);
    CHECK(receiver_out.dark_mask_result[4] == 0);
    CHECK(receiver_out.dark_mask_result[68] == 0);

    CHECK(!receiver_out.status.cancelled);

    // No progress value at the end of the measurement
    REQUIRE(!service.GetProgress().has_value());
}