Jungfraujoch/writer/HDF5NXmx.cpp

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

#include <algorithm>
#include <cmath>

#include "HDF5NXmx.h"

#include "../common/DetectorOrientation.h"
#include "../common/DiffractionGeometry.h"
#include "../common/GitInfo.h"
#include "../include/spdlog/fmt/fmt.h"
#include "MakeDirectory.h"
#include "../common/time_utc.h"
#include "gemmi/symmetry.hpp"

std::string HDF5Metadata::MasterFileName(const StartMessage &start) {
    if (start.master_suffix.has_value())
        return fmt::format("{:s}_{:s}.h5", start.file_prefix, start.master_suffix.value());
    return fmt::format("{:s}_master.h5", start.file_prefix);
}

NXmx::NXmx(const StartMessage &start)
        : start_message(start),
          filename(HDF5Metadata::MasterFileName(start)) {
    uint64_t tmp_suffix;
    try {
        if (!start.arm_date.empty())
            tmp_suffix = parse_UTC_to_ms(start.arm_date);
    } catch (...) {
        tmp_suffix = std::chrono::system_clock::now().time_since_epoch().count();
    }

    tmp_filename = fmt::format("{}.{:08x}.tmp", filename, tmp_suffix);

    if (start.overwrite.has_value())
        overwrite = start.overwrite.value();

    MakeDirectory(filename);

    bool v1_10 = (start.file_format == FileWriterFormat::NXmxVDS)
                 || !start.hdf5_source_data.empty();

    hdf5_file = std::make_shared<HDF5File>(tmp_filename, v1_10);
    hdf5_file->Attr("file_name", filename);
    hdf5_file->Attr("HDF5_Version", hdf5_version());
    HDF5Group(*hdf5_file, "/entry").NXClass("NXentry").SaveScalar("definition", "NXmx");
    hdf5_file->SaveScalar("/entry/start_time", start.arm_date);

    Facility(start);
    Detector(start);
    Beam(start);
    Attenuator(start);
    UserData(start);
    MX(start);
    ROI(start);
    Fluorescence(start);
}

NXmx::~NXmx() {
    try {
        if (hdf5_file) {
            hdf5_file.reset();
            std::error_code ec;
            std::filesystem::remove(tmp_filename, ec);
        }
    } catch (...) {}
}



std::string HDF5Metadata::DataFileName(const StartMessage &msg, int64_t file_number) {
    if (file_number < 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "File number cannot be negative");

    if (msg.source_name == "SwissFEL") {
        if (file_number >= 10000)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Format doesn't allow for 10'000 or more files");
        else if (msg.detector_serial_number.empty())
            return fmt::format("{:s}{:04d}.JF.h5", msg.file_prefix, file_number + 1);
        else
            return fmt::format("{:s}{:04d}.{:s}.h5", msg.file_prefix, file_number + 1, msg.detector_serial_number);
    } else {
        if (file_number >= 1000000)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Format doesn't allow for 1 million or more files");
        else
            return fmt::format("{:s}_data_{:06d}.h5", msg.file_prefix, file_number + 1);
    }
}

void NXmx::LinkToData(const StartMessage &start, const EndMessage &end) {
    hsize_t total_images = end.max_image_number;
    hsize_t images_per_file = start.images_per_file;
    hsize_t file_count = 0;
    if (start.images_per_file > 0) {
        file_count = total_images / images_per_file;
        if (total_images % images_per_file > 0)
            file_count++;
    }

    HDF5Group(*hdf5_file, "/entry/data").NXClass("NXdata");

    for (uint32_t file_id = 0; file_id < file_count; file_id++) {
        char buff[32];
        snprintf(buff,32,"/entry/data/data_%06d", file_id+1);
        hdf5_file->ExternalLink(HDF5Metadata::DataFileName(start, file_id),
                                "/entry/data/data",
                                std::string(buff));
    }
}

void NXmx::LinkToData_VDS(const StartMessage &start, const EndMessage &end) {
    hsize_t total_images = end.max_image_number;
    hsize_t width = start.image_size_x;
    hsize_t height = start.image_size_y;
    if (total_images > 0) {
        HDF5Group(*hdf5_file, "/entry/data").NXClass("NXdata");
        auto data_dataset = VDS(start,
                                "/entry/data/data",
                                {total_images, height, width},
                                HDF5DataType(start.bit_depth_image / 8, start.pixel_signed),
                                start.error_value);
        data_dataset->Attr("image_nr_low", (int32_t) 1)
                .Attr("image_nr_high",(int32_t) total_images);

        if (start.max_spot_count > 0) {
            VDS(start, "/entry/MX/peakXPosRaw",{total_images, start.max_spot_count}, HDF5DataType(0.0f));
            VDS(start, "/entry/MX/peakYPosRaw",{total_images, start.max_spot_count}, HDF5DataType(0.0f));
            VDS(start, "/entry/MX/peakTotalIntensity",{total_images, start.max_spot_count}, HDF5DataType(0.0f));
            VDS(start, "/entry/MX/peakIceRingRes", {total_images, start.max_spot_count}, HDF5DataType(static_cast<int8_t>(0)));
            VDS(start, "/entry/MX/nPeaks", {total_images}, HDF5DataType((uint32_t) 0));
        }

        if (start.indexing_algorithm != IndexingAlgorithmEnum::None) {
            VDS(start, "/entry/MX/peakIndexed", {total_images, start.max_spot_count}, HDF5DataType(static_cast<int8_t>(0)));
            VDS(start, "/entry/MX/peakLattice", {total_images, start.max_spot_count}, HDF5DataType(static_cast<int8_t>(-1)));
            VDS(start, "/entry/MX/peakH", {total_images, start.max_spot_count}, HDF5DataType((int32_t) 0));
            VDS(start, "/entry/MX/peakK", {total_images, start.max_spot_count}, HDF5DataType((int32_t) 0));
            VDS(start, "/entry/MX/peakL", {total_images, start.max_spot_count}, HDF5DataType((int32_t) 0));
            VDS(start, "/entry/MX/peakDistEwaldSphere", {total_images, start.max_spot_count}, HDF5DataType((float) 0));
            VDS(start, "/entry/MX/latticeIndexed", {total_images,9}, HDF5DataType((float) 0))->Units("Angstrom");
            if (start.max_extra_lattices > 0)
                VDS(start, "/entry/MX/latticeIndexedExtra", {total_images, start.max_extra_lattices, 9}, HDF5DataType((float) 0))->Units("Angstrom");
        }

        if (!start.az_int_bin_to_q.empty()) {
            size_t azimuthal_bins = start.az_int_phi_bin_count.value_or(1);
            size_t q_bins = start.az_int_q_bin_count.value_or(1);
            if (q_bins > 0 && azimuthal_bins > 0) {
                VDS(start, "/entry/azint/image",
                    {total_images, azimuthal_bins, q_bins},
                    HDF5DataType(0.0f));
                VDS(start, "/entry/azint/image_count",
                    {total_images, azimuthal_bins, q_bins},
                    HDF5DataType(static_cast<uint64_t>(0UL)));
                // We make the link if we don't know if st.dev is recorded
                VDS(start, "/entry/azint/image_std",
                    {total_images, azimuthal_bins, q_bins},
                    HDF5DataType(0.0f));
            }
        }

        if (!start.rois.empty()) {
            // Per-image ROI results live in the data files; expose them in the master
            // through virtual datasets, one /entry/roi/<name> group per ROI.
            HDF5Group(*hdf5_file, "/entry/roi").NXClass("NXcollection");
            for (const auto &r: start.rois) {
                const std::string base = "/entry/roi/" + r.name;
                HDF5Group(*hdf5_file, base);
                VDS(start, base + "/max",    {total_images}, HDF5DataType((int64_t) 0));
                VDS(start, base + "/sum",    {total_images}, HDF5DataType((int64_t) 0));
                VDS(start, base + "/sum_sq", {total_images}, HDF5DataType((int64_t) 0));
                VDS(start, base + "/npixel", {total_images}, HDF5DataType((int64_t) 0));
                VDS(start, base + "/x",      {total_images}, HDF5DataType((float) 0));
                VDS(start, base + "/y",      {total_images}, HDF5DataType((float) 0));
            }
        }

        if (start.xfel_pulse_id.value_or(false)) {
            HDF5Group(*hdf5_file, "/entry/xfel").NXClass("NXcollection");
            VDS(start, "/entry/xfel/pulseID", {total_images}, HDF5DataType((uint64_t) 0));
            VDS(start, "/entry/xfel/eventCode", {total_images}, HDF5DataType((uint32_t) 0));
        }

        if (start.storage_cell_number)
            VDS(start,
                "/entry/detector/storage_cell_image",
                "/entry/instrument/detector/detectorSpecific/storage_cell_image",
                {total_images},
                HDF5DataType((uint8_t) 0));

        LinkToReflections_VDS(start, end);
    }
}

