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

#include "JFJochReceiver.h"

#include <thread>

#include "../jungfrau/JFPedestalCalc.h"
#include "../image_analysis/MXAnalyzer.h"
#include "../common/DiffractionGeometry.h"
#include "ImageMetadata.h"
#include "../common/time_utc.h"
#include "../common/PixelMask.h"
#include "../common/CUDAWrapper.h"
#include "../common/ModuleSummation.h"

JFJochReceiver::JFJochReceiver(const DiffractionExperiment &in_experiment,
                               const PixelMask &in_pixel_mask,
                               const JFCalibration *in_calibration,
                               AcquisitionDeviceGroup &in_aq_device,
                               ImagePusher &in_image_sender,
                               Logger &in_logger, int64_t in_forward_and_sum_nthreads,
                               const NUMAHWPolicy &in_numa_policy,
                               const SpotFindingSettings &in_spot_finding_settings,
                               PreviewImage &in_preview_image,
                               PreviewImage &in_preview_image_indexed,
                               JFJochReceiverCurrentStatus &in_current_status,
                               JFJochReceiverPlots &in_plots,
                               SendBuffer &buf,
                               ZMQPreviewSocket *in_zmq_preview_socket,
                               ZMQMetadataSocket *in_zmq_metadata_socket)
    : experiment(in_experiment),
      calibration(nullptr),
      pixel_mask(in_pixel_mask),
      send_buf_ctrl(experiment, buf),
      acquisition_device(in_aq_device),
      logger(in_logger),
      image_pusher(in_image_sender),
      pedestal_nthreads((experiment.GetStorageCellNumber() > 2) ? 1 : 4),
      summation_nthreads(2),
      frame_transformation_nthreads(
          in_forward_and_sum_nthreads),
      ndatastreams(experiment.GetDataStreamsNum()),
      data_acquisition_ready(ndatastreams),
      frame_transformation_ready((experiment.GetImageNum() > 0) ? frame_transformation_nthreads : 0),
      numa_policy(in_numa_policy),
      adu_histogram_module(experiment.GetModulesNum()),
      az_int_mapping(experiment),
      spot_finding_settings(in_spot_finding_settings),
      preview_image(in_preview_image),
      preview_image_indexed(in_preview_image_indexed),
      current_status(in_current_status),
      plots(in_plots),
      serialmx_filter(experiment),
      user_mask_raw_coord(in_pixel_mask.GetUserMask(experiment, false)),
      zmq_preview_socket(in_zmq_preview_socket),
      zmq_metadata_socket(in_zmq_metadata_socket) {
    current_status.SetProgress(0);
    current_status.SetStatus(GetStatus());
    current_status.SetEfficiency({});

    plots.Setup(experiment, az_int_mapping);

    if (experiment.GetDetectorSetup().GetDetectorType() == DetectorType::JUNGFRAU)
        calibration = in_calibration;

    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");

    logger.Info("NUMA policy: {}", numa_policy.GetName());

    if ((get_gpu_count() == 0) && (experiment.GetUnitCell().has_value()))
        logger.Warning("GPU absent, indexing is not possible");

    for (int d = 0; d < ndatastreams; d++) {
        acquisition_device[d].PrepareAction(experiment);
        acquisition_device[d].SetSpotFinderParameters(spot_finding_settings);
        logger.Debug("Acquisition device {} prepared", 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");
        }

        switch (experiment.GetDetectorMode()) {
            case DetectorMode::PedestalG1:
                only_2nd_sc_pedestal = !experiment.IsFixedGainG1() && (experiment.GetStorageCellNumber() == 2);
                break;
            case DetectorMode::PedestalG2:
                only_2nd_sc_pedestal = (experiment.GetStorageCellNumber() == 2);
                break;
            default:
                only_2nd_sc_pedestal = false;
                break;
        }

        int64_t pedestal_count = (only_2nd_sc_pedestal)
                                     ? experiment.GetModulesNum()
                                     : experiment.GetModulesNum() * experiment.GetStorageCellNumber();
        for (int i = 0; i < pedestal_count; i++)
            pedestal.emplace_back(std::make_unique<JFPedestalCalc>(experiment));

        for (int s = 0; s < experiment.GetStorageCellNumber(); s++) {
            bool ignore = only_2nd_sc_pedestal && (s == 0);
            for (int d = 0; d < ndatastreams; d++)
                for (int m = 0; m < experiment.GetModulesNum(d); m++)
                    for (int n = 0; n < pedestal_nthreads; n++)
                        frame_transformation_futures.emplace_back(std::async(std::launch::async,
                                                                             &JFJochReceiver::MeasurePedestalThread,
                                                                             this, d, m, s, n, ignore));
        }
        logger.Info("Pedestal threads ready ({} threads/module*SC)", pedestal_nthreads);
    } else if (experiment.GetImageNum() > 0) {
        SendStartMessage(pixel_mask);
        SendCalibration();

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

        // Setup frames summation and forwarding
        for (uint32_t i = 0; i < frame_transformation_nthreads; i++) {
            auto handle = std::async(std::launch::async, &JFJochReceiver::FrameTransformationThread,
                                     this, i);
            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");
    }

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

    data_acquisition_ready.wait();
    start_time = std::chrono::system_clock::now();
    measurement = std::async(std::launch::async, &JFJochReceiver::FinalizeMeasurement, this);
    logger.Info("Receiving data started");
}

void JFJochReceiver::SendStartMessage(const PixelMask &pixel_mask) {
    StartMessage message{};
    experiment.FillMessage(message);
    message.arm_date = time_UTC(std::chrono::system_clock::now());
    message.az_int_bin_number = az_int_mapping.GetBinNumber();
    message.az_int_bin_to_q = az_int_mapping.GetBinToQ();
    message.writer_notification_zmq_addr = image_pusher.GetWriterNotificationSocketAddress();

    std::vector<uint32_t> nexus_mask;
    nexus_mask = pixel_mask.GetMask(experiment);

    size_t xpixel = experiment.GetXPixelsNum();
    size_t ypixel = experiment.GetYPixelsNum();

    CompressedImage image{
        .data = (uint8_t *) nexus_mask.data(),
        .size = nexus_mask.size() * sizeof(nexus_mask[0]),
        .xpixel = xpixel,
        .ypixel = ypixel,
        .pixel_depth_bytes = sizeof(nexus_mask[0]),
        .pixel_is_signed = false,
        .pixel_is_float = false,
        .algorithm = CompressionAlgorithm::NO_COMPRESSION,
        .channel = "default"
    };
    message.AddPixelMask(image);

    if (push_images_to_writer)
        image_pusher.StartDataCollection(message);

    if (zmq_preview_socket != nullptr)
        zmq_preview_socket->StartDataCollection(message);

    if (zmq_metadata_socket != nullptr)
        zmq_metadata_socket->StartDataCollection(message);
}

void JFJochReceiver::SendPedestal(const std::string &prefix, const std::vector<uint8_t> &v, int gain, int sc) {
    size_t xpixel = RAW_MODULE_COLS;
    size_t ypixel = experiment.GetModulesNum() * RAW_MODULE_LINES;

    std::string channel;

    if (experiment.GetStorageCellNumber() > 1)
        channel = fmt::format("{:s}_g{:d}_sc{:d}", prefix, gain, sc);
    else
        channel = fmt::format("{:s}_g{:d}", prefix, gain);

    CompressedImage image{
        .data = v.data(),
        .size = v.size(),
        .xpixel = (size_t) xpixel,
        .ypixel = (size_t) ypixel,
        .pixel_depth_bytes = 2,
        .pixel_is_signed = false,
        .pixel_is_float = false,
        .algorithm = CompressionAlgorithm::BSHUF_LZ4,
        .channel = channel
    };

    image_pusher.SendCalibration(image);
}

void JFJochReceiver::SendCalibration() {
    if ((calibration == nullptr) || !experiment.GetSaveCalibration() || !push_images_to_writer)
        return;

    JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4);

    for (int sc = 0; sc < experiment.GetStorageCellNumber(); sc++) {
        for (int gain = 0; gain < 3; gain++) {
            if (experiment.IsFixedGainG1() && (gain != 1))
                continue;
            SendPedestal("pedestal", compressor.Compress(calibration->GetPedestal(gain, sc)), gain, sc);
            SendPedestal("pedestal_rms", compressor.Compress(calibration->GetPedestalRMS(gain, sc)), gain, sc);
        }
    }
}

