// Copyright (2019-2023) Paul Scherrer Institute

#include <catch2/catch.hpp>

#include "../grpc/gRPCServer_Template.h"
#include "../broker/JFJochStateMachine.h"
#include "../writer/StreamWriter.h"
#include "FPGAUnitTest.h"
#include "../acquisition_device/MockAcquisitionDevice.h"
#include "../common/ZMQImagePusher.h"

using namespace std::literals::chrono_literals;

TEST_CASE("JFJochIntegrationTest_ZMQ", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ");
    ZMQContext zmq_context;

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        aq_devices.Add(std::move(test));
    }

    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    
    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);

    REQUIRE(!state_machine.GetMeasurementStatistics().has_value());
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);
    services.Receiver(&fpga_receiver);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.images_per_trigger = 5;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_test";
    setup.photon_energy_keV = 12.4;
    setup.data_file_count = 2;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    BrokerStatus status;
    status = state_machine.GetStatus();
    REQUIRE(status.progress == Approx(0.0));
    REQUIRE(status.broker_state == JFJochState::Measuring);

    for (int i = 0; i < ndatastream; i++) {
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int image_num = 1; image_num <= nimages; image_num++)
                dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).AddModule(image_num, m, image.data());
        }
        dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).Terminate();
    }

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    status = state_machine.GetStatus();
    REQUIRE(status.broker_state == JFJochState::Idle);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 5);
    REQUIRE(statistics.max_image_number_sent == 4);
    REQUIRE(!statistics.cancelled);
    REQUIRE(statistics.file_prefix == "integration_test");
    REQUIRE(statistics.detector_width == 2068);
    REQUIRE(statistics.detector_height == 2164);
    REQUIRE(statistics.detector_pixel_depth == 2);
    writer_future.get();
}

TEST_CASE("JFJochIntegrationTest_ZMQ_save_calibration", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_save_calibration");
    ZMQContext zmq_context;

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;
    
    AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        aq_devices.Add(std::move(test));
    }

    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    
    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    
    REQUIRE(!state_machine.GetMeasurementStatistics().has_value());
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);
    services.Receiver(&fpga_receiver);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);;

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.images_per_trigger = 5;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_test_with_calibration";
    setup.photon_energy_keV = 12.4;
    setup.data_file_count = 2;
    setup.save_calibration = true;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    BrokerStatus status;
    status = state_machine.GetStatus();
    REQUIRE(status.progress == Approx(0.0));
    REQUIRE(status.broker_state == JFJochState::Measuring);

    for (int i = 0; i < ndatastream; i++) {
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int image_num = 1; image_num <= nimages; image_num++)
                dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).AddModule(image_num, m, image.data());
        }
        dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).Terminate();
    }

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    status = state_machine.GetStatus();
    REQUIRE(status.broker_state == JFJochState::Idle);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 5);
    REQUIRE(statistics.max_image_number_sent == 4);
    REQUIRE(!statistics.cancelled);
    REQUIRE(statistics.file_prefix == "integration_test_with_calibration");
    REQUIRE(statistics.detector_width == 2068);
    REQUIRE(statistics.detector_height == 2164);
    REQUIRE(statistics.detector_pixel_depth == 2);
    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("JFJochIntegrationTest_ZMQ_2DataStreams_4Devices", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_2DataStreams_4Devices");
    ZMQContext zmq_context;

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < 4; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        aq_devices.Add(std::move(test));
    }

    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    
    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);

    REQUIRE(!state_machine.GetMeasurementStatistics().has_value());

    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);
    services.Receiver(&fpga_receiver);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.images_per_trigger = 5;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_test";
    setup.photon_energy_keV = 12.4;
    setup.data_file_count = 2;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    BrokerStatus status;
    status = state_machine.GetStatus();
    REQUIRE(status.progress == Approx(0.0));
    REQUIRE(status.broker_state == JFJochState::Measuring);

    for (int i = 0; i < ndatastream; i++) {
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int image_num = 1; image_num <= nimages; image_num++)
                dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).AddModule(image_num, m, image.data());
        }
        dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).Terminate();
    }

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    status = state_machine.GetStatus();
    REQUIRE(status.broker_state == JFJochState::Idle);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 5);
    REQUIRE(statistics.max_image_number_sent == 4);
    REQUIRE(!statistics.cancelled);
    REQUIRE(statistics.file_prefix == "integration_test");
    REQUIRE(statistics.detector_width == 2068);
    REQUIRE(statistics.detector_height == 2164);
    REQUIRE(statistics.detector_pixel_depth == 2);

    REQUIRE_NOTHROW(writer_future.get());
}

