// Copyright (2019-2022) Paul Scherrer Institute
// SPDX-License-Identifier: GPL-3.0-or-later

#include "JFJochReceiver.h"

#include <thread>

#include "../image_analysis/GPUImageAnalysis.h"
#include "../jungfrau/JFPedestalCalc.h"
#include "../image_analysis/IndexerWrapper.h"

#ifdef JFJOCH_USE_NUMA
#include <numa.h>
#endif

inline std::string time_UTC(const std::chrono::time_point<std::chrono::system_clock> &input) {
    auto time_ms = std::chrono::duration_cast<std::chrono::milliseconds>(input.time_since_epoch()).count();

    char buf1[255], buf2[255];
    time_t time = time_ms / (1000);
    strftime(buf1, sizeof(buf1), "%FT%T", gmtime(&time));
    snprintf(buf2, sizeof(buf2), ".%06ld", time_ms%1000);
    return std::string(buf1) + std::string(buf2) + "Z";
}

JFJochReceiver::JFJochReceiver(const JFJochProtoBuf::ReceiverInput &settings,
                               std::vector<AcquisitionDevice *> &in_aq_device,
                               ImagePusher &in_image_sender,
                               Logger &in_logger, int64_t in_forward_and_sum_nthreads,
                               int64_t in_send_buffer_count,
                               ZMQPreviewPublisher* in_preview_publisher,
                               ZMQPreviewPublisher* in_preview_publisher_indexed) :
        experiment(settings.jungfraujoch_settings()),
        acquisition_device(in_aq_device),
        logger(in_logger),
        image_pusher(in_image_sender),
        frame_transformation_nthreads((experiment.GetSummation() >= threaded_summation_threshold) ?
                                      2 : in_forward_and_sum_nthreads),
        preview_publisher(in_preview_publisher),
        preview_publisher_indexed(in_preview_publisher_indexed),
        ndatastreams(experiment.GetDataStreamsNum()),
        data_acquisition_ready(ndatastreams),
        frame_transformation_ready((experiment.GetImageNum() > 0) ? frame_transformation_nthreads : 0),
        send_buffer_count(in_send_buffer_count),
        send_buffer(send_buffer_size * send_buffer_count),
        indexing_solution_per_file(experiment.GetDataFileCount())
{
    if (settings.has_calibration()) {
        calib.emplace(settings.calibration());
        one_byte_mask = calib->CalculateOneByteMask(experiment);
    } else {
        one_byte_mask.resize(experiment.GetPixelsNum());
        for (auto &i: one_byte_mask) i = 1;
    }

    for (uint32_t i = 0; i < send_buffer_count; i++) {
        send_buffer_avail.Put(i);
        send_buffer_zero_copy_ret_val.emplace_back(send_buffer_avail, i);
    }

    if (experiment.GetConversionOnCPU())
        PrepareConversionOnCPU();

    if (!experiment.CheckGitSha1Consistent())
        logger.Warning(experiment.CheckGitSha1Msg());

    push_images_to_writer = (experiment.GetImageNum() > 0) && (!experiment.GetFilePrefix().empty());

    if (acquisition_device.size() < ndatastreams)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Number of acquisition devices has to match data streams");
    if (frame_transformation_nthreads <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Number of threads must be more than zero");

    preview_stride = experiment.GetPreviewStride();
    spotfinder_stride = experiment.GetSpotFindingStride();

    logger.Info("Image stride for data analysis:   preview {}, spot finding/radial integration {}",
                preview_stride, spotfinder_stride);
    if (experiment.GetDetectorMode() == DetectorMode::Conversion) {
        if (preview_publisher != nullptr)
            preview_publisher->Start(experiment, calib.value());
        if (preview_publisher_indexed != nullptr)
            preview_publisher_indexed->Start(experiment, calib.value());

        if (!GPUImageAnalysis::GPUPresent())
            logger.Info("GPU support missing");

        rad_int_mapping = std::make_unique<RadialIntegrationMapping>(experiment, one_byte_mask.data());
        rad_int_profile = std::make_unique<RadialIntegrationProfile>(*rad_int_mapping, experiment);

        for (int i = 0; i < experiment.GetDataFileCount(); i++)
            rad_int_profile_per_file.emplace_back(std::make_unique<RadialIntegrationProfile>(*rad_int_mapping, experiment));

        spot_finder_mask = calib->CalculateOneByteMask(experiment);
    }

    for (int d = 0; d < ndatastreams; d++) {
        if (calib)
            acquisition_device[d]->InitializeCalibration(experiment, calib.value());
        acquisition_device[d]->PrepareAction(experiment);

        logger.Debug("Acquisition device {} prepared", d);
    }

    for (int d = 0; d < ndatastreams; d++)
        data_acquisition_futures.emplace_back(std::async(std::launch::async, &JFJochReceiver::AcquireThread,
                                                         this, d));

    logger.Info("Data acquisition devices ready");

    if ((experiment.GetDetectorMode() == DetectorMode::PedestalG0)
        || (experiment.GetDetectorMode() == DetectorMode::PedestalG1)
        || (experiment.GetDetectorMode() == DetectorMode::PedestalG2)) {

        if (experiment.GetImageNum() > 0) {
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Saving and calculating pedestal is not supported for the time being");
        }

        if (experiment.GetDetectorMode() == DetectorMode::PedestalG0) {
            pedestal_result.resize(experiment.GetModulesNum() * experiment.GetStorageCellNumber());
            for (int s = 0; s < experiment.GetStorageCellNumber(); s++) {
                for (int d = 0; d < ndatastreams; d++) {
                    for (int m = 0; m < experiment.GetModulesNum(d); m++) {
                        auto handle = std::async(std::launch::async, &JFJochReceiver::MeasurePedestalThread, this, d, m,
                                                 s);
                        frame_transformation_futures.emplace_back(std::move(handle));
                    }
                }
            }
        } else {
            pedestal_result.resize(experiment.GetModulesNum());
            for (int d = 0; d < ndatastreams; d++) {
                for (int m = 0; m < experiment.GetModulesNum(d); m++) {
                    auto handle = std::async(std::launch::async, &JFJochReceiver::MeasurePedestalThread, this, d, m, 0);
                    frame_transformation_futures.emplace_back(std::move(handle));
                }
            }
        }

        logger.Info("Pedestal threads ready");
    }

    if (experiment.GetImageNum() > 0) {
        logger.Info("Data file count {}", experiment.GetDataFileCount());

        if (push_images_to_writer) {
            StartMessage message{};
            experiment.FillMessage(message);
            message.arm_date = time_UTC(std::chrono::system_clock::now());

            if (calib)
                message.pixel_mask["sc0"] = calib->CalculateNexusMask(experiment, 0);

            if (rad_int_mapping) {
                message.rad_int_bin_number = rad_int_mapping->GetBinNumber();
                message.rad_int_bin_to_q = rad_int_mapping->GetBinToQ();
                message.rad_int_solid_angle_corr = rad_int_mapping->GetSolidAngleCorr();
            } else
                message.rad_int_bin_number = 0;

            image_pusher.StartDataCollection(message);
        }

        for (int i = 0; i < experiment.GetImageNum(); i++)
            images_to_go.Put(i);

        // Setup frames summation and forwarding
        for (int i = 0; i < frame_transformation_nthreads; i++) {
            auto handle = std::async(std::launch::async, &JFJochReceiver::FrameTransformationThread, this);
            frame_transformation_futures.emplace_back(std::move(handle));
        }

        logger.Info("Image compression/forwarding threads started");

        frame_transformation_ready.wait();
        logger.Info("Image compression/forwarding threads ready");
    }

    data_acquisition_ready.wait();

    logger.Info("Acquisition devices ready");

    start_time = std::chrono::system_clock::now();
    logger.Info("Receiving data started");

    measurement = std::async(std::launch::async, &JFJochReceiver::FinalizeMeasurement, this);
}

void JFJochReceiver::AcquireThread(uint16_t data_stream) {
    PinThreadToDevice(data_stream);

    try {
        frame_transformation_ready.wait();
        logger.Debug("Device thread {} start FPGA action", data_stream);
        acquisition_device[data_stream]->StartAction(experiment);
    } catch (const JFJochException &e) {
        Cancel(e);
        data_acquisition_ready.count_down();
    }

    data_acquisition_ready.count_down();

    try {
        logger.Debug("Device thread {} wait for FPGA action complete", data_stream);
        acquisition_device[data_stream]->WaitForActionComplete();
    } catch (const JFJochException &e) {
        Cancel(e);
    }

    logger.Info("Device thread {} done", data_stream);
}

void JFJochReceiver::MeasurePedestalThread(uint16_t data_stream, uint16_t module_number, uint16_t storage_cell) {
    PinThreadToDevice(data_stream);

    JFPedestalCalc pedestal_calc(experiment);

    bool storage_cell_G1G2 = (experiment.GetStorageCellNumber() > 1)
                             && ((experiment.GetDetectorMode() == DetectorMode::PedestalG1)
                                 || (experiment.GetDetectorMode() == DetectorMode::PedestalG2));
    size_t staring_frame;
    size_t frame_stride;

    size_t offset = experiment.GetFirstModuleOfDataStream(data_stream) + module_number;
    if (experiment.GetDetectorMode() == DetectorMode::PedestalG0) {
        staring_frame = storage_cell;
        frame_stride = experiment.GetStorageCellNumber();
        offset += experiment.GetModulesNum() * storage_cell;
    } else {
        staring_frame = 0;
        frame_stride = 1;
    }

    uint32_t storage_cell_header = UINT32_MAX;
    logger.Debug("Pedestal calculation thread for data stream {} module {} storage cell {} starting",
                 data_stream, module_number, storage_cell);
    try {
        for (size_t frame = staring_frame; frame < experiment.GetFrameNum(); frame += frame_stride) {
            // Frame will be processed only if one already collects frame+2
            acquisition_device[data_stream]->WaitForFrame(frame + 2, module_number);
            if (!storage_cell_G1G2 || (frame % 2 == 1)) {
                // Partial packets will bring more problems, than benefit
                if (acquisition_device[data_stream]->IsFullModuleCollected(frame, module_number)) {
                    pedestal_calc.AnalyzeImage((uint16_t *) acquisition_device[data_stream]->GetFrameBuffer(frame, module_number));
                }
                auto tmp = acquisition_device[data_stream]->GetJFInfo(frame, module_number);
                storage_cell_header = (tmp >> 8) & 0xF;
            }

            acquisition_device[data_stream]->FrameBufferRelease(frame, module_number);

            UpdateMaxDelay(acquisition_device[data_stream]->CalculateDelay(frame, module_number));
            UpdateMaxImage(frame);
        }

        if (experiment.GetDetectorMode() == DetectorMode::PedestalG0)
            pedestal_calc.Export(pedestal_result[offset], 2);
        else
            pedestal_calc.Export(pedestal_result[offset]);
        pedestal_result[offset].SetFrameCount(experiment.GetFrameNum());
        pedestal_result[offset].SetCollectionTime(start_time.time_since_epoch().count() / 1e9);
    } catch (const JFJochException &e) { Cancel(e); }
    logger.Info("Pedestal calculation thread for data stream {} module {} storage cell {} -> header {} done",
                 data_stream, module_number, storage_cell, storage_cell_header);
}

void JFJochReceiver::MiniSummationThread(int d, int m, size_t image_number, bool &send_image,
                                            FrameTransformation &transformation, DataMessage &message) {
    for (int j = 0; j < experiment.GetSummation(); j++) {
        size_t frame_number = image_number * experiment.GetSummation() + j;
        acquisition_device[d]->WaitForFrame(frame_number + 2);

        const int16_t *src;

        if (acquisition_device[d]->IsFullModuleCollected(frame_number, m)) {
            src = acquisition_device[d]->GetFrameBuffer(frame_number, m);

            if (!send_image) {
                // the information is for first module/frame that was collected in full
                message.bunch_id = acquisition_device[d]->GetBunchID(frame_number, m);
                message.jf_info  = acquisition_device[d]->GetJFInfo(frame_number, m);
                message.storage_cell = (message.jf_info >> 8) & 0xF;
                message.timestamp = acquisition_device[d]->GetTimestamp(frame_number, m);
                message.exptime = acquisition_device[d]->GetExptime(frame_number, m);
            }
            send_image = true;
        } else
            src = acquisition_device[d]->GetErrorFrameBuffer();

        if (experiment.GetConversionOnCPU()) {
            auto &conv = fixed_point_conversion.at(experiment.GetFirstModuleOfDataStream(d) + m);
            transformation.ProcessModule(conv, src, m, d);
        } else
            transformation.ProcessModule(src, m, d);

        acquisition_device[d]->FrameBufferRelease(frame_number, m);
        UpdateMaxDelay(acquisition_device[d]->CalculateDelay(frame_number, m));
    }
}

void JFJochReceiver::FrameTransformationThread() {
    FrameTransformation transformation(experiment);

     std::unique_ptr<GPUImageAnalysis> spot_finder;

    try {
        if (rad_int_mapping)
            spot_finder = std::make_unique<GPUImageAnalysis>(experiment.GetXPixelsNum(),
                                                             experiment.GetYPixelsNum(),
                                                             one_byte_mask, *rad_int_mapping);
        else
            spot_finder = std::make_unique<GPUImageAnalysis>(experiment.GetXPixelsNum(),
                                                             experiment.GetYPixelsNum(),
                                                             one_byte_mask);
        spot_finder->SetInputBuffer(transformation.GetPreview16BitImage());
        spot_finder->RegisterBuffer();
    } catch (const JFJochException& e) {
        frame_transformation_ready.count_down();
        logger.Error("Error creating GPU spot finder");
        Cancel(e);
    }

    std::vector<char> writer_buffer(experiment.GetMaxCompressedSize());
    std::vector<int16_t> conversion_buffer(RAW_MODULE_SIZE);

    uint64_t image_number;

    frame_transformation_ready.count_down();

    std::unique_ptr<IndexerWrapper> indexer;
    if (experiment.HasUnitCell()) {
        indexer = std::make_unique<IndexerWrapper>();
        indexer->Setup(experiment.GetUnitCell());
    }

    while (images_to_go.Get(image_number) != 0) {
        try {
            DataMessage message{};
            message.number = image_number;
            message.timestamp_base = 10*1000*1000;
            message.exptime_base = 10*1000*1000;
            message.indexing_result = 0;

            bool send_preview = false;
            bool send_bkg_estimate = false;
            bool calculate_spots = false;

            if ((preview_publisher != nullptr) && (preview_stride > 0) && (image_number % preview_stride == 0))
                send_preview = true;

            bool send_image = false; // We send image if at least one module was collected in full

            if (GPUImageAnalysis::GPUPresent()) {
                if (((spotfinder_stride > 0) && (image_number % spotfinder_stride == 0)) || send_preview) {
                    calculate_spots = true;
                    if (rad_int_mapping)
                        send_bkg_estimate = true;
                }
            }

            if (experiment.GetSummation() >= threaded_summation_threshold) {
                std::vector<std::future<void>> mini_summation_threads;
                for (int d = 0; d < ndatastreams; d++)
                    for (int m = 0; m < experiment.GetModulesNum(d); m++)
                        mini_summation_threads.emplace_back(std::async(std::launch::async, &JFJochReceiver::MiniSummationThread,
                                                                       this, d, m, image_number,
                                                                       std::ref(send_image), std::ref(transformation),
                                                                       std::ref(message)));
                message.receiver_aq_dev_delay = max_delay;
            } else {
                for (int j = 0; j < experiment.GetSummation(); j++) {
                    size_t frame_number = image_number * experiment.GetSummation() + j;

                    for (int d = 0; d < ndatastreams; d++) {
                        acquisition_device[d]->WaitForFrame(frame_number + 2);

                        for (int m = 0; m < experiment.GetModulesNum(d); m++) {
                            const int16_t *src;

                            if (acquisition_device[d]->IsFullModuleCollected(frame_number, m)) {
                                src = acquisition_device[d]->GetFrameBuffer(frame_number, m);

                                if (!send_image) {
                                    // the information is for first module/frame that was collected in full
                                    message.bunch_id = acquisition_device[d]->GetBunchID(frame_number, m);
                                    message.jf_info = acquisition_device[d]->GetJFInfo(frame_number, m);
                                    message.timestamp = acquisition_device[d]->GetTimestamp(frame_number, m);
                                    message.exptime = acquisition_device[d]->GetExptime(frame_number, m);
                                    message.storage_cell = (message.jf_info >> 8) & 0xF;
                                }
                                send_image = true;
                            } else
                                src = acquisition_device[d]->GetErrorFrameBuffer();

                            if (experiment.GetConversionOnCPU()) {
                                auto &conv = fixed_point_conversion.at(experiment.GetFirstModuleOfDataStream(d) + m);
                                transformation.ProcessModule(conv, src, m, d);
                            } else
                                transformation.ProcessModule(src, m, d);

                            acquisition_device[d]->FrameBufferRelease(frame_number, m);
                        }
                        auto delay = acquisition_device[d]->CalculateDelay(frame_number);
                        UpdateMaxDelay(delay);
                        if (delay > message.receiver_aq_dev_delay)
                            message.receiver_aq_dev_delay = delay;
                    }
                }
            }

            if (send_image) {
                transformation.Pack();

                std::vector<DiffractionSpot> spots;
                auto local_data_processing_settings = GetDataProcessingSettings();

                // Spot finding is async, so it can be sandwiched between sending image and other tasks
                if (calculate_spots || send_bkg_estimate)
                    spot_finder->LoadDataToGPU(!experiment.GetApplyPixelMaskInFPGA());

                if (calculate_spots)
                    spot_finder->RunSpotFinder(local_data_processing_settings);

                if (send_bkg_estimate)
                    spot_finder->RunRadialIntegration();

                if (calculate_spots) {
                    spot_finder->GetSpotFinderResults(experiment, GetDataProcessingSettings(), spots);
                    for (const auto &spot: spots)
                        message.spots.push_back(spot);
                    spot_count.AddElement(image_number, spots.size());

                    if (indexer) {
                        std::vector<Coord> recip;
                        for (const auto &i: spots)
                            recip.push_back(i.ReciprocalCoord(experiment));

                        auto indexer_result = indexer->Run(recip);

                        if (!indexer_result.empty()) {
                            message.indexing_result = 2;
                            indexing_solution.AddElement(image_number, 1);
                            indexing_solution_per_file.Add(image_number % experiment.GetDataFileCount(), 1);
                            for (int i = 0; i < recip.size(); i++)
                                message.spots[i].indexed = indexer_result[0].indexed_spots[i];
                            indexer_result[0].l.Save(message.indexing_lattice);
                            if (preview_publisher_indexed)
                                preview_publisher_indexed->Publish(experiment,
                                                                   transformation.GetPreview16BitImage(),
                                                                   message);
                        } else {
                            message.indexing_result = 1;
                            indexing_solution.AddElement(image_number, 0);
                            indexing_solution_per_file.Add(image_number % experiment.GetDataFileCount(), 0);
                        }
                    }
                }

                if (send_bkg_estimate) {
                    uint16_t rad_int_min_bin = std::floor(
                            rad_int_mapping->QToBin(local_data_processing_settings.bkg_estimate_low_q()));
                    uint16_t rad_int_max_bin = std::ceil(
                            rad_int_mapping->QToBin(local_data_processing_settings.bkg_estimate_high_q()));
                    float bkg_estimate_val = spot_finder->GetRadialIntegrationRangeValue(rad_int_min_bin, rad_int_max_bin);
                    bkg_estimate.AddElement(image_number, bkg_estimate_val);

                    spot_finder->GetRadialIntegrationProfile(message.rad_int_profile);
                    rad_int_profile->Add(spot_finder->GetRadialIntegrationSum(),
                                         spot_finder->GetRadialIntegrationCount());

                    if (image_number % experiment.GetDataFileCount() < rad_int_profile_per_file.size())
                        rad_int_profile_per_file[image_number % experiment.GetDataFileCount()]
                                ->Add(spot_finder->GetRadialIntegrationSum(),
                                     spot_finder->GetRadialIntegrationCount());
                }

                if (send_preview)
                    preview_publisher->Publish(experiment,
                                               transformation.GetPreview16BitImage(),
                                               message);

                if (push_images_to_writer) {
                    message.receiver_available_send_buffers = GetAvailableSendBuffers();

                    auto send_buffer_handle = send_buffer_avail.GetBlocking();
                    auto ptr = send_buffer.data() + send_buffer_size * send_buffer_handle;
                    JFJochFrameSerializer serializer(ptr, send_buffer_size);
                    PrepareCBORImage(message, experiment, nullptr, 0);
                    serializer.SerializeImage(message);
                    if (serializer.GetRemainingBuffer() < experiment.GetMaxCompressedSize())
                        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                                              "Not enough memory to save image");
                    size_t image_size = transformation.SaveCompressedImage(ptr + serializer.GetImageAppendOffset());
                    serializer.AppendImage(image_size);
                    image_pusher.SendImage(ptr, serializer.GetBufferSize(), image_number,
                                           &send_buffer_zero_copy_ret_val[send_buffer_handle]);
                    compressed_size += image_size;
                }

                UpdateMaxImage(image_number);
                images_sent++;
            }
        } catch (const JFJochException &e) { Cancel(e); }
    }

    spot_finder->UnregisterBuffer();

    logger.Debug("Sum&compression thread done");
}