void JFJochReceiver::AcquireThread(uint16_t data_stream) {
    try {
        NUMAHWPolicy::RunOnNode(acquisition_device[data_stream].GetNUMANode());
    } catch (const JFJochException &e) {
        logger.Warning("NUMA bind error {} for device thread {} - continuing without binding", e.what(), data_stream);
    }

    try {
        LoadCalibrationToFPGA(data_stream);
        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();
        logger.ErrorException(e);
        logger.Warning("Device thread {} done due to an error", data_stream);
        return;
    }

    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) {
        logger.ErrorException(e);
        Cancel(e);
        logger.ErrorException(e);
        logger.Warning("Device thread {} done due to an error", data_stream);
        return;
    }
    logger.Info("Device thread {} done", data_stream);
}


void JFJochReceiver::MeasurePedestalThread(uint16_t data_stream, uint16_t module_number, uint16_t storage_cell,
                                           uint32_t threadid, bool ignore) {
    try {
        NUMAHWPolicy::RunOnNode(acquisition_device[data_stream].GetNUMANode());
    } catch (const JFJochException &e) {
        logger.Error("HW bind error {}", e.what());
    }

    JFPedestalCalc pedestal_calc(experiment);

    uint64_t starting_frame = storage_cell + threadid * experiment.GetStorageCellNumber();
    uint64_t frame_stride = experiment.GetStorageCellNumber() * pedestal_nthreads;
    uint32_t storage_cell_header = UINT32_MAX;

    try {
        for (size_t frame = starting_frame; frame < experiment.GetFrameNum(); frame += frame_stride) {
            // Frame will be processed only if one already collects frame+2
            acquisition_device[data_stream].Counters().WaitForFrame(frame + 2, module_number);
            if (acquisition_device[data_stream].Counters().IsFullModuleCollected(frame, module_number) && !ignore) {
                auto output = acquisition_device[data_stream].GetDeviceOutput(frame, module_number);
                // Partial packets will bring more problems, than benefit
                pedestal_calc.AnalyzeImage((uint16_t *) output->pixels);
                storage_cell_header = (output->module_statistics.debug >> 8) & 0xF;
            }
            acquisition_device[data_stream].FrameBufferRelease(frame, module_number);

            UpdateMaxDelay(acquisition_device[data_stream].Counters().CalculateDelay(frame, module_number));
            current_status.SetProgress(GetProgress());
            current_status.SetStatus(GetStatus());
        }

        uint64_t offset = experiment.GetFirstModuleOfDataStream(data_stream) + module_number;
        if (!only_2nd_sc_pedestal)
            offset += experiment.GetModulesNum() * storage_cell;

        if (!ignore)
            *pedestal[offset] += pedestal_calc;
    } catch (const JFJochException &e) {
        Cancel(e);
    }
    logger.Debug("Pedestal calculation thread for data stream {} module {} storage cell {} -> header {} done",
                 data_stream, module_number, storage_cell, storage_cell_header);
}

