Jungfraujoch/tests/JFJochReceiverProcessingTest.cpp

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

#include <catch2/catch_all.hpp>

#include <fstream>

#include "../receiver/JFJochReceiverTest.h"
#include "../acquisition_device/HLSSimulatedDevice.h"
#include "../writer/HDF5Objects.h"
#include "../receiver/JFJochReceiverService.h"
#include "../preview/JFJochTIFF.h"
#include "../image_pusher/ZMQStream2Pusher.h"
#include "../writer/StreamWriter.h"
#include "../image_pusher/NonePusher.h"
#include "../image_pusher/HDF5FilePusher.h"

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(2)
            .FilePrefix("lyso_test").JungfrauConvPhotonCnt(false).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);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);

    aq_devices.Add(std::move(test));

    ZMQStream2Pusher pusher({"ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");

    ZMQImagePuller puller(pusher.GetAddress()[0]);
    StreamWriter writer(logger, puller);
    auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer);

    JFJochReceiverService service(aq_devices, 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.Start(experiment, pixel_mask, nullptr);
    auto receiver_out = service.Stop();

    std::string jpeg;
    PreviewJPEGSettings jpeg_settings{
            .saturation_value = 10,
            .jpeg_quality = 100,
            .show_spots = true
    };
    REQUIRE_NOTHROW(jpeg = service.GetJPEGFromBuffer(jpeg_settings));
    REQUIRE(!jpeg.empty());
    std::ofstream f("lyso_processing_test_min_pix_1.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    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 measurement
    REQUIRE(!service.GetProgress().has_value());

    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_FFT", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(2)
            .FilePrefix("lyso_test_fft").JungfrauConvPhotonCnt(false).SetFileWriterFormat(FileWriterFormat::NXmxVDS)
            .OverwriteExistingFiles(true)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4)
            .IndexingAlgorithm(IndexingAlgorithmEnum::FFT)
            .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);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);

    aq_devices.Add(std::move(test));

    ZMQStream2Pusher pusher({"ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");

    ZMQImagePuller puller(pusher.GetAddress()[0]);
    StreamWriter writer(logger, puller);
    auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer);

    JFJochReceiverService service(aq_devices, 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.Start(experiment, pixel_mask, nullptr);
    auto receiver_out = service.Stop();

    std::string jpeg;
    PreviewJPEGSettings jpeg_settings{
            .saturation_value = 10,
            .jpeg_quality = 100,
            .show_spots = true
    };
    REQUIRE_NOTHROW(jpeg = service.GetJPEGFromBuffer(jpeg_settings));
    REQUIRE(!jpeg.empty());
    std::ofstream f("lyso_processing_test_min_pix_1.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    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 measurement
    REQUIRE(!service.GetProgress().has_value());

    REQUIRE_NOTHROW(writer_future.get());
}


TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_grid", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    GridScanSettings grid_scan(3, 10.0f, 25.0f, false, true);
    DatasetSettings dataset_settings;
    dataset_settings.ImagesPerTrigger(7).NumTriggers(1).DetectorDistance_mm(75)
            .BeamY_pxl(1136).BeamX_pxl(1090)
            .PhotonEnergy_keV(12.4).GridScan(grid_scan)
            .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90});

    DiffractionExperiment experiment(DetJF4M());
    experiment.UseInternalPacketGenerator(true).ImagesPerFile(2).ImportDatasetSettings(dataset_settings)
            .FilePrefix("lyso_test").JungfrauConvPhotonCnt(false)
            .SetFileWriterFormat(FileWriterFormat::NXmxVDS).OverwriteExistingFiles(true);
    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);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);

    aq_devices.Add(std::move(test));

    ZMQStream2Pusher pusher({"ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");

    ZMQImagePuller puller(pusher.GetAddress()[0]);
    StreamWriter writer(logger, puller);
    auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer);

    JFJochReceiverService service(aq_devices, 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.Start(experiment, pixel_mask, nullptr);
    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);

    auto grid_scan_out = service.GetDataProcessingPlot(
            PlotRequest{.type = PlotType::IndexingRate, .experimental_coord = true});

    CHECK(grid_scan_out.GetUnits() == MultiLinePlotUnits::Grid_um);
    REQUIRE(grid_scan_out.GetPlots().size() == 1);
    REQUIRE(grid_scan_out.GetPlots()[0].z.size() == 7);

    CHECK(grid_scan_out.GetPlots()[0].z[0] == 1.0f);
    CHECK(grid_scan_out.GetPlots()[0].z[1] == 1.0f);
    CHECK(grid_scan_out.GetPlots()[0].z[2] == 1.0f);
    CHECK(grid_scan_out.GetPlots()[0].z[3] == 1.0f);
    CHECK(grid_scan_out.GetPlots()[0].z[4] == 1.0f);
    CHECK(grid_scan_out.GetPlots()[0].z[5] == 1.0f);
    CHECK(grid_scan_out.GetPlots()[0].z[6] == 1.0f);

    CHECK(grid_scan_out.GetPlots()[0].x[5] == 1.5 * 10.0f);
    CHECK(grid_scan_out.GetPlots()[0].y[5] == 2.5 * 25.0f);
    CHECK(grid_scan_out.GetPlots()[0].x[6] == 2.5 * 10.0f);
    CHECK(grid_scan_out.GetPlots()[0].y[6] == 0.5 * 25.0f);

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

    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_32", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(2)
            .FilePrefix("lyso_test_32").JungfrauConvPhotonCnt(false).SetFileWriterFormat(FileWriterFormat::NXmxLegacy).OverwriteExistingFiles(true)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4).BitDepthImage(32).PixelSigned(true)
            .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);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);

    aq_devices.Add(std::move(test));

    ZMQStream2Pusher pusher({"ipc://*"});
    pusher.WriterNotificationSocket("ipc://*");

    ZMQImagePuller puller(pusher.GetAddress()[0]);
    StreamWriter writer(logger, puller);
    auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer);

    JFJochReceiverService service(aq_devices, 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.Start(experiment, pixel_mask, nullptr);
    auto receiver_out = service.Stop();

    std::string jpeg;
    PreviewJPEGSettings jpeg_settings{
            .saturation_value = 10,
            .jpeg_quality = 100,
            .show_spots = true
    };
    REQUIRE_NOTHROW(jpeg = service.GetJPEGFromBuffer(jpeg_settings));
    REQUIRE(!jpeg.empty());
    std::ofstream f("lyso_processing_test_min_pix_1.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    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 measurement
    REQUIRE(!service.GetProgress().has_value());

    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_min_pix_2", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).OverwriteExistingFiles(true)
            .FilePrefix("lyso_test_min_pix_2").JungfrauConvPhotonCnt(false).SetFileWriterFormat(
                    FileWriterFormat::NXmxVDS)
            .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})
            .PixelSigned(true);
    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 = {2,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);
    aq_devices.Add(std::move(test));

    ZMQStream2Pusher pusher({"ipc://*"});
    ZMQImagePuller puller(pusher.GetAddress()[0]);
    StreamWriter writer(logger, puller);
    auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer);

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

    SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings();
    settings.signal_to_noise_threshold = 2.5;
    settings.photon_count_threshold = 3;
    settings.min_pix_per_spot = 2;
    settings.max_pix_per_spot = 200;
    service.SetSpotFindingSettings(settings);

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

    std::string jpeg, jpeg_indexed;
    PreviewJPEGSettings jpeg_settings{
        .saturation_value = 10,
        .jpeg_quality = 100,
        .show_spots = true
    };
    REQUIRE_NOTHROW(jpeg = service.GetJPEGFromBuffer(jpeg_settings));
    REQUIRE(!jpeg.empty());
    std::ofstream f("lyso_processing_test_min_pix_2.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    REQUIRE_NOTHROW(jpeg_indexed = service.GetJPEGFromBuffer(jpeg_settings, ImageBuffer::MaxIndexedImage));
    REQUIRE(!jpeg_indexed.empty());

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

    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("GenerateResolutionMap") {
    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true)
            .FilePrefix("lyso_test").JungfrauConvPhotonCnt(false)
            .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});

    std::vector<float> spot_finder_resolution_map(experiment.GetModulesNum() * RAW_MODULE_SIZE, 1.0);

    for (int m = 0; m < experiment.GetModulesNum(); m++)
        experiment.CalcSpotFinderResolutionMap(spot_finder_resolution_map.data() + m * RAW_MODULE_SIZE, m);

    std::vector<uint32_t> spot_finder_resolution_map_int(spot_finder_resolution_map.size());
    for (int i = 0; i < spot_finder_resolution_map.size(); i++)
        spot_finder_resolution_map_int[i] = static_cast<uint32_t>(spot_finder_resolution_map[i] * 100);
    std::vector<uint32_t> spot_finder_resolution_map_int_conv(experiment.GetPixelsNum(), 0);

    RawToConvertedGeometry(experiment, spot_finder_resolution_map_int_conv.data(), spot_finder_resolution_map_int.data());

    CompressedImage image(spot_finder_resolution_map_int_conv, experiment.GetXPixelsNum(),
                          experiment.GetYPixelsNum() );
    WriteTIFFToFile("ResolutionMap.tiff", image);
}

