Jungfraujoch/receiver/JFJochReceiverService.cpp

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

#include "JFJochReceiverService.h"
#include "JFJochReceiverRugnux.h"
#include "../common/ThreadAffinity.h"
#include "../rugnux/RugnuxDefaults.h"
#include "JFJochReceiverFPGA.h"
#include "JFJochReceiverLite.h"
#include "../preview/JFJochJPEG.h"
#include "../preview/JFJochTIFF.h"

JFJochReceiverService::JFJochReceiverService(AcquisitionDeviceGroup &in_aq_devices,
                                             Logger &in_logger, ImagePusher &pusher,
                                             size_t send_buffer_size_MiB)
        : aq_devices(in_aq_devices),
          logger(in_logger),
          image_buffer(send_buffer_size_MiB * 1024 * 1024),
          image_pusher(pusher),
          spot_finding_settings(DiffractionExperiment::DefaultDataProcessingSettings()) {
}

JFJochReceiverService &JFJochReceiverService::RugnuxLive(bool enable, const std::string &model_path,
                                                         int reserved_cpus, double deadline_s) {
    std::unique_lock ul_state(state_mutex);
    if (state != ReceiverState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Cannot change rugnux settings during data collection");
    rugnux_enable = enable;
    rugnux_model_path = model_path;
    rugnux_reserved_cpus = reserved_cpus;
    rugnux_deadline_s = deadline_s;
    return *this;
}

RugnuxLiveSnapshot JFJochReceiverService::GetRugnuxStatus() const {
    return rugnux_status.Get();
}

JFJochReceiverService &JFJochReceiverService::NumThreads(int64_t input) {
    if (input <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Thread number must be above zero");
    nthreads = input;
    return *this;
}

void JFJochReceiverService::FinalizeMeasurementChangeState() {
    std::unique_lock ul(state_mutex);
    state = ReceiverState::Idle;
    measurement_done.notify_all();
}

void JFJochReceiverService::FinalizeMeasurement() {
    try {
        receiver->StopReceiver();
    } catch (...) {
        FinalizeMeasurementChangeState();
        throw;
    }
    FinalizeMeasurementChangeState();
}

std::optional<JFJochReceiverStatus> JFJochReceiverService::GetStatus() {
    return receiver_status.GetStatus();
}

void JFJochReceiverService::Preflight(const DiffractionExperiment &experiment) {
    // Same condition JFJochReceiver uses to decide whether it pushes to a writer at all: with no
    // images or no file prefix nothing is written, so there is nothing to check.
    if ((experiment.GetImageNum() <= 0) || experiment.GetFilePrefix().empty())
        return;

    // FillMessage alone, which is what makes this frame small: the per-pixel arrays that dominate a
    // real start message - the pixel mask, the azimuthal-integration map, the ROI map - are added
    // by JFJochReceiver::SendStartMessage and are not here. Measured on a JUNGFRAU 9M: 1.4 kB
    // against 543 kB for the start message that follows. It is deliberately the whole of
    // FillMessage rather than the handful of fields the check reads today, so that the writer sees
    // the same description of the run it will be given at start - a hand-picked subset would
    // answer for a different set of file names the day the check learns to read one more field.
    StartMessage message{};
    experiment.FillMessage(message);
    image_pusher.Preflight(message);
}

void JFJochReceiverService::Start(const DiffractionExperiment &experiment,
                                  const PixelMask &pixel_mask,
                                  const JFCalibration *calibration,
                                  std::shared_ptr<ImagePuller> puller) {
    std::unique_lock ul_state(state_mutex); // unique lock, as it will destroy and create receiver object
    if (state != ReceiverState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState, "Receiver not idle, cannot start");

    try {
        auto nthreads_local = nthreads;
        if (experiment.IsCPUSummation())
            nthreads_local = 4;

        // First clean-up old measurement
        receiver.reset();

        preview_image.Configure(experiment, pixel_mask);

        switch (experiment.GetDetectorType()) {
            case DetectorType::EIGER:
            case DetectorType::JUNGFRAU:
                receiver = std::make_unique<JFJochReceiverFPGA>(experiment, pixel_mask,
                                                                calibration,
                                                                aq_devices, image_pusher,
                                                                logger,
                                                                nthreads_local,
                                                                spot_finding_settings,
                                                                receiver_status,
                                                                plots,
                                                                image_buffer,
                                                                zmq_preview_socket.get(),
                                                                zmq_metadata_socket.get(),
                                                                indexer_thread_pool.get());
                break;
            case DetectorType::DECTRIS:
                if (puller)
                    image_puller = puller;
                else {
                    image_puller = std::make_shared<ZMQImagePuller>(
                            experiment.GetDetectorSetup().GetDECTRISStream2Addr());
                }
                if (rugnux_enable && RugnuxRotationData(experiment)) {
                    receiver = std::make_unique<JFJochReceiverRugnux>(
                            experiment, pixel_mask, *image_puller, image_pusher, logger, nthreads_local,
                            std::max<int>(1, static_cast<int>(AvailableCpus()) - rugnux_reserved_cpus),
                            rugnux_model_path, rugnux_deadline_s, rugnux_status, spot_finding_settings, receiver_status, plots, image_buffer,
                            zmq_preview_socket.get(), zmq_metadata_socket.get());
                    break;
                }
                receiver = std::make_unique<JFJochReceiverLite>(experiment,
                                                                pixel_mask,
                                                                *image_puller,
                                                                image_pusher,
                                                                logger,
                                                                nthreads_local,
                                                                spot_finding_settings,
                                                                receiver_status,
                                                                plots,
                                                                image_buffer,
                                                                zmq_preview_socket.get(),
                                                                zmq_metadata_socket.get(),
                                                                indexer_thread_pool.get());
                break;

        }
        measurement = std::async(std::launch::async, &JFJochReceiverService::FinalizeMeasurement, this);
        state = ReceiverState::Running;
    } catch (const std::exception &e) {
        // The receiver never started, so drop the status its base constructor had already reset -
        // otherwise /status and /statistics keep reporting a zero-progress run that never happened,
        // until the next start overwrites it.
        receiver_status.Clear();
        receiver_status.SetProgress({});
        logger.ErrorException(e);
        throw;
    }
}

void JFJochReceiverService::Cancel(bool silent) {
    std::unique_lock ul(state_mutex);
    if (state == ReceiverState::Running)
        receiver->Cancel(silent);
}

JFJochReceiverOutput JFJochReceiverService::Stop() {
    std::unique_lock ul(state_mutex);

    measurement_done.wait(ul, [this] { return (state != ReceiverState::Running); });

    if (state != ReceiverState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongReceiverState, "Receiver in weird state");

    try {
        if (measurement.valid())
            measurement.get();
    } catch (JFJochException &e) {
        logger.ErrorException(e);
        throw;
    }

    if (!receiver) {
        logger.Warning("Request to stop while receiver not running");
        throw JFJochException(JFJochExceptionCategory::WrongReceiverState, "Receiver idle, cannot stop");
    }
    return receiver->GetFinalStatistics();
}

void JFJochReceiverService::SetSpotFindingSettings(const SpotFindingSettings &settings) {
    try {
        std::unique_lock ul(state_mutex);
        DiffractionExperiment::CheckDataProcessingSettings(settings);
        spot_finding_settings = settings;
        if (state != ReceiverState::Idle)
            receiver->SetSpotFindingSettings(settings);
    } catch (std::exception &e) {
        logger.ErrorException(e);
        throw;
    }
}

MultiLinePlot JFJochReceiverService::GetDataProcessingPlot(const PlotRequest &request) {
    return plots.GetPlots(request);
}

void JFJochReceiverService::GetPlotRaw(std::vector<float> &v, PlotType type, const std::string &roi) {
    plots.GetPlotRaw(v, type, roi);
}


std::vector<AcquisitionDeviceNetConfig> JFJochReceiverService::GetNetworkConfig() {
    return aq_devices.GetNetworkConfig();
}

void JFJochReceiverService::LoadInternalGeneratorImage(const DiffractionExperiment &experiment,
                                                       const std::vector<uint16_t> &image,
                                                       uint64_t image_number) {
    std::vector<uint16_t> raw_geom, eiger_geom;

    const uint16_t *frame;
    if (image.size() == RAW_MODULE_SIZE * experiment.GetModulesNum()) {
        frame = image.data();
    } else if (image.size() == experiment.GetPixelsNum()) {
        raw_geom.resize(RAW_MODULE_SIZE * experiment.GetModulesNum());
        ConvertedToRawGeometry(experiment, raw_geom.data(), image.data());
        frame = raw_geom.data();
    } else
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Size of input array with raw expected image is wrong");

    for (int i = 0; i < experiment.GetDataStreamsNum(); i++) {
        uint32_t module0 = experiment.GetFirstModuleOfDataStream(i);
        switch (experiment.GetDetectorSetup().GetDetectorType()) {
            case DetectorType::EIGER:
                eiger_geom.resize(RAW_MODULE_SIZE);
                for (int m = 0; m < experiment.GetModulesNum(i); m++) {
                    RawToEigerInput(eiger_geom.data(), frame + (module0 + m) * RAW_MODULE_SIZE);
                    aq_devices[i].SetInternalGeneratorFrame(eiger_geom.data(),
                                                            m + experiment.GetModulesNum(i) * image_number);
                }
                break;
            case DetectorType::JUNGFRAU:
                for (int m = 0; m < experiment.GetModulesNum(i); m++)
                    aq_devices[i].SetInternalGeneratorFrame(frame + (module0 + m) * RAW_MODULE_SIZE,
                                                            m + experiment.GetModulesNum(i) * image_number);
                break;
            default:
                throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Detector not supported");
        }
    }
}

void JFJochReceiverService::GetXFELEventCode(std::vector<uint64_t> &v) const {
    plots.GetXFELEventCode(v);
}

void JFJochReceiverService::GetXFELPulseID(std::vector<uint64_t> &v) const {
    plots.GetXFELPulseID(v);
}

std::vector<DeviceStatus> JFJochReceiverService::GetDeviceStatus() const {
    return aq_devices.GetDeviceStatus();
}

std::optional<float> JFJochReceiverService::GetProgress() const {
    return receiver_status.GetProgress();
}

JFJochReceiverService &JFJochReceiverService::PreviewSocket(const std::string &addr, const std::optional<int32_t> &watermark) {
    if (!addr.empty()) {
        logger.Info("ZeroMQ preview socket available at {}", addr);
        zmq_preview_socket = std::make_unique<ZMQPreviewSocket>(addr, watermark);
    }
    return *this;
}

JFJochReceiverService &JFJochReceiverService::MetadataSocket(const std::string &addr) {
    if (!addr.empty()) {
        logger.Info("ZeroMQ metadata socket available at {}", addr);
        zmq_metadata_socket = std::make_unique<ZMQMetadataSocket>(addr);
    }
    return *this;
}

std::string JFJochReceiverService::GetPreviewSocketAddress() const {
    if (zmq_preview_socket)
        return zmq_preview_socket->GetAddress();
    return "";
}

std::string JFJochReceiverService::GetMetadataSocketAddress() const {
    if (zmq_metadata_socket)
        return zmq_metadata_socket->GetAddress();
    return "";
}

JFJochReceiverService &JFJochReceiverService::PreviewSocketSettings(const ZMQPreviewSettings &input) {
    if (zmq_preview_socket)
        zmq_preview_socket->ImportSettings(input);
    return *this;
}

JFJochReceiverService &JFJochReceiverService::MetadataSocketSettings(const ZMQMetadataSettings &input) {
    if (zmq_metadata_socket)
        zmq_metadata_socket->ImportSettings(input);
    return *this;
}

ZMQPreviewSettings JFJochReceiverService::GetPreviewSocketSettings() const {
    if (zmq_preview_socket)
        return zmq_preview_socket->GetSettings();
    return {};
}

ZMQMetadataSettings JFJochReceiverService::GetMetadataSocketSettings() const {
    if (zmq_metadata_socket)
        return zmq_metadata_socket->GetSettings();
    return {};
}

void JFJochReceiverService::GetStartMessageFromBuffer(std::vector<uint8_t> &v) {
    image_buffer.GetStartMessage(v);
}

bool JFJochReceiverService::GetImageFromBuffer(std::vector<uint8_t> &v, int64_t image_number) {
    return image_buffer.GetImage(v, image_number);
}

std::string JFJochReceiverService::GetJPEGFromBuffer(const PreviewImageSettings &settings, int64_t image_number) {
    std::vector<uint8_t> cbor_image;
    if (!image_buffer.GetImage(cbor_image, image_number))
        return {};
    return preview_image.GenerateImage(settings, cbor_image);
}

std::string JFJochReceiverService::GetTIFFFromBuffer(int64_t image_number) {
    std::vector<uint8_t> cbor_image;
    if (!image_buffer.GetImage(cbor_image, image_number))
        return {};
    return PreviewImage::GenerateTIFF(cbor_image);
}

ImageBufferStatus JFJochReceiverService::GetImageBufferStatus() const {
    return image_buffer.GetStatus();
}

void JFJochReceiverService::ClearResults() {
    std::unique_lock ul(state_mutex);
    if (state != ReceiverState::Idle)
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Cannot clear results during data collection");
    receiver_status.Clear();
    receiver_status.SetProgress({});
    plots.Clear();
    image_buffer.Finalize(std::chrono::milliseconds(2500));
}

void JFJochReceiverService::ClearImageBuffer() {
    std::unique_lock ul(state_mutex);
    // Clearing image buffer during data collection could be catastrophic, so better protect here, even if redundant
    // with JFJochStateMachine
    if (state == ReceiverState::Idle)
        image_buffer.Finalize(std::chrono::milliseconds(2500));
    else
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Cannot clear image buffer during data collection");
}

JFJochReceiverService &JFJochReceiverService::Indexing(const IndexingSettings &input) {
    std::unique_lock ul(state_mutex);
    // Clearing image buffer during data collection could be catastrophic, so better protect here, even if redundant
    // with JFJochStateMachine
    if (state == ReceiverState::Idle) {
        // Release the previous run's receiver first. It holds a raw pointer to the indexer pool and
        // its own GPU resources; keeping it alive while we rebuild the pool means a fresh GPU indexer
        // has to coexist with a stale receiver (e.g. after a failed acquisition start, where the
        // receiver is stopped but not destroyed until the next Start), which can make the GPU indexer
        // initialisation fail. Destroying the receiver before the pool also avoids the dangling
        // pointer. Safe here: state is Idle, so no measurement is using it.
        receiver.reset();

        logger.Info("Resetting indexing thread pool");
        indexer_thread_pool.reset();

        if (input.GetAlgorithm() != IndexingAlgorithmEnum::None) {
            logger.Info("Creating indexing thread pool...");
            indexer_thread_pool = std::make_unique<IndexerThreadPool>(input);
            logger.Info("   ... done");
        }
        return *this;
    } else
        throw JFJochException(JFJochExceptionCategory::WrongDAQState,
                              "Cannot change indexing settings during data collection");
}


ImagePusherStatus JFJochReceiverService::GetImagePusherStatus() const {
    return image_pusher.GetStatus();
}