int64_t JFJochReceiver::SummationThread(uint16_t data_stream,
                                        int64_t image_number,
                                        uint16_t module_number,
                                        uint32_t threadid,
                                        ModuleSummation &summation) {
    try {
        NUMAHWPolicy::RunOnNode(acquisition_device[data_stream].GetNUMANode());
    } catch (const JFJochException &e) {
        logger.Error("HW bind error {}", e.what());
    }

    ModuleSummation local_summation(experiment);
    int64_t starting_frame = image_number * experiment.GetSummation();

    for (int64_t i = threadid; i < experiment.GetSummation(); i += summation_nthreads) {
        const int64_t frame = starting_frame + i;

        // Frame will be processed only if one already collects frame+2
        acquisition_device[data_stream].Counters().WaitForFrame(frame + 2, module_number);
        if (acquisition_device[data_stream].Counters().IsAnyPacketCollected(frame, module_number)) {
            const auto output = acquisition_device[data_stream].GetDeviceOutput(frame, module_number);
            local_summation.AddFPGAOutput(*output);
        } else
            local_summation.AddEmptyOutput();

        acquisition_device[data_stream].FrameBufferRelease(frame, module_number);

        UpdateMaxDelay(acquisition_device[data_stream].Counters().CalculateDelay(frame, module_number));
        current_status.SetProgress(GetProgress());
        current_status.SetStatus(GetStatus());
    }

    if (!summation.empty())
        summation.AddFPGAOutput(local_summation.GetOutput(), 4);
    return 0;
}