void JFJochReceiver::GetStatistics(JFJochProtoBuf::ReceiverOutput &ret) const {
    uint64_t expected_packets = 0;
    uint64_t received_packets = 0;

    for (int d = 0; d < ndatastreams; d++) {
        acquisition_device[d]->SaveStatistics(experiment, *ret.add_device_statistics());
        expected_packets += ret.device_statistics(d).packets_expected();
        received_packets += ret.device_statistics(d).good_packets();
    }

    if ((expected_packets == received_packets) || (expected_packets == 0))
        ret.set_efficiency(1.0);
    else
        ret.set_efficiency(received_packets / static_cast<double>(expected_packets));

    ret.set_compressed_size(compressed_size);
    ret.set_max_image_number_sent(max_image_number_sent);

    if ((experiment.GetImageNum() > 0) && (compressed_size > 0)) {
        ret.set_compressed_ratio( static_cast<double> (images_sent
                                                       * experiment.GetPixelDepth()
                                                       * experiment.GetModulesNum()
                                                       * RAW_MODULE_SIZE)
                                  / static_cast<double> (compressed_size));
    } else
        ret.set_compressed_ratio(0);

    ret.set_max_receive_delay(max_delay);

    ret.set_images_sent(images_sent);
    ret.set_start_time_ms(std::chrono::duration_cast<std::chrono::milliseconds>(start_time.time_since_epoch()).count());
    ret.set_end_time_ms(std::chrono::duration_cast<std::chrono::milliseconds>(end_time.time_since_epoch()).count());
    if (!pedestal_result.empty())
        *ret.mutable_pedestal_result() = {pedestal_result.begin(), pedestal_result.end()};

    ret.set_master_file_name(experiment.GetFilePrefix());
    ret.set_cancelled(cancelled);

    auto tmp = indexing_solution.Mean();
    if (!std::isnan(tmp))
        ret.set_indexing_rate(tmp);

    tmp = bkg_estimate.Mean();
    if (!std::isnan(tmp))
        ret.set_bkg_estimate(tmp);
}