TEST_CASE("JFJochIntegrationTest_ZMQ_ROI", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true)
            .FilePrefix("lyso_test_roi").JungfrauConvPhotonCnt(false)
            .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})
            .SetFileWriterFormat(FileWriterFormat::NXmxVDS).PixelSigned(true).OverwriteExistingFiles(true);

    PixelMask pixel_mask(experiment);

    experiment.ROI().SetROI(ROIDefinition{
        .boxes = {ROIBox("beam", 100, 120, 20, 30)},
        .circles = {ROICircle("roi1", 500, 800, 10)},
        .azimuthal = {ROIAzimuthal("roi2", 4.0, 7.0)}
    });

    // 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 = {5,0,0};
    std::vector<hsize_t> file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]};
    dataset.ReadVector(image_conv, start, file_size);

    uint64_t roi_value = 0;
    uint64_t pixels = 0;
    for (int y = 20; y <= 30; y++) {
        for (int x = 100; x <= 120; x++) {
            int16_t val = image_conv[experiment.GetXPixelsNum() * y + x];
            if ((val != INT16_MIN) && (val != INT16_MAX)) {
                pixels += 1;
                roi_value += val;
            }
        }
    }

    std::vector<int16_t> image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE);
    ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data());
    logger.Info("Loaded image");

    auto mask = pixel_mask.GetMaskRaw(experiment);

    uint64_t sat_pixels = 0;
    uint64_t err_pixels = 0;
    for (int i = 0; i < image_raw_geom.size(); i++) {
        if (mask[i] == 0) {
            if (image_raw_geom[i] == INT16_MAX)
                sat_pixels++;
            else if (image_raw_geom[i] == INT16_MIN)
                err_pixels++;
        }
    }

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    std::unique_ptr<HLSSimulatedDevice> test = std::make_unique<HLSSimulatedDevice>(0, 64);
    for (int m = 0; m < experiment.GetModulesNum(); m++)
        test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m);

    aq_devices.Add(std::move(test));

    ZMQStream2Pusher pusher({"ipc://*"});
    ZMQImagePuller puller(pusher.GetAddress()[0]);
    StreamWriter writer(logger, puller);
    auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer);

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

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

    std::string jpeg;
    PreviewJPEGSettings jpeg_settings{
            .saturation_value = 10,
            .jpeg_quality = 100,
            .show_spots = true,
            .show_roi = true
    };

    REQUIRE_NOTHROW(jpeg = service.GetJPEGFromBuffer(jpeg_settings));

   REQUIRE(!jpeg.empty());
    std::ofstream f("lyso_processing_test_roi.jpeg", std::ios::binary);
    f.write(jpeg.data(), jpeg.size());

    auto plot = service.GetDataProcessingPlot(PlotRequest{.type = PlotType::ROISum, .binning = 1});

    REQUIRE(plot.GetPlots().size() == 3);
    CHECK(plot.GetPlots()[0].title == "beam");
    REQUIRE(!plot.GetPlots()[0].x.empty());
    CHECK(plot.GetPlots()[0].x[0] == 0);
    CHECK(plot.GetPlots()[0].y[0] == roi_value);
    CHECK(plot.GetPlots()[1].title == "roi1");
    CHECK(plot.GetPlots()[2].title == "roi2");

    CHECK(receiver_out.status.roi_beam_sum == roi_value);
    CHECK(receiver_out.status.roi_beam_npixel == pixels);
    REQUIRE(receiver_out.status.error_pixels.has_value());
    REQUIRE(receiver_out.status.saturated_pixels.has_value());

    CHECK(receiver_out.status.error_pixels.value() == err_pixels);
    CHECK(receiver_out.status.saturated_pixels.value() == sat_pixels);
    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.images_sent == experiment.GetImageNum());
    CHECK(!receiver_out.status.cancelled);

    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("JFJochIntegrationTest_ZMQPreview", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true)
            .FilePrefix("").JungfrauConvPhotonCnt(false)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4);

    PixelMask pixel_mask(experiment);

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    aq_devices.Add(std::make_unique<HLSSimulatedDevice>(0, 64));

    NonePusher pusher;

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);
    service.PreviewSocket("tcp://0.0.0.0:*");
    service.PreviewSocketSettings({.period = std::chrono::seconds(0)}); // Send one message

    ZMQSocket sub_socket(ZMQSocketType::Sub);
    sub_socket.ReceiveWaterMark(10);
    sub_socket.Connect(service.GetPreviewSocketAddress());
    sub_socket.SubscribeAll();
    sub_socket.ReceiveTimeout(std::chrono::seconds(1)); // Ensure Subscriber is connected

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

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.images_collected == 5);
    CHECK(receiver_out.status.images_sent == 0);
    CHECK(!receiver_out.status.cancelled);

    ZMQMessage msg1;
    REQUIRE(sub_socket.Receive(msg1, true));

    // start message
    auto out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(out != nullptr);
    CHECK(out->msg_type == CBORImageType::START);

    REQUIRE(sub_socket.Receive(msg1, true));
    out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(out != nullptr);
    CHECK(out->msg_type == CBORImageType::IMAGE);

    // 4 other images
    REQUIRE(sub_socket.Receive(msg1, true));
    REQUIRE(sub_socket.Receive(msg1, true));
    REQUIRE(sub_socket.Receive(msg1, true));
    REQUIRE(sub_socket.Receive(msg1, true));

    // end message
    REQUIRE(sub_socket.Receive(msg1, true));
    REQUIRE(!sub_socket.Receive(msg1, true));
}