namespace {
    void SetFillValue(HDF5Dcpl &dcpl, const HDF5DataType &data_type,
                      const std::optional<int64_t> &fill_value) {
        if (!fill_value.has_value())
            return;
        const int64_t value = fill_value.value();
        if (data_type.IsSigned()) {
            switch (data_type.GetElemSize()) {
                case 1: dcpl.SetFillValue8(static_cast<int8_t>(value)); break;
                case 2: dcpl.SetFillValue16(static_cast<int16_t>(value)); break;
                case 4: dcpl.SetFillValue32(static_cast<int32_t>(value)); break;
                default: break;
            }
        } else {
            switch (data_type.GetElemSize()) {
                case 1: dcpl.SetFillValueU8(static_cast<uint8_t>(value)); break;
                case 2: dcpl.SetFillValueU16(static_cast<uint16_t>(value)); break;
                case 4: dcpl.SetFillValueU32(static_cast<uint32_t>(value)); break;
                default: break;
            }
        }
    }
}

void NXmx::LinkToData_ProcessingVDS(const StartMessage &start, const EndMessage &end) {
    if (start.hdf5_source_data.empty() || end.max_image_number == 0)
        return;

    const hsize_t total_images = end.max_image_number;
    const hsize_t width = start.image_size_x;
    const hsize_t height = start.image_size_y;

    HDF5Group(*hdf5_file, "/entry/data").NXClass("NXdata");

    HDF5DataSpace full_data_space({total_images, height, width});
    HDF5Dcpl dcpl;
    dcpl.SetChunking({1, height, width});
    // Same reason as the master's own virtual dataset: a source file that cannot be resolved reads
    // as the fill value, and HDF5's default fill is zero - which is a legitimate count. Fill with the
    // error marker instead, so an unreadable frame is masked rather than integrated as blank. Must be
    // set before SetVirtual.
    SetFillValue(dcpl, HDF5DataType(start.bit_depth_image / 8, start.pixel_signed), start.error_value);

    for (const auto &mapping: start.hdf5_source_data) {
        if (mapping.image_count == 0)
            continue;

        // The mapping is built from the number of images the run intended to process, while
        // total_images is the number it actually finished. A cancelled run, or one that skipped an
        // unreadable frame, legitimately ends up with fewer - so map what was written and drop the
        // rest. Rejecting the mismatch here would take the whole output file with it.
        if (mapping.virtual_first_image >= total_images)
            continue;

        const hsize_t image_count = std::min(static_cast<hsize_t>(mapping.image_count),
                                             total_images - mapping.virtual_first_image);

        const std::string source_dataset = mapping.dataset.empty()
                                           ? "/entry/data/data"
                                           : mapping.dataset;

        HDF5DataSpace virtual_data_space({total_images, height, width});
        virtual_data_space.SelectHyperslab(
                {static_cast<hsize_t>(mapping.virtual_first_image), 0, 0},
                {image_count, height, width}
        );

        const hsize_t source_extent_images = mapping.source_first_image + image_count;
        HDF5DataSpace source_data_space({source_extent_images, height, width});
        source_data_space.SelectHyperslab(
                {static_cast<hsize_t>(mapping.source_first_image), 0, 0},
                {image_count, height, width}
        );

        dcpl.SetVirtual(mapping.filename,
                        source_dataset,
                        source_data_space,
                        virtual_data_space);
    }

    auto data_dataset = std::make_unique<HDF5DataSet>(
            *hdf5_file,
            "/entry/data/data",
            HDF5DataType(start.bit_depth_image / 8, start.pixel_signed),
            full_data_space,
            dcpl
    );

    data_dataset->Attr("image_nr_low", static_cast<int32_t>(1))
                .Attr("image_nr_high", static_cast<int32_t>(total_images));
}

void NXmx::LinkToReflections_VDS(const StartMessage &start, const EndMessage &end) {
    if (end.integrated_reflections.empty())
        return;

    HDF5Group(*hdf5_file, "/entry/reflections").NXClass("NXcollection");

    for (size_t image = 0; image < end.integrated_reflections.size(); ++image) {
        if (end.integrated_reflections[image] <= 0)
            continue;

        if (start.images_per_file <= 0)
            continue;

        const uint64_t file_id = image / static_cast<uint64_t>(start.images_per_file);
        const uint64_t image_in_file = image % static_cast<uint64_t>(start.images_per_file);

        const std::string local_name = fmt::format("/entry/reflections/image_{:06d}", image);
        const std::string source_name = fmt::format("/entry/reflections/image_{:06d}", image_in_file);

        hdf5_file->ExternalLink(HDF5Metadata::DataFileName(start, file_id),
                                source_name,
                                local_name);
    }
}

std::unique_ptr<HDF5DataSet> NXmx::VDS(const StartMessage &start,
                                       const std::string &name,
                                       const std::vector<hsize_t> &dim,
                                       const HDF5DataType &data_type,
                                       const std::optional<int64_t> &fill_value) {
    return VDS(start, name, name, dim, data_type, fill_value);
}