void JFJochReceiver::Cancel() {
    // Remote abort: This tells FPGAs to stop, but doesn't do anything to CPU code
    logger.Warning("Cancelling on request");

    cancelled = true;

    for (int d = 0; d < ndatastreams; d++)
        acquisition_device[d]->ActionAbort();
}

void JFJochReceiver::Cancel(const JFJochException &e) {
    logger.Error("Cancelling data collection due to exception");
    logger.ErrorException(e);
    // Error abort: This tells FPGAs to stop and also prevents deadlock in CPU code, by setting abort to 1
    cancelled = true;

    for (int d = 0; d < ndatastreams; d++)
        acquisition_device[d]->ActionAbort();
}

float JFJochReceiver::GetIndexingRate() const {
    return indexing_solution.Mean();
}

float JFJochReceiver::GetProgress() const {
    if (experiment.GetImageNum() > 0)
        return static_cast<float>(max_image_number_sent) / static_cast<float>(experiment.GetImageNum()) * 100.0;
    else if (experiment.GetFrameNum() > 0)
        // Pedestal
        return static_cast<float>(max_image_number_sent) / static_cast<float>(experiment.GetFrameNum()) * 100.0;
    else
        return 100.0;
}

void JFJochReceiver::FinalizeMeasurement() {
    if (!frame_transformation_futures.empty()) {
        for (auto &future: frame_transformation_futures)
            future.get();

        logger.Info("All processing threads done");
    }

    if (push_images_to_writer) {
        JFJochProtoBuf::ReceiverOutput output;
        GetStatistics(output);

        EndMessage message{};
        message.number_of_images = output.max_image_number_sent();
        message.max_receiver_delay = output.max_receive_delay();
        message.efficiency = output.efficiency();
        message.end_date = time_UTC(std::chrono::system_clock::now());
        message.write_master_file = true;

        if (rad_int_profile)
            message.rad_int_result["dataset"] = rad_int_profile->GetResult(true);
        for (int i = 0; i < rad_int_profile_per_file.size(); i++)
            message.rad_int_result["file" + std::to_string(i)]
                    = rad_int_profile_per_file[i]->GetResult(true);

        image_pusher.EndDataCollection(message);
        logger.Info("Disconnected from writers");
    }

    if (preview_publisher != nullptr)
        preview_publisher->Stop(experiment);

    if (preview_publisher_indexed != nullptr)
        preview_publisher_indexed->Stop(experiment);

    for (int d = 0; d < ndatastreams; d++)
        acquisition_device[d]->ActionAbort();

    end_time = std::chrono::system_clock::now();

    for (auto &future : data_acquisition_futures)
        future.get();

    for (int i = 0; i < send_buffer_count; i++)
        send_buffer_avail.GetBlocking();

    logger.Info("Devices stopped");
    logger.Info("Receiving data done");
}

