Jungfraujoch/broker/JFJochStateMachine.h

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

#pragma once

#include <string>
#include <mutex>
#include <future>
#include <optional>
#include <exception>

#include "../common/DiffractionExperiment.h"
#include "../jungfrau/JFCalibration.h"
#include "../common/Logger.h"

#include "JFJochServices.h"
#include "../common/ROIMap.h"
#include "../common/BrokerStatus.h"
#include "BearerTokens.h"

struct DetectorListElement {
    std::string description;
    std::string serial_number;
    std::string base_ipv4_addr;
    int64_t udp_interface_count;
    int64_t nmodules;
    int64_t width;
    int64_t height;
    std::chrono::nanoseconds readout_time;
    std::chrono::nanoseconds min_frame_time;
    std::chrono::nanoseconds min_count_time;
    DetectorType detector_type;
    float pixel_size_mm;
};

struct DetectorList {
    std::vector<DetectorListElement> detector;
    int64_t current_id;
};

struct MeasurementStatistics {
    std::string file_prefix;
    std::string experiment_group;
    int64_t run_number;

    int64_t images_expected;
    int64_t images_collected;
    int64_t images_sent;
    int64_t images_skipped;
    std::optional<int64_t> images_written;

    int64_t max_image_number_sent;
    std::optional<float> collection_efficiency;
    std::optional<float> compression_ratio;

    bool cancelled;
    std::optional<int64_t> max_receive_delay;

    std::optional<float> indexing_rate;

    int64_t detector_width;
    int64_t detector_height;
    int64_t detector_pixel_depth;

    std::optional<float> bkg_estimate;
    std::optional<std::pair<float, float>> beam_center_drift_pxl;

    std::string unit_cell;

    std::optional<float> error_pixels;
    std::optional<float> saturated_pixels;
    std::optional<float> roi_beam_npixel;
    std::optional<float> roi_beam_sum;
};

class JFJochStateMachine {
    Logger &logger;
    JFJochServices &services;

    std::future<void> measurement;

    // assuming immutable during normal operation
    std::vector<DetectorSetup> detector_setup;
    std::vector<JFModuleGainCalibration> gain_calibration;

    mutable std::mutex experiment_detector_settings_mutex;
    mutable std::mutex experiment_azimuthal_integration_settings_mutex;
    mutable std::mutex experiment_instrument_metadata_mutex;
    mutable std::mutex experiment_image_format_settings_mutex;
    mutable std::mutex experiment_file_writer_settings_mutex;
    mutable std::mutex experiment_indexing_settings_mutex;
    mutable std::mutex experiment_dark_mask_settings_mutex;
    mutable std::mutex rugnux_mutex;
    DiffractionExperiment experiment;

    // mutex m is protecting:
    mutable std::mutex m;
    std::condition_variable c;
    std::atomic<JFJochState> state = JFJochState::Inactive; // state should not be set directly, but through SetState function
    std::atomic<bool> cancel_sequence = false;
    std::unique_ptr<JFCalibration> calibration;
    PixelMask pixel_mask;
    int64_t current_detector_setup; // Lock only on change
    std::optional<ScanResult> scan_result;
    BearerTokens tokens;
    // Set by MeasurementThread when a Start fails. A synchronous Start() rethrows it directly; an
    // asynchronous one has already returned, so the wait functions rethrow it instead. Every entry
    // point that begins new work clears it, so it is reported to every caller asking in between but
    // never attributed to the operation after it.
    std::exception_ptr start_exception;

    mutable std::mutex calibration_statistics_mutex;
    std::vector<JFCalibrationModuleStatistics> calibration_statistics;

    mutable std::mutex data_processing_settings_mutex;
    SpotFindingSettings data_processing_settings;

    mutable std::mutex pixel_mask_statistics_mutex;
    PixelMaskStatistics pixel_mask_statistics;

    mutable std::mutex broker_status_mutex;
    BrokerStatus broker_status;

    mutable std::mutex roi_mutex;
    ROIDefinition roi;

    const int32_t gpu_count;

    void UpdatePixelMaskStatistics(const PixelMaskStatistics &input);

    // Private functions assume that lock m is acquired
    void SetState(JFJochState curr_state,
        const std::optional<std::string> &message = {},
        BrokerStatus::MessageSeverity message_severity = BrokerStatus::MessageSeverity::Info);
    void MeasurementThread();
    void InitializeThread(std::unique_lock<std::mutex> ul);
    bool ImportPedestalG1G2(const JFJochReceiverOutput &receiver_output, size_t gain_level, size_t storage_cell = 0);
    bool ImportPedestalG0(const JFJochReceiverOutput &receiver_output);
    bool IsRunning() const; // Is state Busy/Pedestal/Measure
    void ResetError() noexcept;
    // The calibration steps report their own outcome through SetState and return false if the
    // sequence was cancelled or the data was not collected properly, so the caller does not
    // overwrite that with success.
    bool TakeDarkMaskInternal(std::unique_lock<std::mutex> &ul);
    void CalibrateDetector(std::unique_lock<std::mutex> ul);
    bool CalibrateJUNGFRAU(std::unique_lock<std::mutex> &ul);
    bool TakePedestalInternalG0(std::unique_lock<std::mutex> &ul);
    bool TakePedestalInternalG1(std::unique_lock<std::mutex> &ul, int32_t storage_cell = 0);
    bool TakePedestalInternalG2(std::unique_lock<std::mutex> &ul, int32_t storage_cell = 0);
    bool ImportDetectorSettings(const DetectorSettings& input);