void JFJochReceiver::FrameTransformationThread(uint32_t threadid) {
    std::unique_ptr<MXAnalyzer> analyzer;

    try {
        numa_policy.Bind(threadid);
        analyzer = std::make_unique<MXAnalyzer>(experiment);
    } catch (const JFJochException &e) {
        frame_transformation_ready.count_down();
        logger.Error("Thread setup error {}", e.what());
        Cancel(e);
        return;
    }

    FrameTransformation transformation(experiment);

    frame_transformation_ready.count_down();

    uint16_t az_int_min_bin = std::floor(az_int_mapping.QToBin(experiment.GetLowQForBkgEstimate_recipA()));
    uint16_t az_int_max_bin = std::ceil(az_int_mapping.QToBin(experiment.GetHighQForBkgEstimate_recipA()));

    uint64_t image_number;
    while (images_to_go.Get(image_number) != 0) {
        try {
            int64_t expected_frame = image_number;
            if (experiment.IsCPUSummation())
                expected_frame *= experiment.GetSummation();

            logger.Debug("Frame transformation thread - trying to get image {}", expected_frame);
            // If data acquisition is finished and fastest frame for the first device is behind
            acquisition_device[0].Counters().WaitForFrame(expected_frame);
            logger.Debug("Frame transformation thread - frame arrived {}", expected_frame);
            if (acquisition_device[0].Counters().IsAcquisitionFinished() &&
                (acquisition_device[0].Counters().GetFastestFrameNumber() < expected_frame)) {
                logger.Debug("Frame transformation thread - skipping image {}", expected_frame);
                continue;
            }

            DataMessage message{};
            message.number = image_number;
            message.original_number = image_number;
            message.user_data = experiment.GetImageAppendix();
            message.run_number = experiment.GetRunNumber();
            message.run_name = experiment.GetRunName();

            ImageMetadata metadata(experiment);

            AzimuthalIntegrationProfile az_int_profile_image(az_int_mapping);

            auto local_spot_finding_settings = GetSpotFindingSettings();

            if (experiment.IsCPUSummation()) {
                std::vector<std::unique_ptr<ModuleSummation>> summation;
                for (int i = 0; i < experiment.GetModulesNum(); i++)
                    summation.emplace_back(std::make_unique<ModuleSummation>(experiment));

                std::vector<std::future<int64_t> > futures;

                for (int d = 0; d < ndatastreams; d++) {
                    for (int m = 0; m < experiment.GetModulesNum(d); m++) {
                        size_t module_abs_number = experiment.GetFirstModuleOfDataStream(d) + m;
                        for (int i = 0; i < summation_nthreads; i++) {
                            futures.emplace_back(
                                std::async(std::launch::async,
                                           &JFJochReceiver::SummationThread,
                                           this,
                                           d, image_number, m, i,
                                           std::ref(*summation[module_abs_number]))
                            );
                        }
                    }
                }

                for (auto &f: futures)
                    f.get();

                for (int d = 0; d < ndatastreams; d++) {
                    for (int m = 0; m < experiment.GetModulesNum(d); m++) {
                        size_t i = experiment.GetFirstModuleOfDataStream(d) + m;

                        if (!summation[i]->empty()) {
                            adu_histogram_module[i].Add(summation[i]->GetOutput());
                            transformation.ProcessModule(&summation[i]->GetOutput(), d);
                            metadata.Process(&summation[i]->GetOutput());
                            az_int_profile_image.Add(summation[i]->GetOutput());
                            analyzer->ReadFromCPU(&summation[i]->GetOutput(), GetSpotFindingSettings(), i);
                        } else
                            transformation.FillNotCollectedModule(m, d);
                    }
                }
            } else {
                logger.Debug("Frame transformation thread - processing image from FPGA {}", image_number);
                for (int d = 0; d < ndatastreams; d++) {
                    for (int m = 0; m < experiment.GetModulesNum(d); m++) {
                        acquisition_device[d].Counters().WaitForFrame(image_number + 2, m);
                        if (acquisition_device[d].Counters().IsAnyPacketCollected(image_number, m)) {
                            const DeviceOutput *output = acquisition_device[d].GetDeviceOutput(image_number, m);

                            metadata.Process(output);

                            size_t module_abs_number = experiment.GetFirstModuleOfDataStream(d) + m;
                            adu_histogram_module[module_abs_number].Add(*output);
                            az_int_profile_image.Add(*output);

                            analyzer->ReadFromFPGA(output, local_spot_finding_settings, module_abs_number);

                            transformation.ProcessModule(output, d);
                        } else
                            transformation.FillNotCollectedModule(m, d);
                        acquisition_device[d].FrameBufferRelease(image_number, m);
                    }
                    auto delay = acquisition_device[d].Counters().CalculateDelay(image_number);
                    UpdateMaxDelay(delay);
                    if (delay > message.receiver_aq_dev_delay)
                        message.receiver_aq_dev_delay = delay;
                }
            }

            metadata.Export(message, 256 * experiment.GetModulesNum() * experiment.GetSummation());
            if (message.image_collection_efficiency == 0.0f) {
                plots.AddEmptyImage(message);
                continue;
            }

            analyzer->Process(message, local_spot_finding_settings);

            message.receiver_free_send_buf = send_buf_ctrl.GetAvailBufLocations();
            message.az_int_profile = az_int_profile_image.GetResult();
            message.bkg_estimate = az_int_profile_image.GetMeanValueOfBins(az_int_min_bin, az_int_max_bin);

            plots.Add(message, az_int_profile_image);

            preview_image.UpdateImage(transformation.GetImage(), message.spots);
            if (message.indexing_result)
                preview_image_indexed.UpdateImage(transformation.GetImage(), message.spots);

            images_collected++;

            if (!serialmx_filter.ApplyFilter(message))
                images_skipped++;
            else {
                auto loc = send_buf_ctrl.GetBufLocation();
                if (loc == nullptr) // No free buffer locations - continue
                    continue;

                auto writer_buffer = (uint8_t *) loc->get_ptr();

                CBORStream2Serializer serializer(writer_buffer, experiment.GetSendBufferLocationSize());
                PrepareCBORImage(message, experiment, nullptr, 0);
                serializer.SerializeImage(message);
                if (experiment.GetSendBufferLocationSize() - serializer.GetImageAppendOffset()
                    < experiment.GetMaxCompressedSize())
                    throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                                          "Not enough memory to save image");
                size_t image_size = transformation.CompressImage(writer_buffer + serializer.GetImageAppendOffset());
                serializer.AppendImage(image_size);
                compressed_size += image_size;

                if (zmq_preview_socket != nullptr)
                    zmq_preview_socket->SendImage(writer_buffer, serializer.GetBufferSize());

                if (zmq_metadata_socket != nullptr)
                    zmq_metadata_socket->AddDataMessage(message);

                if (push_images_to_writer) {
                    image_pusher.SendImage(writer_buffer, serializer.GetBufferSize(), message.number, loc);
                    images_sent++; // Handle case when image not sent properly
                } else
                    loc->release();
                UpdateMaxImage(message.number);
            }

            logger.Debug("Frame transformation thread - done sending image {} / {}", image_number, message.number);

            current_status.SetProgress(GetProgress());
            current_status.SetStatus(GetStatus());
        } catch (const JFJochException &e) {
            logger.ErrorException(e);
            Cancel(e);
        }
    }
    logger.Debug("Sum&compression thread done");
}

