Jungfraujoch/reader/JFJochStartMessageReader.cpp

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

#include "JFJochStartMessageReader.h"

#include <filesystem>

#include "spdlog/fmt/fmt.h"
#include "../common/JFJochException.h"

namespace {
    // The data file name the writer gives file `file_number` (HDF5Metadata::DataFileName), without
    // the directory: the data files sit next to the master file.
    std::string DataFileName(const StartMessage &start, uint64_t file_number) {
        const std::string prefix = std::filesystem::path(start.file_prefix).filename().string();
        return fmt::format("{:s}_data_{:06d}.h5", prefix, file_number + 1);
    }
}

std::shared_ptr<JFJochReaderDataset> DatasetFromStartMessage(const StartMessage &start,
                                                             const DiffractionExperiment &defaults) {
    auto dataset = std::make_shared<JFJochReaderDataset>();
    dataset->arm_date = start.arm_date;
    dataset->jfjoch_release = start.jfjoch_release;
    dataset->error_value = start.error_value;
    dataset->file_detect_ice_rings = start.detect_ice_rings;

    auto &x = dataset->experiment;
    x = defaults;
    // The reader hands every image out as signed 32-bit, as JFJochHDF5Reader does.
    x.BitDepthImage(32);
    x.PixelSigned(true);
    x.FilePrefix(std::filesystem::path(start.file_prefix).filename().string());
    x.BeamX_pxl(start.beam_center_x);
    x.BeamY_pxl(start.beam_center_y);
    x.DetectorDistance_mm(start.detector_distance * 1000.0f);
    x.PoniRot1_rad(start.poni_rot1.value_or(0.0f));
    x.PoniRot2_rad(start.poni_rot2.value_or(0.0f));
    x.PoniRot3_rad(start.poni_rot3.value_or(0.0f));
    x.IncidentEnergy_keV(WVL_1A_IN_KEV / start.incident_wavelength);
    if (start.incident_wavelength_spread)
        x.BandwidthFWHM(start.incident_wavelength_spread.value() / start.incident_wavelength);
    x.DetectIceRings(start.detect_ice_rings.value_or(false));

    DetectorSetup detector = DetDECTRIS(start.image_size_x, start.image_size_y, start.detector_description, {});
    detector.PixelSize_um(start.pixel_size_x * 1e6f);
    detector.MirrorY(start.mirror_y);
    detector.ImageOrientation(DetectorOrientation(start.detector_orientation_mirror_y,
                                                  start.detector_orientation_quarter_turns));
    detector.SensorThickness_um(start.sensor_thickness * 1e6f);
    detector.SensorMaterial(start.sensor_material);
    detector.SaturationLimit(SaturationLimitFromValue(start.saturation_value));
    detector.BitDepthImage(32);
    detector.MinFrameTime(std::chrono::microseconds(0));
    detector.MinCountTime(std::chrono::microseconds(0));
    detector.ReadOutTime(std::chrono::nanoseconds(0));
    x.Detector(detector);

    x.FrameTime(std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<float>(start.frame_time)),
                std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::duration<float>(start.count_time)));

    x.Goniometer(start.goniometer);
    x.GridScan(start.grid_scan);
    x.Smargon(start.smargon_position);
    x.NumTriggers(1);
    x.ImagesPerTrigger(static_cast<int64_t>(start.number_of_images));

    InstrumentMetadata instrument;
    instrument.SourceName(start.source_name);
    instrument.SourceType(start.source_type);
    instrument.InstrumentName(start.instrument_name);
    x.ImportInstrumentMetadata(instrument);
    x.SampleName(start.sample_name);
    x.SampleTemperature_K(start.sample_temperature_K);
    x.RingCurrent_mA(start.ring_current_mA);
    x.TotalFlux(start.total_flux);
    x.AttenuatorTransmission(start.attenuator_transmission);
    if (start.beam_size_x && start.beam_size_y) {
        x.BeamSizeX_um(start.beam_size_x.value() * 1e6f);
        x.BeamSizeY_um(start.beam_size_y.value() * 1e6f);
    }

    if (!start.pixel_mask.empty())
        dataset->pixel_mask = std::make_shared<const PixelMask>(start.pixel_mask.begin()->second);
    else
        dataset->pixel_mask = std::make_shared<const PixelMask>(
                std::vector<uint32_t>(start.image_size_x * start.image_size_y));
    return dataset;
}