    void UpdateROIDefinition();
public:
    JFJochStateMachine(const DiffractionExperiment& experiment,
                       JFJochServices &in_services,
                       Logger &logger,
                       const SpotFindingSettings &spot_finding_settings = SpotFindingSettings());
    ~JFJochStateMachine();

    void Initialize();
    void Pedestal();
    void Deactivate();
    // tokens: the bearer tokens that protect the new dataset's read endpoints; they replace the
    // previous dataset's tokens only once the start is accepted, and an empty list lifts the
    // protection.
    void Start(const DatasetSettings& settings, bool async = false, std::vector<std::string> tokens = {});
    // Does the Authorization header (if any) of a request entitle it to the current dataset?
    bool AcceptToken(const std::optional<std::string> &authorization) const;
    BrokerStatus WaitTillNotBusy(std::chrono::milliseconds timeout);

    BrokerStatus WaitTillMeasurementDone();
    BrokerStatus WaitTillMeasurementDone(std::chrono::milliseconds timeout);
    void Trigger();

    void Cancel();

    void SetCalibrationStatistics(const std::vector<JFCalibrationModuleStatistics> &input);

    DetectorSettings GetDetectorSettings() const;
    void LoadDetectorSettings(const DetectorSettings& settings);

    InstrumentMetadata GetInstrumentMetadata() const;
    void LoadInstrumentMetadata(const InstrumentMetadata& settings);

    FileWriterSettings GetFileWriterSettings() const;
    void LoadFileWriterSettings(const FileWriterSettings& settings);

    ImageFormatSettings GetImageFormatSettings() const;
    void LoadImageFormatSettings(const ImageFormatSettings& settings);
    void RawImageFormatSettings();
    void ConvImageFormatSettings();

    // return by value to ensure thread safety
    std::optional<MeasurementStatistics> GetMeasurementStatistics() const;
    std::vector<JFCalibrationModuleStatistics> GetCalibrationStatistics() const;

    BrokerStatus GetStatus() const;
    MultiLinePlot GetPlots(const PlotRequest &request) const;
    void GetPlotRaw(std::vector<float> &v,PlotType type, const std::string &roi) const;

    void SetSpotFindingSettings(const SpotFindingSettings& settings);
    SpotFindingSettings GetSpotFindingSettings() const;

    DetectorList GetDetectorsList() const;
    void SelectDetector(int64_t id);
    std::optional<DetectorStatus> GetDetectorStatus() const;

    void SetRadialIntegrationSettings(const AzimuthalIntegrationSettings& settings);
    AzimuthalIntegrationSettings GetRadialIntegrationSettings() const;

    std::string GetPreviewJPEG(const PreviewImageSettings& settings, int64_t image_number) const;
    std::string GetPreviewTIFF(int64_t image_number) const;
    std::string GetPedestalTIFF(size_t gain_level, size_t sc) const;

    void LoadInternalGeneratorImage(const void *data, size_t size, uint64_t image_number);
    void LoadInternalGeneratorImageTIFF(const std::string &s, uint64_t image_number);

    // Not thread safe - only for configuration in serial context
    DiffractionExperiment Experiment();

    // Function for debug only - UNSAFE for real operation
    void DebugOnly_SetState(JFJochState state,
        const std::optional<std::string> &message = {},
        BrokerStatus::MessageSeverity message_severity = BrokerStatus::MessageSeverity::Info);

    void SetROIDefinition(const ROIDefinition& input);
    ROIDefinition GetROIDefintion() const;

    std::vector<uint64_t> GetXFELPulseID() const;
    std::vector<uint64_t> GetXFELEventCode() const;

    std::string GetFullPixelMaskTIFF() const;
    std::string GetUserPixelMaskTIFF() const;
    std::vector<uint32_t> GetFullPixelMask() const;
    std::vector<uint32_t> GetUserPixelMask() const;

    void SetUserPixelMask(const std::vector<uint32_t> &v);
    void SetUserPixelMask(const CompressedImage &image);

    std::vector<DeviceStatus> GetDeviceStatus() const;

    void SetPreviewSocketSettings(const ZMQPreviewSettings &input);
    ZMQPreviewSettings GetPreviewSocketSettings();

    void SetMetadataSocketSettings(const ZMQMetadataSettings &input);
    ZMQMetadataSettings GetMetadataSocketSettings();

    void SetIndexingSettings(const IndexingSettings &input);
    IndexingSettings GetIndexingSettings() const;

    void SetRugnuxLive(bool enable, int64_t reserved_cpus, double deadline_s);
    void SetRugnuxModel(std::shared_ptr<const RugnuxModel> model, int64_t size_bytes);   // null = remove the model
    RugnuxLiveSettings GetRugnuxLive() const;
    RugnuxLiveSnapshot GetRugnuxStatus() const;

    void SetBraggIntegrationSettings(const BraggIntegrationSettings &input);
    BraggIntegrationSettings GetBraggIntegrationSettings() const;
    PixelMaskStatistics GetPixelMaskStatistics() const;

    void GetStartMessageFromBuffer(std::vector<uint8_t> &v);
    void GetImageFromBuffer(std::vector<uint8_t> &v, int64_t image_number = -1);
    ImageBufferStatus GetImageBufferStatus() const;
    void ClearImageBuffer() const;
    void AddDetectorSetup(const DetectorSetup& setup); // Not thread safe, only during setup

    std::optional<ScanResult> GetScanResult() const;

    void SetDarkMaskSettings(const DarkMaskSettings& settings);
    DarkMaskSettings GetDarkMaskSettings() const;

    ImagePusherStatus GetImagePusherStatus() const;
};