float JFJochReceiver::GetEfficiency() const {
    uint64_t expected_packets = 0;
    uint64_t received_packets = 0;

    for (int d = 0; d < ndatastreams; d++) {
        expected_packets += acquisition_device[d].Counters().GetExpectedPackets();
        received_packets += acquisition_device[d].Counters().GetTotalPackets();
    }

    if ((expected_packets == received_packets) || (expected_packets == 0))
        return 1.0;
    else
        return received_packets / static_cast<double>(expected_packets);
}

void JFJochReceiver::Cancel(bool silent) {
    if (!silent) {
        // 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].Cancel();
}

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].Cancel();
}

float JFJochReceiver::GetProgress() const {
    int64_t frames = experiment.GetImageNum();
    if (experiment.IsCPUSummation())
        frames *= experiment.GetSummation();

    if (frames == 0)
        frames = experiment.GetFrameNum();
    if ((frames == 0) || (acquisition_device[0].Counters().IsAcquisitionFinished()))
        return 1.0;

    return static_cast<float>(acquisition_device[0].Counters().GetSlowestFrameNumber()) / static_cast<float>(frames);
}

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

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

    current_status.SetProgress(1.0);
    current_status.SetStatus(GetStatus());

    EndMessage message{};
    message.max_image_number = max_image_number_sent;
    message.images_collected_count = images_collected;
    message.images_sent_to_write_count = images_sent;

    message.max_receiver_delay = max_delay;
    message.efficiency = GetEfficiency();
    message.end_date = time_UTC(std::chrono::system_clock::now());
    message.run_number = experiment.GetRunNumber();
    message.run_name = experiment.GetRunName();

    message.bkg_estimate = plots.GetBkgEstimate();
    message.indexing_rate = plots.GetIndexingRate();

    message.az_int_result["dataset"] = plots.GetAzIntProfile();

    for (int i = 0; i < adu_histogram_module.size(); i++)
        message.adu_histogram["module" + std::to_string(i)] = adu_histogram_module[i].GetHistogram();

    if (push_images_to_writer) {
        if (!image_pusher.EndDataCollection(message))
            logger.Error("End message not sent via ZeroMQ (time-out)");

        logger.Info("Disconnected from writers");
    }

    if (zmq_metadata_socket != nullptr)
        zmq_metadata_socket->EndDataCollection(message);

    if (zmq_preview_socket != nullptr)
        zmq_preview_socket->EndDataCollection(message);

    if (experiment.GetImageNum() > 0) {
        for (int d = 0; d < ndatastreams; d++)
            acquisition_device[d].Cancel();
    }

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

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

    logger.Info("Devices stopped");

    if (!send_buf_ctrl.CheckIfBufferReturned(std::chrono::seconds(10))) {
        logger.Error("Send commands not finalized in 10 seconds");
        throw JFJochException(JFJochExceptionCategory::ZeroMQ, "Send commands not finalized in 10 seconds");
    }
    logger.Info("Receiving data done");

    if (push_images_to_writer)
        image_pusher.Finalize();

    current_status.SetProgress({});
    current_status.SetStatus(GetStatus());

    logger.Info("Writing process finalized");
}

