Jungfraujoch/common/JFJochMessages.h

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

#pragma once

#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 "DetectorTransformation.h"
#include "GridScanSettings.h"
#include "Reflection.h"
#include "CrystalLattice.h"
#include "IndexingSettings.h"
#include "XrayFluorescenceSpectrum.h"
#include "../gemmi_gph/gemmi/symmetry.hpp"

constexpr const uint64_t user_data_release = 6;
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,
    NXmxIntegrated = 3,
    // 4 (CBF) and 5 (TIFF) removed - only HDF5 is written now. The values are kept
    // (deprecated) in the OpenAPI enum in broker/jfjoch_api.yaml for back compatibility.
    NoFile = 6
};

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

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

struct ROIConfigAzim {
    double qmin, qmax;
    // Azimuthal-angle sector in degrees; phi_min == phi_max means full ring.
    float phi_min, phi_max;
};

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 SmargonPosition {
    float phi_deg = 0;
    float chi_deg = 0;
    Coord chi_axis = {0, 0, 1};
    Coord phi_axis = {1, 0, 0}; // Default omega axis
};

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> azint_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::optional<float> preprocessing_time_s;
    std::optional<float> compression_time_s;
    std::optional<float> index_analysis_time_s;
    std::optional<float> image_scale_time_s;

    std::vector<SpotToSave> spots;
    std::optional<int64_t> spot_count;
    std::optional<int64_t> spot_count_ice_rings;
    // Spots in the ICE-FREE control flanks either side of the hexagonal rings, rescaled to the ring
    // bands' own q width - so spot_count_ice_rings / spot_count_ice_control is 1 when spots are spread
    // evenly and > 1 when they pile up on the rings. Pooled over the run (a ratio of two per-image
    // counts this small is meaningless on its own image), it is the second ice channel: TEXTURED ice
    // arrives as discrete spots and leaves the radial profile - and so ice_ring_score - flat.
    std::optional<float> spot_count_ice_control;
    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::vector<float> az_int_profile_std;
    std::vector<uint64_t> az_int_profile_count;

    std::optional<float> bkg_estimate;
    std::optional<float> ice_ring_score;   // strongest ice ring over the smooth radial background (1 = none)

    std::optional<bool> indexing_result;
    std::optional<CrystalLattice> indexing_lattice;
    std::vector<CrystalLattice> indexing_extra_lattices;
    std::optional<int64_t> indexing_lattice_count;
    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;

    std::optional<int64_t> integrated_reflections;

    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;

    std::optional<float> image_scale_factor;
    std::optional<float> image_scale_cc;
    std::optional<float> image_scale_mosaicity;

    std::optional<float> compression_ratio;
};

struct HDF5DataSourceMessage {
    std::string filename;
    std::string dataset = "/entry/data/data";

    uint64_t source_first_image = 0;
    uint64_t virtual_first_image = 0;
    uint64_t image_count = 0;
};

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;
    // Whether the assembled image is mirrored in Y relative to the detector's raw readout. True is
    // the MX convention - row 0 at the top of the detector, seen from the sample - and is what
    // Jungfraujoch has always produced, so it is also what absence of the field means.
    bool mirror_y = true;
    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;
    std::optional<float> incident_wavelength_spread; // NXmx incident_wavelength_spread: FWHM of dlambda (Angstrom)

    float frame_time;
    float count_time;

    int64_t saturation_value;
    std::optional<int64_t> error_value;
    std::optional<int64_t> underload_value; // NXmx: lowest valid 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
    uint64_t max_extra_lattices = 0;

    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::vector<uint16_t> az_int_map;

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

    // Per-pixel ROI bitmask (converted geometry), bit i set for ROI i in rois.
    std::vector<uint16_t> roi_map;

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

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

    std::optional<bool> write_master_file;
    std::optional<bool> write_images;

    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;

    std::vector<HDF5DataSourceMessage> hdf5_source_data;
    std::optional<std::string> master_suffix;

    std::optional<SmargonPosition> smargon_position;
};

struct EndMessage {
    // The sample transformation chain, base first (see DetectorTransformation). Optional: when it is
    // empty the writer builds the same chain itself from the start message, so a producer that does
    // not send it loses nothing and the wire format needs no version bump. It lives in the END
    // message because a future producer will want to report where the stage actually WENT, and that
    // is only known once the run is over.
    std::vector<DetectorTransformation> transformations;
    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<CrystalLattice> rotation_extra_lattices;
    std::optional<UnitCell> unit_cell;
    // Space group determined by the offline analysis (overrides the start message when writing the
    // master, since it is only known after merging).
    std::optional<uint64_t> space_group_number;

    // Geometry refined by the offline analysis. When present these overwrite the master-file geometry
    // at finalization (rugnux writes the refined values; the broker leaves them empty and keeps the
    // user-provided StartMessage geometry). Beam centre in pixels, detector tilt (PONI rot1/2/3) in
    // radians, rotation axis as a unit vector.
    std::optional<float> refined_beam_center_x;
    std::optional<float> refined_beam_center_y;
    std::optional<float> refined_poni_rot1;
    std::optional<float> refined_poni_rot2;
    std::optional<float> refined_poni_rot3;
    std::optional<Coord> refined_rotation_axis;

    // Vectors with end result:
    std::vector<float> data_collection_efficiency;
    std::vector<int32_t> spot_count;
    std::vector<int32_t> spot_count_ice_ring;
    std::vector<float> spot_count_ice_control;
    std::vector<int32_t> spot_count_low_res;
    std::vector<int32_t> spot_count_indexed;
    std::vector<uint8_t> image_indexed;
    std::vector<int32_t> indexed_lattice_count;
    std::vector<float> v_bkg_estimate;
    std::vector<float> profile_radius;
    std::vector<float> mosaicity;
    std::vector<float> bFactor;
    std::vector<float> resolution_estimate;
    std::vector<int64_t> min_viable_pixel_value;
    std::vector<int64_t> max_viable_pixel_value;
    std::vector<int32_t> saturated_pixel_count;
    std::vector<int32_t> error_pixel_count;
    std::vector<int32_t> integrated_reflections;
    std::vector<uint8_t> niggli_class;
    std::vector<int64_t> pixel_sum;

    std::vector<float> image_scale_factor;
    std::vector<float> image_scale_cc;
    std::vector<float> image_scale_mosaicity;
    // Per-image ice strength. Note the name carries no v_ prefix, unlike v_bkg_estimate above - it
    // was shipped that way and the CBOR key is part of the stream format.
    std::vector<float> ice_ring_score;
    // Run mean of the above, the single "how icy was this dataset" number (1 = no ice). The
    // bkg_estimate scalar's counterpart; written to /entry/MX/iceRingScoreMean.
    std::optional<float> ice_ring_score_mean;

    // Per-image sweep-quality code: 0 = the image falls in no flagged range, otherwise 1 + the
    // SweepQualityReason of the range it falls in (image_analysis/scale_merge/Merge.h). Written to
    // /entry/MX/sweepQuality, with the vocabulary alongside it in /entry/MX/sweepQualityReasons.
    // Empty when the sweep-quality diagnostic did not run, so nothing could be looked for - which is
    // not the same as every image being clean. Offline (rugnux) only; the broker does not merge.
    std::vector<uint8_t> sweep_quality;
    std::vector<std::string> sweep_quality_reasons;
};

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