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 <array>
#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 "MtzTable.h"
#include "GoniometerAxis.h"
#include "DetectorTransformation.h"
#include "GridScanSettings.h"
#include "Reflection.h"
#include "CrystalLattice.h"
#include "IndexingSettings.h"
#include "XrayFluorescenceSpectrum.h"
#include "DetectorModules.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)
    // How much of a single sweep's blind cone this orientation makes unrecoverable: 0 = one sweep
    // about the spindle reaches everything the point group can give, 1 = a short lattice row lies on
    // the spindle and the whole cone is lost coherently - each shell losing its own 1 - cos theta(d),
    // theta(d) = asin(lambda/2d), up to theta_max at the resolution edge. Absent when the frame gives
    // nothing to decide on - too few spots, no lattice rows, no spindle - which is a third state, not
    // a value of one half: automation must treat absence as "engage" (see SpindleTrigger in
    // image_analysis/indexing/SpindleBlindFraction.h).
    std::optional<float> spindle_blind_fraction;

    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;
    // Set when the source is 4D, [image, channel, y, x]: the channel linked to
    std::optional<uint64_t> source_channel;
};

struct StartMessage {
    float detector_distance;
    float beam_center_x;
    float beam_center_y;
    // Where the undeflected beam lands on the detector, in pixels. This is NOT beam_center_x/y above:
    // that is the PONI, the foot of the perpendicular from the sample, and the two part company as
    // soon as the detector is tilted. Derived from the rest of the geometry
    // (DiffractionGeometry::GetDirectBeam_pxl), and carried so a consumer that wants the beam
    // position - which is what most of them mean, XDS among them - does not have to redo the tilt
    // arithmetic. Absent in a stream written before these fields existed.
    std::optional<float> direct_beam_x;
    std::optional<float> direct_beam_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;
    // How the assembled image sits in the detector plane, relative to the frame the PONI angles below
    // are stated in: mirrored in Y, and/or turned by this many quarter turns about the beam. A
    // different thing from mirror_y above - see DetectorSetup. Absence means the identity, which is
    // what every stream written before these fields existed carries.
    bool detector_orientation_mirror_y = false;
    int64_t detector_orientation_quarter_turns = 0;
    uint64_t bit_depth_image; // user data
    std::optional<uint64_t> bit_depth_readout;
    bool pixel_signed; // user data

    bool countrate_correction_enabled;
    // Maps a measured count c to its corrected value [c]; sent by a DECTRIS detector
    std::vector<uint32_t> countrate_correction_lookup_table;
    std::optional<bool> virtual_pixel_interpolation_enabled;

    float incident_energy;
    float incident_wavelength;
    std::optional<float> incident_wavelength_spread; // NXmx incident_wavelength_spread: FWHM of dlambda (Angstrom)
    // NXmx incident_beam_size: size of the beam where it meets the sample, x then y (m)
    std::optional<float> beam_size_x;
    std::optional<float> beam_size_y;

    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;

    // The detector's panels (DetectorModules.h), where the image's pixels are on the detector plane. Empty
    // for one flat detector, the image being the plane; then the CBOR key is absent.
    std::vector<DetectorModule> detector_modules;

    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;
};

// The atomic model as placed against the data, one vector per atom field in file order, for
// WritePlacedModel to write <prefix>_model.cif. A new residue starts wherever the model, chain or
// residue key changes. unit_cell and space_group_name are what WritePlacedModel is given (the data's);
// model_space_group_name is the group the coordinates were built in.
struct PlacedModelMessage {
    std::string name;
    UnitCell unit_cell;
    std::string space_group_name;
    std::string model_space_group_name;
    std::vector<int32_t> model_number;
    std::vector<std::string> chain;
    std::vector<std::string> residue_name;
    std::vector<int32_t> seq_num;
    std::vector<uint8_t> icode;            // ' ' where there is none
    std::vector<uint8_t> het_flag;         // 'A', 'H', or 0 where unknown
    std::vector<std::string> atom_name;
    std::vector<std::string> element;
    std::vector<uint8_t> altloc;           // 0 where there is none
    std::vector<int32_t> charge;
    std::vector<double> x, y, z;
    std::vector<float> b_iso, occupancy;
    std::vector<float> aniso;              // 6 per atom (u11 u22 u33 u12 u13 u23), empty when none has one
    // mmCIF label_asym_id, label_entity_id and label_seq_id per atom (label_seq INT32_MIN where none); empty
    // in a block from before they were carried
    std::vector<std::string> subchain, entity_id;
    std::vector<int32_t> label_seq;
    // The input model's metadata - entities, sequences, references, authors, software, secondary structure -
    // as the mmCIF categories gemmi writes for them, without coordinates; empty where there is none
    std::string metadata_mmcif;
};

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;
    // Run mean of the per-image spindle_blind_fraction, over the frames that had a value; absent
    // when none did. Written to /entry/MX/spindleBlindFractionMean.
    std::optional<float> spindle_blind_fraction;
    // The exact run-level counterpart, offline only (rugnux): the fraction (0-1) of unique
    // reflections to the run's resolution limit that the MEASURED point group could not recover
    // from the sweep's blind cone, in the crystal's indexed orientation. Not the per-image
    // worst-case bound. Written to /entry/MX/spindleLostUniqueFraction.
    std::optional<float> spindle_lost_unique_fraction;

    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). Both spellings are carried and a reader prefers
    // the name: only the name keeps the setting - "P 1 1 2" and "R 3:R" both come back from a number
    // as the reference setting - while the number is what a reader written before the name existed
    // understands. Set them together; a name without its number is a half-filled message.
    std::optional<std::string> space_group_name;
    std::optional<uint64_t> space_group_number;
    // Change of basis (3x3 integers, row major) from the setting the PER-IMAGE reflections and
    // lattices were written in to the setting of unit_cell / rotation_lattice above:
    // hkl_cell = M . hkl_written, and the same M takes each per-image lattice across. The per-image
    // data go to file as each image is processed, but the space group - and with it the conventional
    // setting the merge works in - is only chosen afterwards, so the two settings are not always the
    // same one. Absent means they are, i.e. the identity.
    std::optional<std::array<int32_t, 9>> reindex_matrix;

    // 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;
    // Per-image spindle_blind_fraction; NaN where the frame had no value. NaN is the CANNOT-SAY
    // state, which automation must treat as "engage" - it is not a zero (see SpindleTrigger in
    // image_analysis/indexing/SpindleBlindFraction.h).
    std::vector<float> v_spindle_blind_fraction;
    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;

    // Per-image disposition: what became of the image's observations in the merged data, as an index
    // into frame_disposition_codes (FrameDisposition in image_analysis/scale_merge/Merge.h - merged,
    // downgraded, rejected). Written to /entry/MX/frameDisposition with the vocabulary alongside it in
    // /entry/MX/frameDispositionCodes. The sweep-quality code above says what was SEEN over a stretch;
    // this says what was DONE about it. Offline (rugnux) only; the broker does not merge.
    std::vector<uint8_t> frame_disposition;
    std::vector<std::string> frame_disposition_codes;

    // The merged result of an analysis run inside the broker, for the writer to put beside the master
    // file as <file_prefix>.mtz, _maps.mtz and _model.cif. Absent where nothing was merged.
    std::optional<MtzTable> merged;             // MergedMtzTable (image_analysis/WriteReflections.h)
    std::optional<MtzTable> map_coefficients;   // ModelMaps::coefficients (ModelValidation.h)
    std::optional<PlacedModelMessage> placed_model;
};

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