std::unique_ptr<HDF5DataSet> NXmx::VDS(const StartMessage &start,
                                       const std::string &name_src,
                                       const std::string &name_dest,
                                       const std::vector<hsize_t> &dim,
                                       const HDF5DataType &data_type,
                                       const std::optional<int64_t> &fill_value) {
    if (dim.empty() || dim.size() > 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Dimension must be in range 1-3");

    hsize_t images_per_file = start.images_per_file;
    hsize_t file_count = 0;
    if (start.images_per_file > 0) {
        file_count = dim[0] / images_per_file;
        if (dim[0] % images_per_file > 0)
            file_count++;
    }


    HDF5DataSpace full_data_space(dim);
    HDF5Dcpl dcpl;

    if (dim.size() == 3)
        dcpl.SetChunking({1, dim[1], dim[2]});

    // Where a virtual dataset has no source file to read from, HDF5 hands back the fill value, and
    // that defaults to zero - so a data file that was not copied alongside the master reads as
    // frames of zero counts, with no error and no warning. Fill with the error marker instead: it is
    // outside underload_value..saturation_value, so a reader masks those frames rather than
    // integrating them.
    SetFillValue(dcpl, data_type, fill_value);

    for (hsize_t file_id = 0; file_id < file_count; file_id++) {
        hsize_t images_in_file = images_per_file;
        if (file_id == file_count - 1)
            images_in_file = dim[0] - (file_count - 1) * images_per_file;

        HDF5DataSpace virtual_data_space(dim);

        auto dim_src = dim;
        dim_src[0] = images_in_file;
        HDF5DataSpace src_data_space(dim_src);

        std::vector<hsize_t> start_dim(dim.size());
        start_dim[0] = file_id * images_per_file;
        virtual_data_space.SelectHyperslab(start_dim, dim_src);
        dcpl.SetVirtual(HDF5Metadata::DataFileName(start, file_id),
                        name_src,src_data_space, virtual_data_space);
    }

    return std::make_unique<HDF5DataSet>(*hdf5_file, name_dest, data_type, full_data_space, dcpl);
}

void NXmx::Detector(const StartMessage &start) {
    HDF5Group group(*hdf5_file, "/entry/instrument/detector");
    group.NXClass("NXdetector");
    SaveScalar(group, "depends_on", "/entry/instrument/detector/transformations/translation");

    // beam_center_x/y and the transformations chain (translation + rot1/2/3) are the refinable geometry;
    // they are written once at Finalize (see Metrology) from the values refined by the offline analysis.
    SaveScalar(group, "distance", start.detector_distance)->Units("m");
    SaveScalar(group, "detector_distance", start.detector_distance)->Units("m");

    SaveScalar(group, "count_time", start.count_time)->Units("s");
    SaveScalar(group, "frame_time", start.frame_time)->Units("s");

    SaveScalar(group, "sensor_thickness", start.sensor_thickness)->Units("m");

    if (start.threshold_energy.size() == 1)
        SaveScalar(group, "threshold_energy", start.threshold_energy.begin()->second)->Units("eV");

    SaveScalar(group, "x_pixel_size", start.pixel_size_x)->Units("m");
    SaveScalar(group, "y_pixel_size", start.pixel_size_y)->Units("m");
    SaveScalar(group, "sensor_material", start.sensor_material);
    SaveScalar(group, "description", start.detector_description);

    if (!start.detector_serial_number.empty()) {
        SaveScalar(group, "detector_number", start.detector_serial_number);
        SaveScalar(group, "serial_number", start.detector_serial_number);
    }

    SaveScalar(group, "bit_depth_image", start.bit_depth_image);
    if (start.bit_depth_readout)
        SaveScalar(group, "bit_depth_readout", start.bit_depth_readout.value());
    SaveScalar(group, "saturation_value", start.saturation_value);
    if (start.underload_value)
        SaveScalar(group, "underload_value", start.underload_value.value());
    if (start.error_value)
        SaveScalar(group, "error_value", start.error_value.value()); // this is not NXmx
    SaveScalar(group, "flatfield_applied", start.flatfield_enabled);
    SaveScalar(group, "pixel_mask_applied", start.pixel_mask_enabled);

    if (start.jungfrau_conversion_enabled)
        SaveScalar(group, "jungfrau_conversion_applied", start.jungfrau_conversion_enabled.value());
    if (start.jungfrau_conversion_factor)
        SaveScalar(group, "jungfrau_conversion_factor", start.jungfrau_conversion_factor.value())->Units("eV");

    SaveScalar(group, "geometry_transformation_applied", start.geometry_transformation_enabled.value_or(true));

    SaveScalar(group, "acquisition_type", "triggered");
    SaveScalar(group, "countrate_correction_applied", start.countrate_correction_enabled);
    SaveScalar(group, "number_of_cycles", start.summation);

    HDF5Group det_specific(group, "detectorSpecific");
    // Not NXmx: NXmx states the row direction only through the module axis vectors, which say what
    // the geometry is but not how it was arrived at. This records the assembly setting itself, so a
    // re-opened file knows whether the stored image was mirrored rather than having to infer it.
    SaveScalar(det_specific, "mirror_y", start.mirror_y);
    // Likewise for the discrete image orientation: the module axis vectors below carry its effect,
    // these two carry the setting.
    SaveScalar(det_specific, "detector_orientation_mirror_y", start.detector_orientation_mirror_y);
    SaveScalar(det_specific, "detector_orientation_quarter_turns",
               start.detector_orientation_quarter_turns);
    det_specific.NXClass("NXcollection");

    if (!start.jfjoch_release.empty())
        SaveScalar(det_specific, "jfjoch_release", start.jfjoch_release);
    SaveScalar(det_specific, "jfjoch_writer_release", jfjoch_version());

    if (start.summation_mode.has_value())
        SaveScalar(det_specific, "summation_mode", start.summation_mode.value());

    if (start.detect_ice_rings.has_value())
        SaveScalar(det_specific, "detect_ice_rings", start.detect_ice_rings.value());

    SaveScalar(det_specific, "x_pixels_in_detector", static_cast<uint32_t>(start.image_size_x));
    SaveScalar(det_specific, "y_pixels_in_detector", static_cast<uint32_t>(start.image_size_y));
    SaveScalar(det_specific, "software_git_commit", jfjoch_git_sha1());
    SaveScalar(det_specific, "software_git_date", jfjoch_git_date());
    if (start.storage_cell_number) {
        SaveScalar(det_specific, "storage_cell_number", static_cast<uint32_t>(start.storage_cell_number.value()));
        if (start.storage_cell_number.value() > 1)
            SaveScalar(det_specific, "storage_cell_delay", static_cast<uint32_t>(start.storage_cell_delay_ns))->Units(
                    "ns");
    }

    if (start.data_reduction_factor_serialmx)
        det_specific.SaveScalar("data_reduction_factor_serialmx", start.data_reduction_factor_serialmx.value());

    if (!start.gain_file_names.empty())
        det_specific.SaveVector("gain_file_names", start.gain_file_names);

    if (start.pixel_mask.size() == 1) {
        // Currently only handling single pixel mask
        CompressionAlgorithm mask_alg = CompressionAlgorithm::BSHUF_LZ4;
        if (start.file_format == FileWriterFormat::NXmxLegacy)
            mask_alg = CompressionAlgorithm::NO_COMPRESSION;
        std::vector<hsize_t> dims = {start.image_size_y, start.image_size_x};

        group.SaveVector("pixel_mask", start.pixel_mask.begin()->second, dims, mask_alg);
        hdf5_file->HardLink("/entry/instrument/detector/pixel_mask",
                            "/entry/instrument/detector/detectorSpecific/pixel_mask");
    }
}

void NXmx::Detector(const StartMessage &start, const EndMessage &end) {
    if (start.images_per_trigger.has_value() && start.images_per_trigger.value() > 0) {
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages", start.images_per_trigger.value());
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/ntrigger", (end.max_image_number + start.images_per_trigger.value() - 1)/ start.images_per_trigger.value());
    } else {
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages", end.max_image_number);
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/ntrigger", 1);
    }
    if (end.images_collected_count)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages_collected", end.images_collected_count.value());
    if (end.images_sent_to_write_count)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/nimages_written", end.images_sent_to_write_count.value());
    if (end.efficiency)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/data_collection_efficiency", end.efficiency.value());
    if (end.max_receiver_delay)
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/max_receiver_delay", end.max_receiver_delay.value());
}

void NXmx::MX(const StartMessage &start) {
    HDF5Group(*hdf5_file, "/entry/MX").NXClass("NXcollection");
    switch (start.indexing_algorithm) {
        case IndexingAlgorithmEnum::FFBIDX:
            hdf5_file->SaveScalar("/entry/MX/indexing_algorithm", "FFBIDX");
            break;
        case IndexingAlgorithmEnum::FFTW:
            hdf5_file->SaveScalar("/entry/MX/indexing_algorithm", "FFT (FFTW)");
            break;
        case IndexingAlgorithmEnum::FFT:
            hdf5_file->SaveScalar("/entry/MX/indexing_algorithm", "FFT (CUDA)");
            break;
        default:
            break;
    }

    switch (start.geom_refinement_algorithm) {
        case GeomRefinementAlgorithmEnum::BeamCenter:
            hdf5_file->SaveScalar("/entry/MX/geom_refinement_algorithm", "beam_center");
            break;
        case GeomRefinementAlgorithmEnum::Flex:
            hdf5_file->SaveScalar("/entry/MX/geom_refinement_algorithm", "flex");
            break;
        default:
            break;
    }
}

static void WriteROIDefinition(const HDF5Object &group, const ROIConfig &def) {
    switch (def.type) {
        case ROIConfig::ROIType::Box:
            SaveScalar(group, "type", "box");
            SaveScalar(group, "min_x_pxl", def.box.xmin);
            SaveScalar(group, "max_x_pxl", def.box.xmax);
            SaveScalar(group, "min_y_pxl", def.box.ymin);
            SaveScalar(group, "max_y_pxl", def.box.ymax);
            break;
        case ROIConfig::ROIType::Circle:
            SaveScalar(group, "type", "circle");
            SaveScalar(group, "center_x_pxl", def.circle.x);
            SaveScalar(group, "center_y_pxl", def.circle.y);
            SaveScalar(group, "radius_pxl", def.circle.r);
            break;
        case ROIConfig::ROIType::Azim:
            SaveScalar(group, "type", "azim");
            SaveScalar(group, "q_min_recipA", def.azim.qmin);
            SaveScalar(group, "q_max_recipA", def.azim.qmax);
            // phi_min == phi_max means a full ring; only record a sector.
            if (def.azim.phi_min != def.azim.phi_max) {
                SaveScalar(group, "phi_min_deg", def.azim.phi_min);
                SaveScalar(group, "phi_max_deg", def.azim.phi_max);
            }
            break;
    }
}

void NXmx::ROI(const StartMessage &start) {
    if (start.rois.empty())
        return;

    // ROI definitions go in /entry/roi_defs, kept separate from the per-image ROI
    // results (/entry/roi, written by the data-file plugin) so that older readers
    // iterating /entry/roi are not disturbed by the bitmap and definition subgroups.
    HDF5Group roi_group(*hdf5_file, "/entry/roi_defs");
    roi_group.NXClass("NXcollection");

    if (!start.roi_map.empty()) {
        // Per-pixel ROI bitmask: bit i (the bit_index below) marks pixels in ROI i.
        CompressionAlgorithm roi_alg = (start.file_format == FileWriterFormat::NXmxLegacy)
                                       ? CompressionAlgorithm::NO_COMPRESSION
                                       : CompressionAlgorithm::BSHUF_LZ4;
        std::vector<hsize_t> dims = {start.image_size_y, start.image_size_x};
        roi_group.SaveVector("roi_map", start.roi_map, dims, roi_alg);
    }

    for (size_t i = 0; i < start.rois.size(); i++) {
        HDF5Group g(roi_group, start.rois[i].name);
        SaveScalar(g, "bit_index", static_cast<uint16_t>(i));
        WriteROIDefinition(g, start.rois[i]);
    }
}

void NXmx::DetectorModule(const std::string &name, const std::vector<int32_t> &origin, const std::vector<int32_t> &size,
                          const std::vector<double> &fast_axis, const std::vector<double> &slow_axis,
                          const std::string &nx_axis, double pixel_size_mm) {
    HDF5Group module_group(*hdf5_file, "/entry/instrument/detector/" + name);

    module_group.NXClass("NXdetector_module");

    module_group.SaveVector("data_origin", origin);
    module_group.SaveVector("data_size", size);

    SaveScalar(module_group, "fast_pixel_direction", pixel_size_mm)->
            Transformation("m", "/entry/instrument/detector/transformations/" + nx_axis,
                           "", "", "translation", fast_axis,
                           {0,0,0}, "");

    SaveScalar(module_group, "slow_pixel_direction", pixel_size_mm)->
            Transformation("m", "/entry/instrument/detector/transformations/" + nx_axis,
                           "", "", "translation", slow_axis,
                           {0,0,0}, "");

    // The module origin coincides with the detector origin, so the offset is zero - but it is still
    // a translation, and NXmx types module_offset NX_FLOAT. Write a float with a proper unit vector
    // rather than an integer with a zero-length one, which is degenerate: the direction of a
    // zero-magnitude translation is arbitrary, not absent.
    SaveScalar(module_group, "module_offset", 0.0f)->
            Transformation("m", "/entry/instrument/detector/transformations/" + nx_axis,
                           "", "", "translation", {0, 0, 1});
}

void NXmx::Facility(const StartMessage &start) {
    HDF5Group(*hdf5_file, "/entry/source").NXClass("NXsource");
    SaveScalar(*hdf5_file, "/entry/source/name", start.source_name);

    if (!start.source_type.empty())
        SaveScalar(*hdf5_file, "/entry/source/type", start.source_type);

    if (start.ring_current_mA) {
        SaveScalar(*hdf5_file, "/entry/source/current", start.ring_current_mA.value() / 1000.0)->Units("A");
    }
    HDF5Group(*hdf5_file, "/entry/instrument").NXClass("NXinstrument");
    SaveScalar(*hdf5_file, "/entry/instrument/name", start.instrument_name);
}

void NXmx::Beam(const StartMessage &start) {
    HDF5Group group(*hdf5_file, "/entry/instrument/beam");
    group.NXClass("NXbeam");
    SaveScalar(group, "incident_wavelength", start.incident_wavelength)->Units("angstrom");
    if (start.incident_wavelength_spread)
        SaveScalar(group, "incident_wavelength_spread", start.incident_wavelength_spread.value())->Units("angstrom");
    if (start.total_flux)
        SaveScalar(group, "total_flux", start.total_flux.value())->Units("Hz");
    // NXmx asks for the beam size as one two-element array in the order x, y, so it is written
    // only when both have been given.
    if (start.beam_size_x && start.beam_size_y)
        group.SaveVector("incident_beam_size",
                         std::vector<float>{start.beam_size_x.value(), start.beam_size_y.value()})->Units("m");
}

void NXmx::Fluorescence(const StartMessage &start) {
    if (start.fluorescence_spectrum.empty())
        return;

    HDF5Group group(*hdf5_file, "/entry/instrument/fluorescence");
    group.NXClass("NXcollection");
    group.SaveVector("energy", start.fluorescence_spectrum.GetEnergy_eV())->Units("eV");
    group.SaveVector("data", start.fluorescence_spectrum.GetData());
}

void NXmx::Metrology(const StartMessage &start, const EndMessage &end) {
    // The beam centre and detector rotations may have been refined by the offline analysis (rugnux).
    // The master file is streamed but never read before Finalize, so the whole geometry is written
    // once here, at the end, from the refined values when present and the StartMessage values
    // otherwise - simpler than an open-time write followed by an in-place overwrite. The broker
    // leaves the refined fields empty, so the user-provided StartMessage geometry is used unchanged.
    const float beam_center_x = end.refined_beam_center_x.value_or(start.beam_center_x);
    const float beam_center_y = end.refined_beam_center_y.value_or(start.beam_center_y);
    const double rot1 = end.refined_poni_rot1.value_or(start.poni_rot1.value_or(0.0f));
    const double rot2 = end.refined_poni_rot2.value_or(start.poni_rot2.value_or(0.0f));
    const double rot3 = end.refined_poni_rot3.value_or(start.poni_rot3.value_or(0.0f));

    HDF5Group detector(*hdf5_file, "/entry/instrument/detector");
    SaveScalar(detector, "beam_center_x", beam_center_x)->Units("pixel");
    SaveScalar(detector, "beam_center_y", beam_center_y)->Units("pixel");

    // Where the beam actually lands. beam_center_x/y is the PONI - the foot of the perpendicular from
    // the sample - so on a tilted detector it is not the beam position, and a program that wants the
    // beam position (XDS's ORGX/ORGY, for one) would otherwise have to redo the tilt arithmetic. Built
    // from the very values written just above and just below, so the three can never disagree.
    //
    // Only where there is a geometry to compute it from. DiffractionGeometry refuses a distance under
    // 1 mm or a pixel size of zero, and a message can legitimately carry neither - a writer told to
    // prove it can create the files never describes a detector. This dataset is a convenience for
    // whoever reads the file later, so it is absent there rather than fatal.
    if (start.detector_distance * 1e3f >= 1.0f && start.pixel_size_x > 0.0f) {
        DiffractionGeometry geometry;
        geometry.BeamX_pxl(beam_center_x)
                .BeamY_pxl(beam_center_y)
                .DetectorDistance_mm(start.detector_distance * 1e3f)
                .PixelSize_mm(start.pixel_size_x * 1e3f)
                .Orientation(DetectorOrientation(start.detector_orientation_mirror_y,
                                                 start.detector_orientation_quarter_turns))
                .PoniRot1_rad(static_cast<float>(rot1))
                .PoniRot2_rad(static_cast<float>(rot2))
                .PoniRot3_rad(static_cast<float>(rot3));
        const auto direct_beam = geometry.GetDirectBeam_pxl();
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/direct_beam_x",
                   direct_beam.first)->Units("pixel");
        SaveScalar(*hdf5_file, "/entry/instrument/detector/detectorSpecific/direct_beam_y",
                   direct_beam.second)->Units("pixel");
    }

    HDF5Group transformations(*hdf5_file, "/entry/instrument/detector/transformations");
    transformations.NXClass("NXtransformations");

    // Internal frame (x = column, y = row downward, z = beam) -> McStas, a half turn about z. Written
    // as a subtraction from zero rather than a negation so that a zero component stays a positive
    // zero, and an untilted, unturned detector writes the same axis vectors it always has.
    auto to_mcstas = [](const Coord &v) {
        return std::vector<double>{0.0 - v.x, 0.0 - v.y, 0.0 + v.z};
    };

    // The discrete image orientation turns the offset from the PONI before the rot1/rot2/rot3 chain
    // acts, so it belongs to the module axes and to the translation, not to the arm rotations.
    const RotMatrix delta = DetectorOrientation(start.detector_orientation_mirror_y,
                                                start.detector_orientation_quarter_turns).Matrix();

    // Sample -> module origin (pixel 0, 0), which is where LabCoord(0, 0) puts it.
    std::vector<double> vector = to_mcstas(delta * Coord(-beam_center_x * start.pixel_size_x,
                                                         -beam_center_y * start.pixel_size_y,
                                                         start.detector_distance));

    double vector_length = sqrt(vector[0] * vector[0] + vector[1] * vector[1] + vector[2] * vector[2]);
    std::vector<double> vector_norm{vector[0] / vector_length, vector[1]/vector_length, vector[2]/vector_length};

    // The translation hangs off rot1, not off ".", so the tilt pivots about the SAMPLE: the rotations
    // then act on the whole sample->pixel vector including the distance, which is what
    // DiffractionGeometry does (poni_rot * {dx, dy, distance}). With the translation at the root the
    // panel origin stays put and only its orientation turns, displacing the origin by tens of mm at a
    // degree of tilt.
    SaveScalar(transformations, "translation", vector_length)->
            Transformation("m", "/entry/instrument/detector/transformations/rot1",
                           "detector", "detector_arm", "translation", vector_norm);

    // https://manual.nexusformat.org/classes/base_classes/NXdetector_module.html?highlight=nxdetector_module
    // The order of indices (i, j or i, j, k) is slow to fast.
    // though EIGER has is the other way round
    // Confusing....
    std::vector<int32_t> origin = {0, 0};
    std::vector<int32_t> size = {static_cast<int32_t>(start.image_size_y),
                                 static_cast<int32_t>(start.image_size_x)};

    // Jungfraujoch holds the tilt in the PyFAI PONI convention, poni_rot = Rz(-rot3)*Rx(-rot2)*Ry(+rot1)
    // (DiffractionGeometry::UpdatePoniRotMatrix), written in the internal frame: x along increasing
    // column, y along increasing row, z along the beam. NXmx uses McStas, which is that frame turned
    // 180 degrees about z - a proper rotation, not a mirror - so rotations about x and y reverse sense
    // while those about z keep it. Hence the axis vectors below: internal +y becomes (0,-1,0), the
    // -x of Rx(-rot2) becomes (1,0,0), and the -z of Rz(-rot3) is unchanged. (Do not confuse this
    // with the internal-to-imgCIF relation, which IS 180 degrees about x - that is the frame DIALS
    // reports in, one step further on.) A depends_on chain applies the
    // DEEPEST dependency first, so rot3 sits at the root to make the product R_rot3*R_rot2*R_rot1.
    SaveScalar(transformations, "rot1", rot1)->
            Transformation("rad",
                           "/entry/instrument/detector/transformations/rot2",
                           "detector", "detector_arm",
                           "rotation",
                           std::vector<double>{0.0, -1.0, 0.0});

    SaveScalar(transformations, "rot2", rot2)->
            Transformation("rad",
                           "/entry/instrument/detector/transformations/rot3",
                           "detector", "detector_arm",
                           "rotation",
                           std::vector<double>{1.0, 0.0, 0.0});

    SaveScalar(transformations, "rot3", rot3)->
            Transformation("rad",
                           ".",
                           "detector", "detector_arm",
                           "rotation",
                           std::vector<double>{0.0, 0.0, -1.0});

    DetectorModule("module", origin, size,
                   to_mcstas(delta * Coord(1, 0, 0)), to_mcstas(delta * Coord(0, 1, 0)),
                   "translation", start.pixel_size_x);
}