TEST_CASE("JFJochIntegrationTest_ZMQMetadata", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true)
            .FilePrefix("").JungfrauConvPhotonCnt(false)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4);

    PixelMask pixel_mask(experiment);

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    aq_devices.Add(std::make_unique<HLSSimulatedDevice>(0, 64));

    NonePusher pusher;

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);
    service.MetadataSocket("tcp://0.0.0.0:*");
    service.MetadataSocketSettings({.period = std::chrono::seconds(2500)}); // Send all metadata as one message

    ZMQSocket sub_socket(ZMQSocketType::Sub);
    sub_socket.ReceiveWaterMark(10);
    sub_socket.Connect(service.GetMetadataSocketAddress());
    sub_socket.SubscribeAll();
    sub_socket.ReceiveTimeout(std::chrono::seconds(1)); // Ensure Subscriber is connected

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

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.images_collected == 5);
    CHECK(receiver_out.status.images_sent == 0);
    CHECK(!receiver_out.status.cancelled);

    ZMQMessage msg1;
    REQUIRE(sub_socket.Receive(msg1, true));

    auto out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(out != nullptr);
    CHECK(out->msg_type == CBORImageType::START);

    REQUIRE(sub_socket.Receive(msg1, true));
    out = CBORStream2Deserialize(msg1.data(), msg1.size());
    REQUIRE(out != nullptr);
    CHECK(out->msg_type == CBORImageType::METADATA);
    CHECK(out->metadata);
    CHECK(out->metadata->images.size() == 5);

    // end message
    REQUIRE(sub_socket.Receive(msg1, true));
    REQUIRE(!sub_socket.Receive(msg1, true));
}

