Jungfraujoch/reader/JFJochCBFReader.cpp

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

#include "JFJochCBFReader.h"

#include <algorithm>
#include <cctype>
#include <cstring>
#include <filesystem>
#include <map>
#include <optional>

#include "../common/JFJochException.h"
#include "../common/JFJochMath.h"

namespace {

bool HasCBFExtension(const std::filesystem::path &p) {
    std::string ext = p.extension().string();
    std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
    return ext == ".cbf";
}

// The sweep a file belongs to, as a template: everything before the trailing run of digits, the
// number of digits, and the extension. "o8_1_0042.cbf" -> {"o8_1_", 4}. A directory can hold several
// sweeps ("o8_1_*" beside "o8_2_*"), so collecting every .cbf in it would silently splice two
// crystals together; matching the template is what makes "point at any frame" safe.
struct Template {
    std::string prefix;
    size_t digits = 0;

    bool Matches(const std::string &name) const {
        if (name.size() != prefix.size() + digits + 4)   // + ".cbf"
            return false;
        if (name.compare(0, prefix.size(), prefix) != 0)
            return false;
        for (size_t i = 0; i < digits; i++)
            if (!std::isdigit(static_cast<unsigned char>(name[prefix.size() + i])))
                return false;
        return true;
    }
};

std::optional<Template> TemplateOf(const std::string &filename) {
    const std::filesystem::path p(filename);
    if (!HasCBFExtension(p))
        return {};
    const std::string stem = p.stem().string();
    size_t end = stem.size();
    while (end > 0 && std::isdigit(static_cast<unsigned char>(stem[end - 1])))
        end--;
    if (end == stem.size())
        return {};                                   // no trailing number: not part of a numbered sweep
    return Template{stem.substr(0, end), stem.size() - end};
}

std::vector<std::string> CollectSweep(const std::string &path) {
    std::filesystem::path p(path);
    const bool is_dir = std::filesystem::is_directory(p);
    const std::filesystem::path dir = is_dir ? p : p.parent_path();

    // Naming a frame selects ITS sweep. Naming a directory selects the sweep with the most frames in
    // it, which is the one a user pointing at a data directory means.
    std::optional<Template> want;
    if (!is_dir)
        want = TemplateOf(p.filename().string());

    std::map<std::pair<std::string, size_t>, std::vector<std::string>> sweeps;
    std::error_code ec;
    for (const auto &e : std::filesystem::directory_iterator(dir, ec)) {
        if (!e.is_regular_file() || !HasCBFExtension(e.path()))
            continue;
        const std::string name = e.path().filename().string();
        const auto t = TemplateOf(name);
        if (!t.has_value())
            continue;
        if (want.has_value() && !want->Matches(name))
            continue;
        sweeps[{t->prefix, t->digits}].push_back(e.path().string());
    }

    std::vector<std::string> out;
    for (auto &[key, files] : sweeps)
        if (files.size() > out.size())
            out = std::move(files);

    // The frame number is zero-padded in every PILATUS naming scheme in use, so within one template a
    // plain sort is the collection order.
    std::sort(out.begin(), out.end());
    return out;
}

// The base rotation axis of the instrument, in the internal frame (x along increasing detector
// column, y along increasing row, z along the beam). Both the spindle and the detector arm turn
// about it: the spindle in RotationAxis below, and the 2theta arm in ReadFiles.
const Coord BASE_AXIS(-1.0f, 0.0f, 0.0f);

// The axis a miniCBF sweep turns about, in the internal frame (x along increasing detector column,
// y along increasing row, z along the beam).
//
// The base axis is a convention: a miniCBF names its rotation axis but never states a direction, so
// this is the sign an NXmx master writes for the same instruments, and a file that needs the other
// one is settled from the data by the run's axis-sign rescue. The sense of the omega rotation below
// follows that same convention, so a sweep parked at a non-zero omega inherits whichever sign the
// base axis turns out to have.
//
// The head is base -> chi -> phi, so only the axes OUTSIDE the scanned one can tilt it. An omega
// scan turns about the base axis however the cradle is set - which is why a header carrying a large
// fixed chi still comes out as the base axis here - and only a phi scan is carried by chi and by
// omega. Chi turns about the beam, as the imgCIF axis convention has it, pointing back at the
// source; internal z points the other way, hence the minus. A kappa arm cannot be expressed at all:
// its inclination is a property of the hardware that no miniCBF header states.
Coord RotationAxis(const minicbf::Header &h) {
    const Coord base = BASE_AXIS;
    if (!minicbf::ScansPhi(h))
        return base;

    const auto rad = [](double deg) { return static_cast<float>(deg * PI / 180.0); };
    const Coord chi_axis(0.0f, 0.0f, -1.0f);
    const Coord tilted = RotMatrix(rad(h.chi_deg.value_or(0.0)), chi_axis) * base;
    return RotMatrix(rad(h.omega_deg.value_or(0.0)), base) * tilted;
}

}  // namespace