void JFJochReceiver::SetSpotFindingSettings(const SpotFindingSettings &in_spot_finding_settings) {
    std::unique_lock ul(spot_finding_settings_mutex);
    DiffractionExperiment::CheckDataProcessingSettings(in_spot_finding_settings);
    spot_finding_settings = in_spot_finding_settings;
    for (int i = 0; i < ndatastreams; i++)
        acquisition_device[i].SetSpotFinderParameters(spot_finding_settings);
}

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

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

SpotFindingSettings JFJochReceiver::GetSpotFindingSettings() {
    std::unique_lock ul(spot_finding_settings_mutex);
    return spot_finding_settings;
}

void JFJochReceiver::UpdateMaxImage(uint64_t image_number) {
    std::unique_lock ul(max_image_number_sent_mutex);
    if (image_number + 1 > max_image_number_sent)
        max_image_number_sent = image_number + 1;
}

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


JFJochReceiverOutput JFJochReceiver::GetFinalStatistics() const {
    JFJochReceiverOutput ret;

    for (int d = 0; d < ndatastreams; d++) {
        for (int m = 0; m < acquisition_device[d].Counters().GetModuleNumber(); m++) {
            ret.expected_packets.push_back(acquisition_device[d].Counters().GetExpectedPacketsPerModule());
            ret.received_packets.push_back(acquisition_device[d].Counters().GetTotalPackets(m));
        }
    }
    ret.efficiency = GetEfficiency();

    ret.start_time_ms = std::chrono::duration_cast<std::chrono::milliseconds>(start_time.time_since_epoch()).count();
    ret.end_time_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end_time.time_since_epoch()).count();
    RetrievePedestal(ret.pedestal_result);

    ret.status = GetStatus();
    return ret;
}