TEST_CASE("JFJochIntegrationTest_HDF5FilePusher", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF4M());
    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true)
            .FilePrefix("HDF5FilePusherTest").JungfrauConvPhotonCnt(true)
            .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4);

    JFCalibration calibration(experiment);
    PixelMask pixel_mask(experiment);

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    aq_devices.Add(std::make_unique<HLSSimulatedDevice>(0, 64));

    HDF5FilePusher pusher;

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

    service.Start(experiment, pixel_mask, &calibration);
    auto receiver_out = service.Stop();

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.images_collected == 5);
    CHECK(receiver_out.status.images_sent == 5);
    CHECK(!receiver_out.status.cancelled);
}

TEST_CASE("JFJochIntegrationTest_HDF5FilePusher_Raw", "[JFJochReceiver]") {
    Logger logger(Catch::getResultCapture().getCurrentTestName());

    RegisterHDF5Filter();

    const uint16_t nthreads = 4;

    DiffractionExperiment experiment(DetJF(1));

    experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true)
            .FilePrefix("HDF5FilePusherTest_Raw");
    experiment.Raw();

    PixelMask pixel_mask(experiment);

    // Setup acquisition device
    AcquisitionDeviceGroup aq_devices;
    aq_devices.Add(std::make_unique<HLSSimulatedDevice>(0, 64));

    HDF5FilePusher pusher;

    JFJochReceiverService service(aq_devices, logger, pusher);
    service.NumThreads(nthreads);

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

    CHECK(receiver_out.efficiency == 1.0);
    CHECK(receiver_out.status.images_collected == 5);
    CHECK(receiver_out.status.images_sent == 5);
    CHECK(!receiver_out.status.cancelled);
}