Jungfraujoch/reader/JFJochXFELReader.h

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

#pragma once

#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <vector>

#include "JFJochReader.h"
#include "DetectorLayout.h"
#include "HDF5ImageSource.h"

class HDF5ReadOnlyFile;

// Reads serial (stills) frames from HDF5 files written by the XFEL pipelines, and places their panels on
// one flat image:
//
//  - a SACLA Cheetah file (one group per shot, /tag-N/data, the sensors stacked in the frame): the
//    sensors are placed where the file's /metadata puts them, with its camera length and pixel mask,
//    and the photon energy of every shot;
//  - a Cheetah CXI file of a single-panel detector (the frames stacked at /entry_1/data_1/data): one
//    flat panel with the pixel size, photon energy and camera length the file states, centred on the
//    beam - the run measures the centre;
//  - any of these, or files holding one frame each, laid out by a CrystFEL geometry file instead.
//
// Each panel is put on the flat image whole, as a block, its axes turned to the nearest image axis and its
// first pixel at a whole pixel, with everything between the panels and every masked or bad region masked.
// Where the panel really lies (the fraction of a pixel, a small turn) is its module's placement on the
// detector plane (DetectorModules), which the run maps its spots and predictions through. Pixel values are converted to photons
// with the gain where one is known (the geometry's, or one given or measured later), and are otherwise
// left in detector units.
//
// One wavelength and one camera length hold for the whole run: where the file states one per shot, it
// is the mean over every frame of the run.
//
// A path is one file, or a directory, every .h5/.cxi file of which is read in name order. Only a few
// files are held open at a time, so a run of thousands of files is read as well as one.
class JFJochXFELReader : public JFJochReader {
    std::vector<xfel::Mask> extra_masks_;
    struct Frame {
        size_t file;
        std::string dataset;
        int64_t index;             // along the first dimension of a 3D dataset; -1 for a 2D one
        std::string group;         // the shot group the geometry's % stands for, if any
    };

    xfel::Layout geom_;
    std::string source_;           // where geom_ came from, for messages
    std::vector<std::string> files_;
    std::vector<Frame> frames_;
    std::shared_ptr<JFJochReaderDataset> dataset_;

    // The most recently used files, held open; the least recent is closed past MAX_OPEN_FILES.
    static constexpr size_t MAX_OPEN_FILES = 8;
    mutable std::vector<std::pair<size_t, std::unique_ptr<HDF5ReadOnlyFile>>> open_files_;
    HDF5ReadOnlyFile &FileAt(size_t file) const;   // under hdf5_mutex
    std::shared_ptr<const HDF5ImageSource::RawFile> RawFileAt(size_t file) const;   // under hdf5_mutex

    size_t frame_width_ = 0, frame_height_ = 0;    // the stored frame
    size_t width_ = 0, height_ = 0;                // the flat image
    // For every stored pixel, where it lands on the flat image (-1: nowhere) and what one photon reads.
    std::vector<int64_t> destination_;
    std::vector<float> adu_per_photon_;
    bool counts_in_photons_ = false;
    double clen_m_ = 0, distance_mm_ = 0, photon_energy_eV_ = 0;
    // Masks stored with one plane per frame (a Cheetah CXI's per-shot mask), applied as each frame is read.
    std::vector<xfel::Mask> frame_masks_;

    // Where every frame is stored as one chunk of a plain little-endian number type - deflated (and
    // optionally byte-shuffled), as Cheetah writes it, or bitshuffle/LZ4 (or zstd), as Jungfraujoch does -
    // only the chunk's place in the file is asked of HDF5 under its lock; the bytes are read and decoded
    // outside it, on every reading thread at once, the way the NXmx reader reads its images
    // (HDF5ImageSource: the chunk lookup, its RawFile, and CompressedImage's decoders).
    struct Layout {
        uint64_t height = 0, width = 0;
        bool direct_chunks = false, shuffled = false, deflated = false;
        CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION;   // where not deflated
        bool stored_float = false, stored_signed = false;
        size_t stored_size = 0;
        bool operator==(const Layout &) const = default;
    };
    Layout layout_;
    // Positional-read handles on the files, opened as first needed (under hdf5_mutex), with the size of
    // each file's user block, which HDF5 chunk addresses count from.
    mutable std::vector<std::shared_ptr<const HDF5ImageSource::RawFile>> raw_files_;
    mutable std::vector<uint64_t> user_block_;
    static Layout ReadLayout(HDF5ReadOnlyFile &file, const Frame &frame);

    // The files, in the order given, every frame in them, and the run's mean of every value the geometry
    // takes from the files.
    void Open(const std::vector<std::string> &files);
    // Places the panels, given the gains, on the flat image; builds the dataset.
    void Place();

    bool LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset,
                     DataMessage &message,
                     std::vector<uint8_t> &buffer,
                     int64_t image_number,
                     bool update_dataset) override;

    template <class Buffer>
    void DecodeInto(int64_t image_number, Buffer &buffer) const;

public:
    JFJochXFELReader();
    ~JFJochXFELReader() override;

    // gain, where given, is the detector units per photon of every panel, whatever the geometry says.
    // files are read in the order given, their frames one after another.
    void ReadFiles(const std::vector<std::string> &files, const std::string &geometry_file,
                   std::optional<float> gain = std::nullopt);

    // Cheetah CXI files of a single-panel detector: one flat panel, placed by
    // what the file states, centred on the beam.
    static bool CanReadCXI(const std::string &path);
    // Why a Cheetah CXI file is not read without a geometry file (a tiled detector's slab), or nothing
    // if the path is not one.
    static std::optional<std::string> CXIRefusal(const std::string &path);
    void ReadCXI(const std::vector<std::string> &files, std::optional<float> gain = std::nullopt);

    // Cheetah HDF5 files with one group per shot and the panels' placement in /metadata (posx/posy in um, the
    // in-plane angle, pixel size, camera length) - as SACLA writes them - placed by that table.
    static bool CanReadPanelTable(const std::string &path);
    void ReadPanelTable(const std::vector<std::string> &files, std::optional<float> gain = std::nullopt);

    // Mask arrays in other files, in the frames' own layout (--mask-good / --mask-bad), applied with those
    // the data files state; set before reading.
    void ExtraMasks(const std::vector<xfel::Mask> &masks) { extra_masks_ = masks; }

    // Converts to photons with this gain from here on (detector units per photon, every panel).
    void ApplyGain(float adu_per_photon);

    // False where no gain is known, so the images are in detector units, not photons.
    [[nodiscard]] bool CountsInPhotons() const;
    // False where nothing read states the photon energy (a geometry file over frames of no known layout):
    // the caller gives it.
    [[nodiscard]] bool HasPhotonEnergy() const;

    [[nodiscard]] uint64_t GetNumberOfImages() const override;
    void Close() override;
    bool ReadRawImage(int64_t image_number, JFJochReaderRawImage &image) override;
    [[nodiscard]] std::vector<SpotToSave> ReadSpots(int64_t image) const override;
};