bool JFJochCBFReader::CanRead(const std::string &path) {
    std::error_code ec;
    if (std::filesystem::is_directory(path, ec))
        return !CollectSweep(path).empty();
    if (!HasCBFExtension(std::filesystem::path(path)))
        return false;
    try {
        return minicbf::ReadHeader(path).byte_offset;
    } catch (const JFJochException &) {
        return false;
    }
}

void JFJochCBFReader::ReadFiles(const std::string &path) {
    files_ = CollectSweep(path);
    if (files_.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF images found for " + path);

    header0_ = minicbf::ReadHeader(files_[0]);
    if (!header0_.byte_offset)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Unsupported CBF compression (only x-CBF_BYTE_OFFSET)");

    dataset_ = std::make_shared<JFJochReaderDataset>();
    dataset_->experiment = default_experiment;

    DetectorSetup detector = DetDECTRIS(header0_.nx, header0_.ny, header0_.detector, {});
    detector.PixelSize_um(static_cast<int64_t>(std::lround(header0_.pixel_x_m * 1e6)));
    detector.SensorThickness_um(static_cast<int64_t>(std::lround(header0_.thickness_m * 1e6)));
    detector.SensorMaterial(header0_.material);
    detector.SaturationLimit(SaturationLimitFromValue(header0_.count_cutoff));
    // Images are handed out as signed 32-bit whatever the file stored, so that is the depth the rest
    // of the code must see; the real overflow is the header's Count_cutoff, set above.
    detector.BitDepthImage(32);
    detector.MinFrameTime(std::chrono::microseconds(0));
    detector.MinCountTime(std::chrono::microseconds(0));
    detector.ReadOutTime(std::chrono::nanoseconds(0));
    dataset_->experiment.Detector(detector);

    dataset_->experiment.BeamX_pxl(static_cast<float>(header0_.beam_x_px));
    dataset_->experiment.BeamY_pxl(static_cast<float>(header0_.beam_y_px));
    dataset_->experiment.DetectorDistance_mm(static_cast<float>(header0_.distance_m * 1000.0));

    // A detector swung out on a 2theta arm, which small-molecule collection uses routinely. The arm
    // turns the detector about the sample, so it carries the square-on geometry with it: the header's
    // Detector_distance stays the distance along the detector normal and Beam_xy stays the point of
    // normal incidence, neither of which the swing moves - which is exactly what the PONI convention
    // wants, so the swing is a PONI rotation and nothing else in the header changes. It turns about
    // the base spindle axis, the four-circle geometry these headers describe having the arm and the
    // spindle on one axis; the imgCIF axis table such a header carries states the two with the same
    // vector.
    if (header0_.two_theta_deg != 0.0) {
        float rot1 = 0, rot2 = 0, rot3 = 0;
        PoniAnglesFromMatrix(RotMatrix(static_cast<float>(header0_.two_theta_deg * PI / 180.0), BASE_AXIS),
                             rot1, rot2, rot3);
        dataset_->experiment.PoniRot1_rad(rot1).PoniRot2_rad(rot2).PoniRot3_rad(rot3);
    }

    dataset_->experiment.IncidentEnergy_keV(WVL_1A_IN_KEV / static_cast<float>(header0_.wavelength_A));
    dataset_->experiment.FrameTime(
        std::chrono::duration_cast<std::chrono::nanoseconds>(
            std::chrono::duration<double>(header0_.period_s)),
        std::chrono::duration_cast<std::chrono::nanoseconds>(
            std::chrono::duration<double>(header0_.exposure_s)));

    // The rotation angle of every image, from its own header.
    std::vector<double> angles(files_.size());
    for (size_t i = 0; i < files_.size(); i++)
        angles[i] = minicbf::ReadHeader(files_[i]).start_angle_deg;

    double increment = header0_.angle_increment_deg;
    if (files_.size() > 1) {
        // Prefer the measured step over the header's nominal one, and unwrap a sweep that passes 360.
        double d = angles[1] - angles[0];
        if (d < -180.0) d += 360.0;
        if (std::abs(d) > 1e-6) increment = d;
    }
    dataset_->experiment.Goniometer(GoniometerAxis(header0_.axis_name,
                                                   static_cast<float>(angles.front()),
                                                   static_cast<float>(increment),
                                                   RotationAxis(header0_), {}));

    dataset_->error_value = -1;
    dataset_->experiment.ImagesPerTrigger(static_cast<int64_t>(files_.size()));

    // The untrusted pixels a PILATUS marks with a negative value: module gaps and the bad-pixel map.
    // They are the same on every frame of a sweep, so frame 0 defines the mask.
    std::vector<int32_t> first;
    minicbf::Read(files_[0], first);
    std::vector<uint32_t> mask(first.size(), 0);
    for (size_t i = 0; i < first.size(); i++)
        if (first[i] < 0)
            mask[i] = 1;
    dataset_->pixel_mask = std::make_shared<const PixelMask>(mask);

    SetStartMessage(dataset_);
}

uint64_t JFJochCBFReader::GetNumberOfImages() const {
    return files_.size();
}

void JFJochCBFReader::Close() {
    files_.clear();
    dataset_.reset();
}

template <class Buffer>
CompressedImage JFJochCBFReader::DecodeInto(int64_t image_number, Buffer &buffer) const {
    if (image_number < 0 || static_cast<size_t>(image_number) >= files_.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Image number out of range");

    const size_t npixel = static_cast<size_t>(header0_.nx) * static_cast<size_t>(header0_.ny);
    buffer.resize(npixel * sizeof(int32_t));

    // Decode straight into the caller's bytes: the pixels are plain int32 and nothing downstream has
    // to decompress them, so NO_COMPRESSION over that buffer is the whole image.
    const auto h = minicbf::ReadInto(files_[image_number],
                                     reinterpret_cast<int32_t *>(buffer.data()), npixel);
    if (static_cast<size_t>(h.nelem) != npixel)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CBF image size differs from the first image of the sweep");

    return CompressedImage(buffer.data(), buffer.size(),
                           static_cast<size_t>(header0_.nx), static_cast<size_t>(header0_.ny),
                           CompressedImageMode::Int32, CompressionAlgorithm::NO_COMPRESSION);
}

bool JFJochCBFReader::LoadImage_i(std::shared_ptr<JFJochReaderDataset> &dataset,
                                  DataMessage &message,
                                  std::vector<uint8_t> &buffer,
                                  int64_t image_number,
                                  bool update_dataset) {
    (void) update_dataset;
    if (!dataset)
        return false;

    // The image must outlive this call, so it is decoded straight into the caller's buffer - the same
    // thing the argument is for on the HDF5 path - and message.image only points at it.
    message.image = DecodeInto(image_number, buffer);
    message.number = image_number;
    return true;
}

std::shared_ptr<JFJochReaderRawImage> JFJochCBFReader::GetRawImage(int64_t image_number) {
    auto ret = std::make_shared<JFJochReaderRawImage>();
    ret->image = DecodeInto(image_number, ret->image_buffer);
    return ret;
}

std::vector<SpotToSave> JFJochCBFReader::ReadSpots(int64_t) const {
    return {};   // a raw CBF stores no analysis results
}