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 <cstddef>
#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
// the detector plane:
//
//  - 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.
//
// A detector of several panels is handed out as stored: the image is the stored frame, and each panel a
// module (DetectorModules) placing its box of the frame on the detector plane exactly where the layout puts
// it, which the run maps its pixels, spots and predictions through. Pixels on no panel, and every masked or
// bad region, are masked. A single panel is put on an image of its own whole, as a block, its axes turned to
// the nearest image axis and its first pixel at a whole pixel; the rest of its placement is its module's. Pixel values are
// converted to photons with the gain where one is given or measured, one value for every panel, and are
// otherwise left in detector units: integers then stay integers, a 16-bit signed frame in its own type.
//
// 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 image handed out
    // A single panel: where its pixels land on the image - its first stored pixel, its size, and the image
    // index of its first pixel and the steps of +1 fs and +1 ss there. Empty for a stored frame.
    struct Block {
        size_t src = 0;
        int64_t width = 0, height = 0;
        int64_t dst = 0, step_fs = 1, step_ss = 0;
    };
    std::vector<Block> blocks_;
    // What one photon reads, the same on every panel; 1 where none is known.
    float adu_per_photon_ = 1.0f;
    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_;
    // Marks in bad every stored pixel the mask m, read as v, leaves out (on its panel, where it has one); v holds
    // the n pixels of the frame from first on.
    template <class T>
    void MarkBad(const xfel::Mask &m, const T *v, size_t first, size_t n, std::vector<uint8_t> &bad) const;

    // 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;
        // A 2- or 4-byte little-endian integer or a 4-byte float, read as stored; any other type HDF5 reads
        // as float.
        bool plain_type = false;
        bool operator==(const Layout &) const = default;
    };
    Layout layout_;
    // For each of frame_masks_, how it is stored where it is read as the frames are (ReadStored): plain
    // unsigned integers of the frame's size, stored alike in every file; plain_type false where HDF5 reads it.
    std::vector<Layout> frame_mask_layouts_;
    // 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. Like open_files_, only the most
    // recently used MAX_OPEN_FILES are held (raw_order_); a reader still using a released one keeps it.
    mutable std::vector<std::shared_ptr<const HDF5ImageSource::RawFile>> raw_files_;
    mutable std::vector<uint64_t> user_block_;
    mutable std::vector<size_t> raw_order_;
    static Layout ReadLayout(HDF5ReadOnlyFile &file, const Frame &frame);
    Layout FrameMaskLayout(const xfel::Mask &m) const;   // under hdf5_mutex

    // 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 detector plane; 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;

    // The frame's pixels, stored as layout says, as stored (plain_type), or as float, into stored; with
    // keep_compressed, a chunk decoded by CompressedImage (not deflated) is left as stored in the file, and
    // true returned. Given piece, the pixels are handed to piece(first pixel, their bytes, pixels) instead: a
    // deflated chunk a piece at a time, as it is inflated (InflateInPieces) or unshuffled, any other at once.
    template <class Buffer, class Piece = std::nullptr_t>
    bool ReadStored(const Frame &frame, const Layout &layout, Buffer &stored, bool keep_compressed,
                    Piece piece = nullptr) const;
    // The deflated stream compressed, bytes long once inflated, to piece a piece of whole pixels of size bytes
    // at a time.
    template <class Piece>
    void InflateInPieces(const std::vector<uint8_t> &compressed, size_t bytes, size_t size, const Frame &frame,
                         Piece piece) const;
    // The frame's stored pixels, as In, as the image handed out (as stored, or placed), as Out, each through
    // convert; a pixel masked marks is err.
    template <class Out, class In, class Buffer, class Convert>
    void PlaceInto(const Frame &frame, const std::vector<uint8_t> &masked, Out err, Buffer &buffer, Convert convert) const;
    template <class Buffer>
    CompressedImage 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);
    // Back to detector units: undoes ApplyGain where nothing read states a gain.
    void ClearGain();

    // False where no gain is known, so the images are in detector units, not photons.
    [[nodiscard]] bool CountsInPhotons() const;
    // True where the frames are stored as floats, which are rounded to whole detector units as read.
    [[nodiscard]] bool StoredFloat() 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;
};