void SaveUnitCell( HDF5Group& group, const std::string& name, const UnitCell& unit_cell) {
    std::vector<float> v = {unit_cell.a, unit_cell.b, unit_cell.c,
                            unit_cell.alpha, unit_cell.beta, unit_cell.gamma};
    group.SaveVector(name, v);
}

void NXmx::Sample(const StartMessage &start, const EndMessage &end) {
    HDF5Group group(*hdf5_file, "/entry/sample");
    group.NXClass("NXsample");
    if (!start.sample_name.empty())
        group.SaveScalar("name", start.sample_name);

    // The offline analysis determines the space group only after merging, so it arrives on the end
    // message; prefer it over the (usually empty) start-message value the user supplied. Within the
    // end message the NAME is preferred and the number is the fallback, because only the name keeps
    // the setting - a number always reads back as the reference setting - and a sender written
    // before the name existed fills in only the number.
    const gemmi::SpaceGroup *sg = nullptr;
    if (end.space_group_name)
        sg = gemmi::find_spacegroup_by_name(*end.space_group_name);
    if (sg == nullptr && end.space_group_number)
        sg = gemmi::find_spacegroup_by_number(static_cast<int>(*end.space_group_number));
    if (sg == nullptr && start.space_group_number)
        sg = gemmi::find_spacegroup_by_number(static_cast<int>(*start.space_group_number));
    if (sg != nullptr) {
        group.SaveScalar("space_group_number", static_cast<uint64_t>(sg->number));
        group.SaveScalar("space_group", sg->xhm());
    }

    std::optional<UnitCell> unit_cell;
    std::optional<UnitCell> input_unit_cell;
    if (end.unit_cell)
        unit_cell = end.unit_cell;
    else if (end.rotation_lattice)
        unit_cell = end.rotation_lattice->GetUnitCell();
    else if (start.unit_cell) {
        unit_cell = start.unit_cell;
        input_unit_cell = start.unit_cell;
    }
    
    if (unit_cell)
        SaveUnitCell(group, "unit_cell", unit_cell.value());

    if (input_unit_cell)
        SaveUnitCell(group, "input_unit_cell", input_unit_cell.value());

    if (end.rotation_lattice) {
        group.SaveVector("ub_matrix",
                         end.rotation_lattice->GetUBMatrix(),
                         {1, 3, 3})
                ->Units("Angstrom^-1");
    }

    if (start.sample_temperature_K)
        group.SaveScalar("temperature", start.sample_temperature_K.value())->Units("K");

    std::string depends_on = ".";

    // Smargon chi/phi are static positioners closest to the sample, so they are appended
    // at the innermost end of the transformation chain (whatever depends_on currently is).
    auto write_smargon = [&start, &end](HDF5Group& transformations, std::string& depends_on) {
        if (!start.smargon_position)
            return;

        // One entry per image, even though neither angle moves. A reader takes the number of images
        // from the innermost axis of the sample chain when no axis varies - chi and phi are innermost
        // whenever they are present - so written as scalars, a still with a head position reads back
        // as a single image however many were collected.
        const auto n = static_cast<size_t>(std::max<int64_t>(end.max_image_number, 1));

        SaveVector(transformations, "chi", std::vector<double>(n, start.smargon_position->chi_deg))->
                Transformation("deg", depends_on, "", "smargon", "rotation",
                    {start.smargon_position->chi_axis.x, start.smargon_position->chi_axis.y,
                     start.smargon_position->chi_axis.z}, {0, 0, 0}, "");
        depends_on = "/entry/sample/transformations/chi";
        SaveVector(transformations, "phi", std::vector<double>(n, start.smargon_position->phi_deg))->
                Transformation("deg", depends_on, "", "smargon", "rotation",
                    {start.smargon_position->phi_axis.x, start.smargon_position->phi_axis.y,
                     start.smargon_position->phi_axis.z}, {0, 0, 0}, "");
        depends_on = "/entry/sample/transformations/phi";
    };

    // One chain, built from the base outwards, rather than the goniometer and the grid scan being
    // alternatives. NXmx applies the deepest dependency first, so the order here is the mounting
    // order: the grid stage is a BASE stage (an Aerotech xyz at SLS) that the spindle is mounted on,
    // the spindle carries the head, and the head carries the sample. So
    //     base -> grid -> omega -> chi -> phi -> sample
    // and a grid position therefore does NOT turn with omega. (A head-mounted grid stage exists too
    // - the Smargon can translate - and would sit on the other side of omega; only the base stage is
    // modelled for now, which is what is actually used.)
    const bool write_goniometer = (end.max_image_number > 0) && start.goniometer.has_value();
    const bool write_grid_scan = start.grid_scan.has_value();

    if (write_grid_scan) {
        HDF5Group grid_scan_group(group, "grid_scan");
        grid_scan_group.NXClass("NXcollection");

        SaveScalar(grid_scan_group, "snake_scan", start.grid_scan->IsSnakeScan());
        SaveScalar(grid_scan_group, "vertical_scan", start.grid_scan->IsVerticalScan());
        SaveScalar(grid_scan_group, "n_fast", start.grid_scan->GetNFast());
        SaveScalar(grid_scan_group, "step_x", start.grid_scan->GetGridStepX_um() * 1e-6)->Units("m");
        SaveScalar(grid_scan_group, "step_y", start.grid_scan->GetGridStepY_um() * 1e-6)->Units("m");
    }

    // A producer that sent the chain already gets it written verbatim - no angle is recomputed here,
    // which is what makes it possible to report positions that were measured rather than commanded.
    // Otherwise the same chain is built from the start message, so nothing is lost by not sending it.
    if (!end.transformations.empty()) {
        HDF5Group transformations(group, "transformations");
        transformations.NXClass("NXtransformations");
        hdf5_file->HardLink("/entry/sample/transformations","/entry/sample/goniometer");

        const std::string base = "/entry/sample/transformations/";
        for (const auto &axis: end.transformations) {
            const std::string parent = axis.GetDependsOn().empty() ? "." : base + axis.GetDependsOn();
            const std::vector<double> offset{axis.GetOffset().x, axis.GetOffset().y, axis.GetOffset().z};
            const std::vector<double> vector{axis.GetVector().x, axis.GetVector().y, axis.GetVector().z};
            const std::string type = axis.IsRotation() ? "rotation" : "translation";

            auto written = axis.IsConstant()
                ? SaveScalar(transformations, axis.GetName(),
                             axis.GetValues().empty() ? 0.0f : axis.GetValues().front())
                : SaveVector(transformations, axis.GetName(), axis.GetValues());
            written->Transformation(axis.GetUnits(), parent,
                                    axis.GetEquipment(), axis.GetEquipmentComponent(),
                                    type, vector, offset, "");
            depends_on = base + axis.GetName();

            // AXISNAME_end and the rotation width are derived here rather than carried on the wire.
            // They follow from the values themselves for a constant step, which is every case there
            // is today, so sending them would cost another array per axis and say nothing new.
            const auto &values = axis.GetValues();
            // Over the endpoints rather than the first pair: the values arrive as floats, and a
            // single difference carries that noise straight into the reported rotation width.
            const double step = (values.size() > 1)
                    ? (values.back() - values.front()) / static_cast<double>(values.size() - 1)
                    : 0.0;
            if (axis.IsRotation() && (step != 0.0)) {
                std::vector<double> end(values.size());
                for (size_t i = 0; i < values.size(); i++)
                    end[i] = values[i] + step;
                SaveVector(transformations, axis.GetName() + "_end", end)->Units(axis.GetUnits());
                SaveScalar(transformations, axis.GetName() + "_range_average",
                           static_cast<float>(step))->Units(axis.GetUnits());
                SaveScalar(transformations, axis.GetName() + "_range_total",
                           static_cast<float>(step * static_cast<double>(values.size())))
                        ->Units(axis.GetUnits());
            }
        }
        group.SaveScalar("depends_on", depends_on);
        return;
    }

    if (write_goniometer || write_grid_scan || start.smargon_position) {
        HDF5Group transformations(group, "transformations");
        transformations.NXClass("NXtransformations");
        hdf5_file->HardLink("/entry/sample/transformations","/entry/sample/goniometer");

        // Base stage first: everything else is mounted on it. The position containers hold one entry
        // per image and are empty if the scan stopped at the first image.
        if (write_grid_scan && (end.max_image_number > 0)) {
            SaveVector(transformations,"grid_scan_x", start.grid_scan->GetXContainer_m(end.max_image_number))
                    ->Transformation("m", depends_on, "", "",
                                     "translation", {1, 0, 0}, {0,0,0}, "");
            depends_on = "/entry/sample/transformations/grid_scan_x";

            SaveVector(transformations,"grid_scan_y", start.grid_scan->GetYContainer_m(end.max_image_number))
                    ->Transformation("m", depends_on, "", "",
                                     "translation", {0, 1, 0}, {0,0,0}, "");
            depends_on = "/entry/sample/transformations/grid_scan_y";
        }

        if (write_goniometer) {
            // Prefer the rotation axis refined by the offline analysis (rugnux); the broker leaves it empty
            // and the user-provided goniometer axis stands.
            const std::vector<double> axis_vector = end.refined_rotation_axis
                ? std::vector<double>{end.refined_rotation_axis->x, end.refined_rotation_axis->y,
                                      end.refined_rotation_axis->z}
                : start.goniometer->GetAxisVector();
            SaveVector(transformations, start.goniometer->GetName(),
                       start.goniometer->GetAngleContainer(end.max_image_number))->
                    Transformation("deg", depends_on, "", "",
                                   "rotation", axis_vector, {0,0,0}, "");

            SaveVector(transformations, start.goniometer->GetName() + "_end",
                       start.goniometer->GetAngleContainerEnd(end.max_image_number))
                    ->Units("deg");

            SaveScalar(transformations, start.goniometer->GetName() + "_range_average",
                       start.goniometer->GetIncrement_deg())
                    ->Units("deg");
            SaveScalar(transformations, start.goniometer->GetName() + "_range_total",
                       start.goniometer->GetIncrement_deg() * end.max_image_number)
                    ->Units("deg");
            depends_on = "/entry/sample/transformations/" + start.goniometer->GetName();
        } else if (write_grid_scan) {
            // No axis was given, but the sample still sits on a spindle - it simply does not turn.
            // Say so: a chain of translations alone is not something readers accept (dxtbx raises on
            // it outright, in get_dxtbx_goniometer). NXmx itself would allow a sample with no
            // goniometer, depends_on ".", so this is a reader constraint and not a standard one. At
            // 0 degrees the rotation is the identity whatever the axis points along, so the
            // conventional vector below carries no geometric claim - but 0 is a placeholder for an
            // angle nobody told us. Send the axis with step 0 and the angle the spindle really stood
            // at is written instead, which is how a grid scan records its head position.
            //
            // One entry per image, for the reason write_smargon gives: a reader takes the image
            // count from the innermost axis of the chain, and with no goniometer and no Smargon head
            // this placeholder IS the innermost one - grid_scan_x/y are translations and are passed
            // over. Written as a scalar it made a whole grid scan read back as a single image.
            const auto n = static_cast<size_t>(std::max<int64_t>(end.max_image_number, 1));
            SaveVector(transformations, "omega", std::vector<double>(n, 0.0))->
                    Transformation("deg", depends_on, "", "",
                                   "rotation", {-1, 0, 0}, {0,0,0}, "");
            depends_on = "/entry/sample/transformations/omega";
        }

        // Smargon chi/phi sit between the spindle and the sample.
        write_smargon(transformations, depends_on);

        if (write_goniometer) {
            auto helical = start.goniometer->GetHelicalStep();
            if (helical.has_value()) {
                SaveVector(transformations,
                           start.goniometer->GetName() + "_helical_x",
                           start.goniometer->GetXContainer_m(end.max_image_number))->
                        Transformation("m", depends_on, "", "",
                                       "translation", {1, 0, 0}, {0,0,0}, "");
                depends_on = "/entry/sample/transformations/" + start.goniometer->GetName() + "_helical_x";

                SaveVector(transformations,
                           start.goniometer->GetName() + "_helical_y",
                           start.goniometer->GetYContainer_m(end.max_image_number))->
                        Transformation("m", depends_on, "", "",
                                       "translation", {0, 1, 0}, {0,0,0}, "");
                depends_on = "/entry/sample/transformations/" + start.goniometer->GetName() + "_helical_y";

                SaveVector(transformations,
                           start.goniometer->GetName() + "_helical_z",
                           start.goniometer->GetZContainer_m(end.max_image_number))->
                        Transformation("m", depends_on, "", "",
                                       "translation", {0, 0, 1}, {0,0,0}, "");
                depends_on = "/entry/sample/transformations/" + start.goniometer->GetName() + "_helical_z";
            }
        }

    }

    group.SaveScalar("depends_on", depends_on);
}

