Jungfraujoch/common/JFJochMessages.h

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

#ifndef JUNGFRAUJOCH_JFJOCHMESSAGES_H
#define JUNGFRAUJOCH_JFJOCHMESSAGES_H

#include <string>
#include <cstdint>
#include <map>
#include <vector>
#include <optional>
#include <cstring>
#include <stdexcept>
#include <nlohmann/json.hpp>

#include "CompressedImage.h"
#include "SpotToSave.h"
#include "UnitCell.h"
#include "GoniometerAxis.h"
#include "GridScanSettings.h"
#include "Reflection.h"
#include "CrystalLattice.h"
#include "IndexingSettings.h"
#include "XrayFluorescenceSpectrum.h"
#include "../symmetry/gemmi/symmetry.hpp"

constexpr const uint64_t user_data_release = 5;
constexpr const uint64_t user_data_magic_number = 0x52320000UL | user_data_release;

enum class CBORImageType {START, END, IMAGE, CALIBRATION, METADATA, NONE};

enum class FileWriterFormat : int {
    DataOnly = 0,
    NXmxLegacy = 1,
    NXmxVDS = 2,
    // TODO: NXmxTR = 3
    CBF = 4,
    TIFF = 5,
    NoFile = 6
};

struct ROIConfigBox {
    int64_t xmin, xmax, ymin, ymax;
};

struct ROIConfigCircle {
    double r, x, y;
};

struct ROIConfigAzim {
    double qmin, qmax;
};

struct ROIConfig {
    enum class ROIType {Box, Circle, Azim} type;
    std::string name;
    union {
        ROIConfigBox box;
        ROIConfigCircle circle;
        ROIConfigAzim azim;
    };
};

struct ROIMessage {
    int64_t sum;
    uint64_t sum_square;
    int64_t max_count;
    uint64_t pixels;
    uint64_t pixels_masked; // only used in the viewer for now
    int64_t x_weighted;
    int64_t y_weighted;
};

struct LatticeMessage {
    char centering;
    int64_t niggli_class;
    gemmi::CrystalSystem crystal_system;
};

struct DataMessage {
    int64_t number = INT64_MIN;
    CompressedImage image;

    std::optional<uint64_t> packets_expected;
    std::optional<uint64_t> packets_received;
    std::optional<float> image_collection_efficiency;

    std::optional<float> processing_time_s;
    std::optional<float> spot_finding_time_s;
    std::optional<float> indexing_time_s;
    std::optional<float> refinement_time_s;
    std::optional<float> bragg_prediction_time_s;
    std::optional<float> integration_time_s;

    std::vector<SpotToSave> spots;
    std::optional<int64_t> spot_count;
    std::optional<int64_t> spot_count_ice_rings;
    std::optional<int64_t> spot_count_low_res;

    std::vector<float> spot_plot_count;
    std::vector<float> spot_plot_intensity;
    std::vector<float> spot_plot_one_over_d_square;

    std::vector<float> az_int_profile;
    std::optional<float> bkg_estimate;

    std::optional<bool> indexing_result;
    std::optional<CrystalLattice> indexing_lattice;
    std::optional<UnitCell> indexing_unit_cell;
    std::optional<int64_t> spot_count_indexed;

    std::optional<float> profile_radius;
    std::optional<float> mosaicity_deg;

    std::optional<float> b_factor;

    std::vector<uint64_t> adu_histogram;

    uint64_t timestamp;
    uint32_t timestamp_base;

    uint32_t exptime;
    uint32_t exptime_base;

    std::string run_name;
    uint64_t run_number;

    std::optional<int64_t> pixel_sum;
    std::optional<int64_t> saturated_pixel_count;
    std::optional<int64_t> error_pixel_count;
    std::optional<int64_t> strong_pixel_count;
    std::optional<int64_t> min_viable_pixel_value;
    std::optional<int64_t> max_viable_pixel_value;
    std::optional<float> resolution_estimate;

    nlohmann::json user_data;

    std::optional<uint64_t> jf_info;
    std::optional<uint64_t> receiver_aq_dev_delay;
    std::optional<uint64_t> receiver_buf_available;
    std::optional<uint64_t> receiver_buf_in_sending;
    std::optional<uint64_t> receiver_buf_in_preparation;
    std::optional<uint64_t> storage_cell;

    std::optional<uint64_t> xfel_pulse_id;
    std::optional<uint64_t> xfel_event_code;

    std::map<std::string, ROIMessage> roi;

    std::optional<int64_t> original_number;
    std::vector<Reflection> reflections;

    std::vector<float> integration_B_logI;
    std::vector<float> integration_B_one_over_d_square;