JFJochReceiverStatus JFJochReceiver::GetStatus() const {
    JFJochReceiverStatus ret;

    ret.indexing_rate = plots.GetIndexingRate();
    ret.bkg_estimate = plots.GetBkgEstimate();

    if ((experiment.GetImageNum() > 0) && (compressed_size > 0)) {
        ret.compressed_ratio = static_cast<double>((images_sent + images_skipped)
                                                   * experiment.GetByteDepthImage()
                                                   * experiment.GetModulesNum()
                                                   * RAW_MODULE_SIZE)
                               / static_cast<double>(compressed_size);
    }

    ret.compressed_size = compressed_size;
    ret.max_receive_delay = max_delay;
    ret.max_image_number_sent = max_image_number_sent;
    ret.images_collected = images_collected;
    ret.images_sent = images_sent;
    ret.images_skipped = images_skipped;
    ret.cancelled = cancelled;
    ret.efficiency = GetEfficiency();

    return ret;
}

void JFJochReceiver::RetrievePedestal(std::vector<JFModulePedestal> &output) const {
    time_t curr_time = std::chrono::system_clock::to_time_t(start_time);
    for (const auto &pc: pedestal) {
        JFModulePedestal mp;
        if (experiment.GetDetectorMode() == DetectorMode::PedestalG0)
            pc->Export(mp, PEDESTAL_G0_WRONG_GAIN_ALLOWED_COUNT);
        else
            pc->Export(mp);
        mp.SetCollectionTime(curr_time);
        output.emplace_back(std::move(mp));
    }
}

void JFJochReceiver::LoadCalibrationToFPGA(uint16_t data_stream) {
    if (experiment.IsPedestalRun())
        return; // No calibration loaded for pedestal

    if (calibration != nullptr)
        acquisition_device[data_stream].InitializeCalibration(experiment, *calibration);

    // Initialize pixel_mask
    acquisition_device[data_stream].InitializePixelMask(experiment, pixel_mask);

    // Initialize roi_map
    acquisition_device[data_stream].InitializeROIMap(experiment);

    size_t m0 = experiment.GetFirstModuleOfDataStream(data_stream);
    std::vector<float> tmp(RAW_MODULE_SIZE);
    for (int m = 0; m < experiment.GetModulesNum(data_stream); m++) {
        // Mask pixels for spot finding based on resolution and user mask
        CalcSpotFinderResolutionMap(tmp.data(), experiment, m0 + m);
        for (int i = 0; i < RAW_MODULE_SIZE; i++) {
            if (user_mask_raw_coord[(m0 + m) * RAW_MODULE_SIZE + i] != 0)
                tmp[i] = 0.0f;
        }
        acquisition_device[data_stream].InitializeSpotFinderResolutionMap(tmp.data(), m);

        CalcAzIntCorrRawCoord(tmp.data(), experiment, m0 + m);
        acquisition_device[data_stream].InitializeIntegrationMap(az_int_mapping.GetPixelToBinMappingRaw().data()
                                                                 + (m0 + m) * RAW_MODULE_SIZE, tmp.data(), m);
    }
}