void NXmx::Attenuator(const StartMessage &start) {
    if (start.attenuator_transmission) {
        HDF5Group group(*hdf5_file, "/entry/instrument/attenuator");
        group.NXClass("NXattenuator");
        SaveScalar(group, "attenuator_transmission", start.attenuator_transmission.value());
    }
}

void NXmx::WriteCalibration(const CompressedImage &image) {
    if (!hdf5_file)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "HDF5 file already closed");

    if (!calibration_group_created) {
        calibration_group_created = true;
        HDF5Group(*hdf5_file, "/entry/instrument/detector/calibration").NXClass("NXcollection");
    }
    SaveCBORImage("/entry/instrument/detector/calibration/" + image.GetChannel(), image);
}

void NXmx::SaveCBORImage(const std::string &hdf5_path, const CompressedImage &image) {
    std::vector<hsize_t> dims = {image.GetHeight(), image.GetWidth()};

    HDF5DataType data_type(image.GetMode());
    HDF5Dcpl dcpl;

    if (image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) {
        dcpl.SetCompression(image.GetCompressionAlgorithm(), 0);
        dcpl.SetChunking(dims);
    }

    HDF5DataSpace data_space(dims);
    auto dataset = std::make_unique<HDF5DataSet>(*hdf5_file, hdf5_path, data_type, data_space, dcpl);

    if (image.GetCompressionAlgorithm() == CompressionAlgorithm::NO_COMPRESSION)
        dataset->Write(data_type, image.GetCompressed());
    else
        dataset->WriteDirectChunk(image.GetCompressed(), image.GetCompressedSize(), {0, 0});

    dataset->Close();
}