void JFJochReceiver::SetDataProcessingSettings(const JFJochProtoBuf::DataProcessingSettings &in_data_processing_settings) {
    std::unique_lock<std::mutex> ul(data_processing_settings_mutex);
    DiffractionExperiment::CheckDataProcessingSettings(in_data_processing_settings);
    data_processing_settings = in_data_processing_settings;
}

void JFJochReceiver::StopReceiver() {
    if (measurement.valid()) {
        measurement.get();
        logger.Info("Receiver stopped");
    }
}

JFJochReceiver::~JFJochReceiver() {
    if (measurement.valid())
        measurement.get();
}

JFJochProtoBuf::DataProcessingSettings JFJochReceiver::GetDataProcessingSettings() {
    std::unique_lock<std::mutex> ul(data_processing_settings_mutex);
    return data_processing_settings;
}

JFJochProtoBuf::Plot JFJochReceiver::GetPlots(const JFJochProtoBuf::PlotRequest &request) {
    JFJochProtoBuf::Plot ret;
    auto nbins = experiment.GetSpotFindingBin();
    if (request.binning() > 0)
        nbins = request.binning();

    switch (request.type()) {
        case JFJochProtoBuf::RAD_INT:
            if (rad_int_profile)
                rad_int_profile->GetPlot(ret, request.solid_angle_correction());
            break;
        case JFJochProtoBuf::SPOT_COUNT:
            spot_count.GetPlot(ret, nbins);
            break;
        case JFJochProtoBuf::INDEXING_RATE:
            indexing_solution.GetPlot(ret, nbins);
            break;
        case JFJochProtoBuf::BKG_ESTIMATE:
            bkg_estimate.GetPlot(ret, nbins);
            break;
        case JFJochProtoBuf::INDEXING_RATE_PER_FILE:
            indexing_solution_per_file.GetPlot(ret);
            break;
        default:
            // Do nothing
            break;
    }
    return ret;
}