    std::vector<float> integration_Isigma;
    std::vector<float> integration_Isigma_one_over_d_square;

    std::optional<float> beam_corr_x;
    std::optional<float> beam_corr_y;

    std::optional<LatticeMessage> lattice_type;
};

struct StartMessage {
    float detector_distance;
    float beam_center_x;
    float beam_center_y;

    uint64_t number_of_images;

    std::optional<uint64_t> images_per_trigger;

    uint64_t image_size_x;
    uint64_t image_size_y;
    uint64_t bit_depth_image; // user data
    std::optional<uint64_t> bit_depth_readout;
    bool pixel_signed; // user data

    bool countrate_correction_enabled;

    float incident_energy;
    float incident_wavelength;

    float frame_time;
    float count_time;

    int64_t saturation_value;
    std::optional<int64_t> error_value;

    float pixel_size_x;
    float pixel_size_y;
    float sensor_thickness;
    std::string sensor_material;

    std::optional<UnitCell> unit_cell; // user data
    std::optional<uint64_t> space_group_number; // user data
    uint64_t max_spot_count; // user data

    std::optional<uint64_t> storage_cell_number;
    uint64_t storage_cell_delay_ns;

    bool flatfield_enabled;
    bool pixel_mask_enabled;

    std::string arm_date;

    std::string sample_name; // user data
    std::string file_prefix; // user data
    int64_t images_per_file = 1; // user data

    std::vector<std::string> channels;

    std::string detector_description;
    std::string detector_serial_number;
    std::string run_name;
    uint64_t run_number;

    std::vector<std::string> gain_file_names;

    std::vector<ROIConfig> rois;

    std::optional<GridScanSettings> grid_scan;
    std::optional<GoniometerAxis> goniometer;
    float detector_translation[3];

    std::string source_type;
    std::string source_name;
    std::string instrument_name;


    uint64_t summation;

    std::vector<float> az_int_bin_to_q;
    std::vector<float> az_int_bin_to_two_theta;
    std::vector<float> az_int_bin_to_phi;
    std::optional<uint64_t> az_int_q_bin_count;
    std::optional<uint64_t> az_int_phi_bin_count;

    std::map<std::string, std::vector<uint32_t>> pixel_mask;

    std::map<std::string, float> threshold_energy;

    std::optional<float> total_flux;
    std::optional<float> attenuator_transmission;

    std::optional<bool> write_master_file;

    nlohmann::json user_data;

    std::optional<float> data_reduction_factor_serialmx;
    std::string experiment_group;

    std::string jfjoch_release;

    std::optional<uint64_t> socket_number; // This is number of socket in ZeroMQ
    std::string writer_notification_zmq_addr; // Socket to inform detector on writer done

    std::optional<bool> jungfrau_conversion_enabled;
    std::optional<float> jungfrau_conversion_factor;
    std::optional<bool> geometry_transformation_enabled;
    std::optional<std::string> summation_mode;
    std::optional<bool> overwrite;
    std::optional<FileWriterFormat> file_format;
    std::optional<bool> xfel_pulse_id;
    std::optional<float> ring_current_mA;
    std::optional<float> sample_temperature_K;

    IndexingAlgorithmEnum indexing_algorithm;
    GeomRefinementAlgorithmEnum geom_refinement_algorithm;

    std::optional<float> poni_rot1;
    std::optional<float> poni_rot2;
    std::optional<float> poni_rot3;

    XrayFluorescenceSpectrum fluorescence_spectrum;

    std::optional<bool> detect_ice_rings;
};

struct EndMessage {
    uint64_t max_image_number; // Counting from 1, i.e. 0 = no images collected
    std::optional<uint64_t> images_collected_count;
    std::optional<uint64_t> images_sent_to_write_count;
    std::optional<uint64_t> max_receiver_delay;
    std::optional<float> efficiency;
    std::optional<float> indexing_rate;
    std::optional<float> bkg_estimate;

    std::optional<std::string> end_date;

    std::string run_name;
    uint64_t run_number;

    std::map<std::string, std::vector<float>> az_int_result;
    std::map<std::string, std::vector<uint64_t>> adu_histogram;
    uint64_t adu_histogram_bin_width;

    std::optional<LatticeMessage> rotation_lattice_type;
    std::optional<CrystalLattice> rotation_lattice;

    std::vector<float> scale_factor;
};

struct MetadataMessage {
    std::string run_name;
    uint64_t run_number;
    std::vector<DataMessage> images;
};

#endif //JUNGFRAUJOCH_JFJOCHMESSAGES_H