void NXmx::AzimuthalIntegration(const StartMessage &start, const EndMessage &end) {
    if (!start.az_int_bin_to_q.empty()) {
        size_t phi_bins = start.az_int_phi_bin_count.value_or(1);
        size_t q_bin = start.az_int_q_bin_count.value_or(1);
        std::vector<hsize_t> dim = {phi_bins, q_bin};

        HDF5Group az_int_group(*hdf5_file, "/entry/azint");
        az_int_group.NXClass("NXcollection");
        if (start.file_format != FileWriterFormat::NXmxIntegrated) {
            az_int_group.SaveVector("bin_to_q", start.az_int_bin_to_q, dim)->Units("reciprocal Angstrom");
            if (!start.az_int_bin_to_two_theta.empty())
                az_int_group.SaveVector("bin_to_two_theta", start.az_int_bin_to_two_theta, dim)->Units("degrees");
            if (!start.az_int_bin_to_phi.empty())
                az_int_group.SaveVector("bin_to_phi", start.az_int_bin_to_phi, dim)->Units("degrees");
        }
        for (const auto &[x,y]: end.az_int_result) {
            if (x != "image")
                az_int_group.SaveVector(x, y, dim);
        }

        if (!start.az_int_map.empty() && start.az_int_map.size() == start.image_size_y * start.image_size_x)
            az_int_group.SaveVector("map", start.az_int_map, {start.image_size_y, start.image_size_x},
                CompressionAlgorithm::BSHUF_LZ4);
    }
}