JFJochProtoBuf::RadialIntegrationProfiles JFJochReceiver::GetRadialIntegrationProfiles() {
    JFJochProtoBuf::RadialIntegrationProfiles ret;

    auto &map = *ret.mutable_plots();

    if (rad_int_profile) {
        rad_int_profile->GetPlot(map["dataset"], true);
        auto &corr = rad_int_profile->GetSolidAngleCorr();
        for (const auto &i: corr)
            ret.add_solid_angle_correction(i);
    }

    for (int i = 0; i < rad_int_profile_per_file.size(); i++)
        rad_int_profile_per_file[i]->GetPlot(map["file" + std::to_string(i)], true);


    return ret;
}

void JFJochReceiver::PrepareConversionOnCPU() {
    if (experiment.GetStorageCellNumber() != 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CPU conversion currently doesn't support storage cells");

    if (!calib.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Calibration not provided");

    fixed_point_conversion.resize(experiment.GetModulesNum());
    for (int i = 0 ; i < experiment.GetModulesNum(); i++)
        fixed_point_conversion[i].Setup(calib->GainCalibration(i),
                                        calib->Pedestal(i, 0, 0),
                                        calib->Pedestal(i, 1, 0),
                                        calib->Pedestal(i, 2, 0),
                                        experiment.GetPhotonEnergy_keV());
}

void JFJochReceiver::PinThreadToDevice(uint16_t data_stream) {
#ifdef JFJOCH_USE_NUMA
    if (numa_available() != -1)
        numa_run_on_node(acquisition_device[data_stream]->GetNUMANode());
#endif
}

void JFJochReceiver::UpdateMaxImage(int64_t image_number) {
    std::unique_lock<std::mutex> ul(max_image_number_sent_mutex);
    if (image_number > max_image_number_sent)
        max_image_number_sent = image_number;
}

void JFJochReceiver::UpdateMaxDelay(int64_t delay) {
    std::unique_lock<std::mutex> ul(max_delay_mutex);
    if (delay > max_delay)
        max_delay = delay;
}

float JFJochReceiver::GetAvailableSendBuffers() const {
    return static_cast<float>(send_buffer_avail.Size()) / static_cast<float>(send_buffer_count);
}