TEST_CASE("JFJochIntegrationTest_ZMQ_RAW", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ");
    ZMQContext zmq_context;

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);

    REQUIRE(!state_machine.GetMeasurementStatistics().has_value());
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().Mode(DetectorMode::Conversion);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

   AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        aq_devices.Add(std::move(test));
    }
    
    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    DetectorSettings detector_settings{};
    detector_settings.frame_time_us = 500;
    detector_settings.collect_raw_data = true;
    detector_settings.storage_cell_count = 1;
    REQUIRE_NOTHROW(state_machine.SetDetectorSettings(detector_settings));
    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_raw_test";
    setup.images_per_trigger = 5;
    setup.photon_energy_keV = 12.4;
    setup.data_file_count = 5;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    BrokerStatus status;
    status = state_machine.GetStatus();
    REQUIRE(status.progress == Approx(0.0));
    REQUIRE(status.broker_state == JFJochState::Measuring);

    for (int i = 0; i < ndatastream; i++) {
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int image_num = 1; image_num <= nimages; image_num++)
                dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).AddModule(image_num, m, image.data());
        }
        dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).Terminate();
    }

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    status = state_machine.GetStatus();
    REQUIRE(status.broker_state == JFJochState::Idle);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 5);
    REQUIRE(statistics.file_prefix == "integration_raw_test");
    REQUIRE(statistics.detector_width == 1024);
    REQUIRE(statistics.detector_height == 8 * 512);
    REQUIRE(statistics.detector_pixel_depth == 2);

    REQUIRE_NOTHROW(writer_future.get());
}