void NXmx::ADUHistogram(const EndMessage &end) {
    if (!end.adu_histogram.empty()) {
        HDF5Group adu_histo_group(*hdf5_file, "/entry/instrument/detector/detectorSpecific/adu_histogram");
        adu_histo_group.SaveScalar("bin_width", end.adu_histogram_bin_width);
        for (const auto &[x, y]: end.adu_histogram)
            adu_histo_group.SaveVector(x, y);
    }
}

template<class T>
void SaveVectorIfMissing(HDF5Object &object,
                         const std::string &path,
                         const std::vector<T> &values,
                         const std::string &units = "") {
    if (values.empty() || object.Exists(path))
        return;

    auto dataset = object.SaveVector(path, values);
    if (!units.empty())
        dataset->Units(units);
}

void NXmx::Finalize(const EndMessage &end) {
    try {
        if (!hdf5_file)
            throw JFJochException(JFJochExceptionCategory::FileWriteError, "HDF5 file already closed");
        if (end.end_date) {
            hdf5_file->Attr("file_time", end.end_date.value());
            hdf5_file->SaveScalar("/entry/end_time", end.end_date.value());
            hdf5_file->SaveScalar("/entry/end_time_estimated", end.end_date.value());
        } else {
            std::string time_now = time_UTC(std::chrono::system_clock::now());
            hdf5_file->Attr("file_time", time_now);
            hdf5_file->SaveScalar("/entry/end_time", time_now);
            hdf5_file->SaveScalar("/entry/end_time_estimated", time_now);
        }

        Detector(start_message, end);
        Sample(start_message, end);

        // Write the refinable detector geometry (beam centre + transformations) now, from the values
        // refined by the offline analysis when present and the StartMessage otherwise. The master file
        // is never read before this point, so writing once here is simpler than an open-time write plus
        // an in-place overwrite - and, unlike the old overwrite, it also updates the translation vector
        // (which encodes the beam centre in the NXmx geometry chain).
        Metrology(start_message, end);

        AzimuthalIntegration(start_message, end);
        ADUHistogram(end);
        EndResultVectors(end);

        switch (start_message.file_format.value_or(FileWriterFormat::NXmxLegacy)) {
            case FileWriterFormat::NXmxLegacy:
                LinkToData(start_message, end);
                break;
            case FileWriterFormat::NXmxVDS:
                LinkToData_VDS(start_message, end);
                break;
            case FileWriterFormat::NXmxIntegrated:
                if (!start_message.hdf5_source_data.empty())
                    LinkToData_ProcessingVDS(start_message, end);
                break;
            default:
                break;
        }

        if (end.rotation_lattice)
            SaveVector(*hdf5_file, "/entry/MX/rotationLatticeIndexed", end.rotation_lattice->GetVector())
            ->Units("Angstrom");

        if (!end.rotation_extra_lattices.empty()) {
            std::vector<float> latt_info(9 * end.rotation_extra_lattices.size());
            for (int i = 0; i < end.rotation_extra_lattices.size(); i++) {
                auto vec = end.rotation_extra_lattices[i].GetVector();
                for (int j = 0; j < 9; j++)
                    latt_info[i * 9 + j] = vec[j];
            }
            SaveVector(*hdf5_file,
                       "/entry/MX/rotationLatticeIndexedExtra",
                       latt_info, {end.rotation_extra_lattices.size(), 9})
                    ->Units("Angstrom");
        }

        if (end.rotation_lattice_type)
            SaveScalar(*hdf5_file, "/entry/MX/rotationLatticeNiggliClass", end.rotation_lattice_type->niggli_class);

        // The setting the per-image reflections and lattices were written in is the one they were
        // indexed in, which is not always the setting of the cell above: the space group is chosen
        // after the images have gone to file, and choosing it can re-seat the lattice. Write the
        // change of basis between the two so the file is self-describing - absent means they agree.
        if (end.reindex_matrix)
            SaveVector(*hdf5_file, "/entry/MX/reindexMatrix",
                       std::vector<int32_t>(end.reindex_matrix->begin(), end.reindex_matrix->end()));

        if (end.indexing_rate) {
            SaveScalar(*hdf5_file, "/entry/MX/imageIndexedMean", end.indexing_rate.value());
        }
        if (end.bkg_estimate) {
            SaveScalar(*hdf5_file, "/entry/MX/bkgEstimateMean", end.bkg_estimate.value());
        }
        if (end.spindle_blind_fraction) {
            SaveScalar(*hdf5_file, "/entry/MX/spindleBlindFractionMean", end.spindle_blind_fraction.value());
        }
        if (end.spindle_lost_unique_fraction) {
            SaveScalar(*hdf5_file, "/entry/MX/spindleLostUniqueFraction",
                       end.spindle_lost_unique_fraction.value());
        }
        if (end.ice_ring_score_mean) {
            SaveScalar(*hdf5_file, "/entry/MX/iceRingScoreMean", end.ice_ring_score_mean.value());
        }

        hdf5_file->Close();
        hdf5_file.reset();
    } catch (const JFJochException &e) {
        hdf5_file.reset();
        std::error_code ec;
        std::filesystem::remove(tmp_filename, ec);
        throw;
    }

    if (std::filesystem::exists(filename) && !overwrite)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "File already exists");

    std::error_code ec;
    std::filesystem::rename(tmp_filename, filename, ec);
    if (ec)
        throw JFJochException(JFJochExceptionCategory::FileWriteError,
                              "Cannot rename temporary HDF5 master file " + tmp_filename +
                              " to " + filename + ": " + ec.message());
}