void JFJochStartMessageReader::Open(const StartMessage &start, const std::string &master_path) {
    if (start.images_per_file <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "The start message gives no images per file");

    number_of_images_ = start.number_of_images;
    images_per_file_ = static_cast<uint64_t>(start.images_per_file);
    const uint64_t files = (number_of_images_ + images_per_file_ - 1) / images_per_file_;

    const std::filesystem::path directory = std::filesystem::path(master_path).parent_path();
    HDF5ImageLocator::Layout layout;
    layout.format = FileWriterFormat::NXmxLegacy;
    layout.images_per_file = images_per_file_;
    for (uint64_t f = 0; f < files; f++)
        layout.legacy_files.push_back({(directory / DataFileName(start, f)).string(), "/entry/data/data"});
    {
        std::unique_lock ul(hdf5_mutex);
        image_source_.Configure(std::move(layout));
    }
    {
        std::unique_lock ul(m_);
        file_images_.assign(files, -1);
        ended_ = false;
        lost_.clear();
    }
    SetStartMessage(DatasetFromStartMessage(start, default_experiment));
}

void JFJochStartMessageReader::FileClosed(uint64_t file_number, uint64_t image_count) {
    {
        std::unique_lock ul(m_);
        if (file_number < file_images_.size())
            file_images_[file_number] = static_cast<int64_t>(image_count);
    }
    cv_.notify_all();
}

void JFJochStartMessageReader::Ended() {
    {
        std::unique_lock ul(m_);
        ended_ = true;
    }
    cv_.notify_all();
}

void JFJochStartMessageReader::Lost(const std::string &why) {
    {
        std::unique_lock ul(m_);
        if (!ended_)
            lost_ = why;
    }
    cv_.notify_all();
}

uint64_t JFJochStartMessageReader::GetNumberOfImages() const {
    return number_of_images_;
}

void JFJochStartMessageReader::Close() {
    std::unique_lock ul(hdf5_mutex);
    image_source_.Clear();
    SetStartMessage({});
}

bool JFJochStartMessageReader::WaitForImage(int64_t image_number) {
    if (image_number < 0 || static_cast<uint64_t>(image_number) >= number_of_images_)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Image out of bounds");
    const uint64_t f = static_cast<uint64_t>(image_number) / images_per_file_;
    std::unique_lock ul(m_);
    cv_.wait(ul, [&] { return file_images_[f] >= 0 || ended_ || !lost_.empty(); });
    if (file_images_[f] < 0) {
        if (!lost_.empty())
            throw LiveCollectionLost(lost_);
        return false;   // the collection ended without this file
    }
    return static_cast<int64_t>(static_cast<uint64_t>(image_number) % images_per_file_) < file_images_[f];
}

bool JFJochStartMessageReader::ReadRawImage(int64_t image_number, JFJochReaderRawImage &ret) {
    if (!WaitForImage(image_number))
        return false;
    // As JFJochHDF5Reader: only the chunk lookup under the HDF5 lock, the read after it.
    std::optional<HDF5ImageSource::DirectChunk> chunk;
    {
        std::unique_lock ul(hdf5_mutex);
        const auto loc = image_source_.Resolve(image_number);
        chunk = image_source_.PrepareDirectRead(loc);
        if (!chunk) {
            ret.image = image_source_.ReadImageAt(ret.image_buffer, loc);
            return true;
        }
    }
    ret.image = HDF5ImageSource::ReadDirect(ret.image_buffer, *chunk);
    return true;
}

bool JFJochStartMessageReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset, DataMessage &message,
                                           std::vector<uint8_t> &buffer, int64_t image_number, bool) {
    if (!dataset || !WaitForImage(image_number))
        return false;
    std::unique_lock ul(hdf5_mutex);
    message.image = image_source_.ReadImageAt(buffer, image_source_.Resolve(image_number));
    message.number = image_number;
    return true;
}