TEST_CASE("JFJochIntegrationTest_ZMQ_3Writers", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_3Writers");
    ZMQContext zmq_context;

    RegisterHDF5Filter();

    int64_t nimages = 37;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

   AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        aq_devices.Add(std::move(test));
    }
    
    ZMQImagePusher pusher(zmq_context, {"inproc://#0", "inproc://#1", "inproc://#2"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    StreamWriter writer_0(zmq_context, logger, "inproc://#0");
    auto writer_0_future = writer_0.RunFuture();

    StreamWriter writer_1(zmq_context, logger, "inproc://#1");
    auto writer_1_future = writer_1.RunFuture();

    StreamWriter writer_2(zmq_context, logger, "inproc://#2");
    auto writer_2_future = writer_2.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_test_3writers";
    setup.images_per_trigger = nimages;
    setup.photon_energy_keV = 12.4;
    setup.data_file_count = 5;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    auto status = state_machine.GetStatus();
    REQUIRE(status.progress == Approx(0.0));
    REQUIRE(status.broker_state == JFJochState::Measuring);

    for (int i = 0; i < ndatastream; i++) {
        for (int image_num = 1; image_num <= nimages; image_num++) {
            for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++)
                dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).AddModule(image_num, m, image.data());
        }
        dynamic_cast<MockAcquisitionDevice&>(aq_devices[i]).Terminate();
    }

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    status = state_machine.GetStatus();
    REQUIRE(status.broker_state == JFJochState::Idle);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == nimages);

    REQUIRE_NOTHROW(writer_0_future.get());
    REQUIRE_NOTHROW(writer_1_future.get());
    REQUIRE_NOTHROW(writer_2_future.get());
}

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

    ZMQContext zmq_context;

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

   AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        test->EnableLogging(&logger);
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int frame = 1; frame <= nimages; frame++)
                test->AddModule(frame, m, image.data());
        }
        aq_devices.Add(std::move(test));
    }
    
    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_test";
    setup.images_per_trigger = 2 * nimages;
    setup.photon_energy_keV = 12.4;
    setup.data_file_count = 1;
    setup.save_calibration = false;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    logger.Info("Cancelling");
    REQUIRE_NOTHROW(state_machine.Cancel());
    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 0.5);
    REQUIRE(statistics.images_collected == 5);
    REQUIRE(statistics.max_image_number_sent == 4);
    REQUIRE(statistics.cancelled);

    REQUIRE_NOTHROW(writer_future.get());
}

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

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0).PreviewPeriod(
            5ms);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int image_num = 1; image_num <= nimages; image_num++)
                test->AddModule(image_num, m, image.data());
        }
        test->Terminate();
        aq_devices.Add(std::move(test));
    }
    ZMQContext zmq_context;


    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    ZMQPreviewPublisher preview(zmq_context, "inproc://#2");
    fpga_receiver.PreviewPublisher(&preview).NumThreads(1);

    ZMQSocket rcv_preview_socket(zmq_context, ZMQSocketType::Sub);
    REQUIRE_NOTHROW(rcv_preview_socket.Connect("inproc://#2"));
    rcv_preview_socket.SubscribeAll();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 100;
    setup.file_prefix = "integration_test_preview";
    setup.images_per_trigger = 5;
    setup.photon_energy_keV = 12.4;
    setup.beam_x_pxl = 123.0;
    setup.data_file_count = 1;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    std::string s;

    JFJochFrameDeserializer deserializer;

    // Pixel mask
    REQUIRE(rcv_preview_socket.Receive(s, false) > 0);
    deserializer.Process((uint8_t *) s.data(), s.size());
    REQUIRE(deserializer.GetType() ==  JFJochFrameDeserializer::Type::START);
    auto start_msg = deserializer.GetStartMessage();
    REQUIRE(start_msg.image_size_x == state_machine.NotThreadSafe_Experiment().GetXPixelsNum());
    REQUIRE(start_msg.image_size_y == state_machine.NotThreadSafe_Experiment().GetYPixelsNum());

    // First frame
    REQUIRE(rcv_preview_socket.Receive(s, false) > 0);

    deserializer.Process((uint8_t *) s.data(), s.size());
    REQUIRE(deserializer.GetType() ==  JFJochFrameDeserializer::Type::IMAGE);
    auto image_msg = deserializer.GetDataMessage();

    REQUIRE(image_msg.number == 0);

    // Check no more frames waiting
    REQUIRE(rcv_preview_socket.Receive(s, false) == -1);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 5);

    REQUIRE_NOTHROW(writer_future.get());
}

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

    RegisterHDF5Filter();

    int64_t nimages = 5;
    int64_t ndatastream = 2;
    int64_t nmodules = 4;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0).PreviewPeriod(
            5ms);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    AcquisitionDeviceGroup aq_devices;

    for (int i = 0; i < ndatastream; i++) {
        auto test = std::make_unique<MockAcquisitionDevice>(i, 256);
        for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(i); m++) {
            for (int image_num = 1; image_num <= nimages; image_num++)
                test->AddModule(image_num, m, image.data());
        }
        test->Terminate();
        aq_devices.Add(std::move(test));
    }
    ZMQContext zmq_context;

    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    ZMQPreviewPublisher preview(zmq_context, "inproc://#2");
    fpga_receiver.PreviewPublisher(&preview).NumThreads(1);

    ZMQSocket rcv_preview_socket(zmq_context, ZMQSocketType::Sub);
    REQUIRE_NOTHROW(rcv_preview_socket.Connect("inproc://#2"));
    rcv_preview_socket.SubscribeAll();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 100;
    setup.images_per_trigger = 5;
    setup.photon_energy_keV = 12.4;
    setup.beam_x_pxl = 123.0;
    setup.beam_y_pxl = 878.0;
    setup.data_file_count = 1;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    std::string s;
    JFJochFrameDeserializer deserializer;

    // Pixel mask
    REQUIRE(rcv_preview_socket.Receive(s, false) > 0);
    deserializer.Process((uint8_t *) s.data(), s.size());
    REQUIRE(deserializer.GetType() ==  JFJochFrameDeserializer::Type::START);
    auto start_msg = deserializer.GetStartMessage();
    REQUIRE(start_msg.image_size_x == state_machine.NotThreadSafe_Experiment().GetXPixelsNum());
    REQUIRE(start_msg.image_size_y == state_machine.NotThreadSafe_Experiment().GetYPixelsNum());

    // First frame
    REQUIRE(rcv_preview_socket.Receive(s, false) > 0);

    deserializer.Process((uint8_t *) s.data(), s.size());
    REQUIRE(deserializer.GetType() ==  JFJochFrameDeserializer::Type::IMAGE);
    auto image_msg = deserializer.GetDataMessage();

    REQUIRE(image_msg.number == 0);

    // Check no more frames waiting
    REQUIRE(rcv_preview_socket.Receive(s, false) == -1);

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 5);
}

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

    RegisterHDF5Filter();

    size_t nimages = 1;
    int64_t ndatastream = 1;
    int64_t nmodules = 8;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);

    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));
    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    HDF5File data("../../tests/test_data/compression_benchmark.h5", false, false, false);
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    std::vector<int16_t> image_conv(nimages * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

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

    std::vector<int16_t> image_raw_geom(
            nimages * state_machine.NotThreadSafe_Experiment().GetModulesNum() * RAW_MODULE_SIZE * sizeof(int16_t));
    for (int i = 0; i < nimages; i++)
        ConvertedToRawGeometry(state_machine.NotThreadSafe_Experiment(),
                               image_raw_geom.data() +
                               i * RAW_MODULE_SIZE * state_machine.NotThreadSafe_Experiment().GetModulesNum(),
                               image_conv.data() + i * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);


    AcquisitionDeviceGroup aq_devices;

    auto test = std::make_unique<MockAcquisitionDevice>(0, 256);

    for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(0); m++) {
        test->AddModule(1, m, (uint16_t *) (image_raw_geom.data() + m * RAW_MODULE_SIZE));
    }
    test->Terminate();

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

    ZMQContext zmq_context;


    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 75;
    setup.file_prefix = "spot_finding_test";
    setup.images_per_trigger = nimages;
    setup.photon_energy_keV = 12.4;
    setup.beam_x_pxl = 1090;
    setup.beam_y_pxl = 1136;
    setup.unit_cell = UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha = 90.0, .beta = 90.0, .gamma = 90.0};
    setup.data_file_count = 1;

    DataProcessingSettings settings{};
    settings.signal_to_noise_threshold = 4;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 3;
    settings.max_pix_per_spot = 200;
    settings.low_resolution_limit = 80.0;
    settings.high_resolution_limit = 2.5;

    REQUIRE_NOTHROW(state_machine.SetDataProcessingSettings(settings));

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == 1);

    REQUIRE_NOTHROW(writer_future.get());
}
/*
TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index");

    RegisterHDF5Filter();

    size_t nimages = 24;
    int64_t ndatastream = 1;
    int64_t nmodules = 8;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0)
            .SpotFindingPeriod(10ms);
    services.Writer("unix:writer_test", "inproc://#1").Receiver(&fpga_receiver);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    HDF5File data("../../tests/test_data/compression_benchmark.h5", false, false, false);
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    std::vector<int16_t> image_conv(nimages * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

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

    std::vector<int16_t> image_raw_geom(
            nimages * state_machine.NotThreadSafe_Experiment().GetModulesNum() * RAW_MODULE_SIZE * sizeof(int16_t));
    for (int i = 0; i < nimages; i++)
        ConvertedToRawGeometry(state_machine.NotThreadSafe_Experiment(),
                               image_raw_geom.data() +
                               i * RAW_MODULE_SIZE * state_machine.NotThreadSafe_Experiment().GetModulesNum(),
                               image_conv.data() + i * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);


    AcquisitionDeviceGroup aq_devices;

    auto *test = new MockAcquisitionDevice(0, 512);

    for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(0); m++) {
        for (int image_num = 1; image_num <= nimages; image_num++)
            test->AddModule(image_num, m, (uint16_t *) (image_raw_geom.data()
                                                        + (m + (image_num - 1) *
                                                               state_machine.NotThreadSafe_Experiment().GetModulesNum(
                                                                       0))
                                                          * RAW_MODULE_SIZE));
    }
    test->Terminate();

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

    ZMQContext zmq_context;


    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(tmp_devices, logger, pusher);
    JFJochWriterService writer(zmq_context, logger);;

    auto writer_server = gRPCServer("unix:writer_test", writer);

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 75;
        setup.file_prefix = "spot_finding_test";
    setup.images_per_trigger = nimages;
    setup.photon_energy_keV = 12.4;
    setup.beam_x_pxl = 1090;
    setup.beam_y_pxl = 1136;
    setup.mutable_unit_cell()->set_a(36.9);
    setup.mutable_unit_cell()->set_b(78.95);
    setup.mutable_unit_cell()->set_c(78.95);
    setup.mutable_unit_cell()->set_alpha(90.0);
    setup.mutable_unit_cell()->set_beta(90.0);
    setup.mutable_unit_cell()->set_gamma(90.0);
    setup.data_file_count = 1;

    JFJochProtoBuf::DataProcessingSettings settings;
    settings.signal_to_noise_threshold = 4;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 3;
    settings.max_pix_per_spot = 200;
    settings.low_resolution_limit = 80.0;
    settings.high_resolution_limit = 2.5;
    settings.local_bkg_size = 5;

    REQUIRE_NOTHROW(state_machine.SetDataProcessingSettings(settings));

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.has_indexing_rate);
    REQUIRE(statistics.indexing_rate == 1.0);
    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == nimages);


    writer_server->Shutdown();
}*/


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

    RegisterHDF5Filter();

    size_t nimages = 24;
    int64_t ndatastream = 1;
    int64_t nmodules = 8;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0);
    state_machine.NotThreadSafe_Experiment().LowQForRadialInt_recipA(0.5).HighQForRadialInt_recipA(3.5)
            .QSpacingForRadialInt_recipA(1.0);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    HDF5File data("../../tests/test_data/compression_benchmark.h5", false, false, false);
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    std::vector<int16_t> image_conv(nimages * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

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

    std::vector<int16_t> image_raw_geom(
            nimages * state_machine.NotThreadSafe_Experiment().GetModulesNum() * RAW_MODULE_SIZE * sizeof(int16_t));
    for (int i = 0; i < nimages; i++)
        ConvertedToRawGeometry(state_machine.NotThreadSafe_Experiment(),
                               image_raw_geom.data() +
                               i * RAW_MODULE_SIZE * state_machine.NotThreadSafe_Experiment().GetModulesNum(),
                               image_conv.data() + i * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);


    AcquisitionDeviceGroup aq_devices;

    auto test = std::make_unique<MockAcquisitionDevice>(0, 512);

    for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(0); m++) {
        for (int image_num = 1; image_num <= nimages; image_num++)
            test->AddModule(image_num, m, (uint16_t *) (image_raw_geom.data()
                                                        + (m + (image_num - 1) *
                                                               state_machine.NotThreadSafe_Experiment().GetModulesNum(
                                                                       0))
                                                          * RAW_MODULE_SIZE));
    }
    test->Terminate();

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

    ZMQContext zmq_context;


    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(aq_devices, logger, pusher);
    services.Receiver(&fpga_receiver);

    StreamWriter writer(zmq_context, logger, "inproc://#1");
    auto writer_future = writer.RunFuture();

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 75;
    setup.file_prefix = "spot_finding_test";
    setup.images_per_trigger = nimages;
    setup.photon_energy_keV = 12.4;
    setup.beam_x_pxl = 1090;
    setup.beam_y_pxl = 1136;
    setup.unit_cell = UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha = 90.0, .beta = 90.0, .gamma = 90.0};
    setup.data_file_count = 4;

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    auto rad_int = state_machine.GetRadialIntegrationProfiles();
    REQUIRE(rad_int.profiles.size() == 4+1);
    auto &plot_map = rad_int.profiles;

    CHECK(plot_map[1].title == "file0");
    CHECK(plot_map[1].plot.x.size() == 3);
    CHECK(plot_map[1].plot.y.size() == 3);

    CHECK(plot_map[4].title == "file3");
    REQUIRE(plot_map[4].plot.x.size() == 3);
    CHECK(plot_map[4].plot.x[0] == Approx(1.0));
    CHECK(plot_map[4].plot.y.size() == 3);

    CHECK(plot_map[0].title == "dataset");
    CHECK(plot_map[0].plot.x.size() == 3);

    REQUIRE_NOTHROW(writer_future.get());
}