void NXmx::UserData(const StartMessage &start) {
    if (!start.user_data.empty()
        && start.user_data.contains("hdf5")
        && start.user_data["hdf5"].is_object()) {
        HDF5Group group(*hdf5_file, "/entry/user");
        group.NXClass("NXcollection");

        for (const auto &[x,y]: start.user_data["hdf5"].items()) {
            if (y.is_number())
                group.SaveScalar(x, y.get<double>());
            else if (y.is_string())
                group.SaveScalar(x, y.get<std::string>());
        }
    }
}

std::shared_ptr<HDF5File> NXmx::GetFile() {
    return hdf5_file;
}

void NXmx::EndResultVectors(const EndMessage &end) {
    if (!end.data_collection_efficiency.empty()) {
        HDF5Group det_specific(*hdf5_file, "/entry/instrument/detector/detectorSpecific");
        det_specific.NXClass("NXcollection");
        SaveVectorIfMissing(*hdf5_file,
                            "/entry/instrument/detector/detectorSpecific/data_collection_efficiency_image",
                            end.data_collection_efficiency);
    }

    if (!end.max_viable_pixel_value.empty() ||
        !end.min_viable_pixel_value.empty() ||
        !end.error_pixel_count.empty() ||
        !end.saturated_pixel_count.empty() ||
        !end.pixel_sum.empty()) {
        HDF5Group image_group(*hdf5_file, "/entry/image");
        image_group.NXClass("NXcollection");

        SaveVectorIfMissing(*hdf5_file, "/entry/image/max_value", end.max_viable_pixel_value);
        SaveVectorIfMissing(*hdf5_file, "/entry/image/min_value", end.min_viable_pixel_value);
        SaveVectorIfMissing(*hdf5_file, "/entry/image/error_pixels", end.error_pixel_count);
        SaveVectorIfMissing(*hdf5_file, "/entry/image/saturated_pixels", end.saturated_pixel_count);
        SaveVectorIfMissing(*hdf5_file, "/entry/image/pixel_sum", end.pixel_sum);
    }

    HDF5Group mx_group(*hdf5_file, "/entry/MX");
    mx_group.NXClass("NXcollection");

    SaveVectorIfMissing(*hdf5_file, "/entry/MX/peakCountUnfiltered", end.spot_count);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/peakCountIceRingRes", end.spot_count_ice_ring);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/peakCountIceRingControl", end.spot_count_ice_control);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/peakCountLowRes", end.spot_count_low_res);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/peakCountIndexed", end.spot_count_indexed);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/imageIndexed", end.image_indexed);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/indexedLatticeCount", end.indexed_lattice_count);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/bkgEstimate", end.v_bkg_estimate);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/spindleBlindFraction", end.v_spindle_blind_fraction);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/iceRingScore", end.ice_ring_score);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/profileRadius", end.profile_radius, "Angstrom^-1");
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/mosaicity", end.mosaicity, "deg");
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/bFactor", end.bFactor, "Angstrom^2");
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/resolutionEstimate", end.resolution_estimate, "Angstrom");
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/imageScaleFactor", end.image_scale_factor);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/integratedReflections", end.integrated_reflections);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/imageScaleCC", end.image_scale_cc);
    SaveVectorIfMissing(*hdf5_file, "/entry/MX/imageScaleMosaicity", end.image_scale_mosaicity, "deg");
    if (!end.niggli_class.empty())
        SaveVectorIfMissing(*hdf5_file, "/entry/MX/niggliClass", end.niggli_class);
    // Per-image sweep-quality code, with the vocabulary next to it so the codes can be read without
    // this source: sweepQuality[i] == 0 means the image is in no flagged range, otherwise it indexes
    // sweepQualityReasons from 1. Absent when the diagnostic did not run.
    if (!end.sweep_quality.empty() && !end.sweep_quality_reasons.empty()) {
        SaveVectorIfMissing(*hdf5_file, "/entry/MX/sweepQuality", end.sweep_quality);
        if (!hdf5_file->Exists("/entry/MX/sweepQualityReasons"))
            hdf5_file->SaveVector("/entry/MX/sweepQualityReasons", end.sweep_quality_reasons);
    }
}