/*
TEST_CASE("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_sum", "[JFJochReceiver]") {
    Logger logger("JFJochIntegrationTest_ZMQ_lysozyme_spot_and_index_sum");

    RegisterHDF5Filter();

    size_t nimages = 24;
    int64_t ndatastream = 1;
    int64_t nmodules = 8;

    JFJochServices services(logger);
    JFJochStateMachine state_machine(services, logger);
    state_machine.AddDetectorSetup(DetectorGeometry(ndatastream * nmodules, 2, 8, 36));

    state_machine.NotThreadSafe_Experiment().DataStreams(ndatastream);
    state_machine.NotThreadSafe_Experiment().PedestalG0Frames(0).PedestalG1Frames(0).PedestalG2Frames(0)
            .SpotFindingPeriod(10ms);
    services.Writer("unix:writer_test", "inproc://#1").Receiver(&fpga_receiver);

    logger.Verbose(true);

    std::vector<uint16_t> image(RAW_MODULE_SIZE);

    HDF5File data("../../tests/test_data/compression_benchmark.h5", false, false, false);
    HDF5DataSet dataset(data, "/entry/data/data");
    HDF5DataSpace file_space(dataset);

    std::vector<int16_t> image_conv(nimages * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);

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

    std::vector<int16_t> image_raw_geom(
            nimages * state_machine.NotThreadSafe_Experiment().GetModulesNum() * RAW_MODULE_SIZE * sizeof(int16_t));
    for (int i = 0; i < nimages; i++)
        ConvertedToRawGeometry(state_machine.NotThreadSafe_Experiment(),
                               image_raw_geom.data() +
                               i * RAW_MODULE_SIZE * state_machine.NotThreadSafe_Experiment().GetModulesNum(),
                               image_conv.data() + i * file_space.GetDimensions()[1] * file_space.GetDimensions()[2]);


    AcquisitionDeviceGroup aq_devices;

    auto *test = new MockAcquisitionDevice(0, 512);

    std::vector<uint16_t> pedestal(RAW_MODULE_SIZE, 3500);

    for (int m = 0; m < state_machine.NotThreadSafe_Experiment().GetModulesNum(0); m++) {
        for (int image_num = 1; image_num <= nimages; image_num++)
            test->AddModule(image_num, m, (uint16_t *) (image_raw_geom.data()
                                                        + (m + (image_num - 1) *
                                                               state_machine.NotThreadSafe_Experiment().GetModulesNum(
                                                                       0))
                                                          * RAW_MODULE_SIZE));
    }
    test->Terminate();

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

    ZMQContext zmq_context;


    ZMQImagePusher pusher(zmq_context, {"inproc://#1"});
    JFJochReceiverService fpga_receiver(tmp_devices, logger, pusher);
    JFJochWriterService writer(zmq_context, logger);

    auto writer_server = gRPCServer("unix:writer_test", writer);

    REQUIRE_NOTHROW(state_machine.Initialize());
    logger.Info("Initialized");

    DatasetSettings setup{};
    setup.ntrigger = 1;
    setup.detector_distance_mm = 75;
    setup.summation = 3;
    setup.file_prefix = "spot_finding_test";
    setup.set_images_per_trigger(nimages / 3);
    setup.photon_energy_keV = 12.4;
    setup.beam_x_pxl = 1090;
    setup.beam_y_pxl = 1136;
    setup.mutable_unit_cell()->set_a(36.9);
    setup.mutable_unit_cell()->set_b(78.95);
    setup.mutable_unit_cell()->set_c(78.95);
    setup.mutable_unit_cell()->set_alpha(90.0);
    setup.mutable_unit_cell()->set_beta(90.0);
    setup.mutable_unit_cell()->set_gamma(90.0);
    setup.data_file_count = 1;

    JFJochProtoBuf::DataProcessingSettings settings;
    settings.signal_to_noise_threshold = 4;
    settings.photon_count_threshold = 5;
    settings.min_pix_per_spot = 3;
    settings.max_pix_per_spot = 200;
    settings.low_resolution_limit = 80.0;
    settings.high_resolution_limit = 2.5;
    settings.local_bkg_size = 5;

    REQUIRE_NOTHROW(state_machine.SetDataProcessingSettings(settings));

    REQUIRE_NOTHROW(state_machine.Start(setup));
    logger.Info("Started measurement");

    REQUIRE_NOTHROW(state_machine.Stop());
    logger.Info("Stopped measurement");

    auto tmp = state_machine.GetMeasurementStatistics();
    REQUIRE(tmp.has_value());
    auto statistics = tmp.value();

    REQUIRE(statistics.collection_efficiency == 1.0);
    REQUIRE(statistics.images_collected == nimages / 3);
    REQUIRE(statistics.has_indexing_rate);
    REQUIRE(statistics.indexing_rate == 1.0);


    writer_server->Shutdown();
}*/
//#endif