Jungfraujoch/reader/HDF5ImageLocator.cpp

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

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

#include <algorithm>

namespace {
    // Coalesce consecutive single-image mappings into one contiguous range when the source and
    // virtual images stay contiguous in the same file/dataset.
    void AppendOrExtendSourceMapping(std::vector<HDF5DataSourceMessage> &ret,
                                     const std::string &filename,
                                     const std::string &dataset,
                                     uint64_t source_first_image,
                                     uint64_t virtual_first_image,
                                     uint64_t image_count,
                                     std::optional<uint64_t> source_channel = {}) {
        if (image_count == 0)
            return;

        if (!ret.empty()) {
            auto &last = ret.back();
            if (last.filename == filename
                && last.dataset == dataset
                && last.source_channel == source_channel
                && last.source_first_image + last.image_count == source_first_image
                && last.virtual_first_image + last.image_count == virtual_first_image) {
                last.image_count += image_count;
                return;
            }
        }

        ret.push_back(HDF5DataSourceMessage{
            .filename = filename,
            .dataset = dataset,
            .source_first_image = source_first_image,
            .virtual_first_image = virtual_first_image,
            .image_count = image_count,
            .source_channel = source_channel
        });
    }

    // A 4D mapping, [image, channel, y, x], may cover only some of the channels; the reader wants
    // the first one.
    bool CoversFirstChannel(const HDF5VirtualDatasetMapping &mapping) {
        return mapping.virtual_start.size() != 4 || mapping.virtual_start[1] == 0;
    }

    std::optional<hsize_t> SourceChannel(const HDF5VirtualDatasetMapping &mapping) {
        if (mapping.source_start.size() == 4)
            return mapping.source_start[1];
        return {};
    }
}

void HDF5ImageLocator::Configure(Layout layout) {
    file_cache_.clear();
    layout_ = std::move(layout);
}

void HDF5ImageLocator::Clear() {
    file_cache_.clear();
    layout_ = Layout{};
}

std::shared_ptr<HDF5ReadOnlyFile> HDF5ImageLocator::OpenCached(const std::string &path) const {
    auto it = file_cache_.find(path);
    if (it != file_cache_.end())
        return it->second;
    auto file = std::make_shared<HDF5ReadOnlyFile>(path);
    file_cache_[path] = file;
    return file;
}

HDF5ImageLocator::Location HDF5ImageLocator::Resolve(int64_t global_image) const {
    if (global_image < 0)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Image out of bounds");

    if (layout_.format == FileWriterFormat::NXmxLegacy) {
        const uint32_t file_id = global_image / layout_.images_per_file;
        const uint32_t local_index = global_image % layout_.images_per_file;
        const auto &data_file = layout_.legacy_files.at(file_id);
        return {OpenCached(data_file.path), local_index, data_file.path, data_file.dataset};
    }

    if (layout_.format == FileWriterFormat::NXmxVDS
        && layout_.data_layout == HDF5DataSetLayout::VIRTUAL) {
        const auto image = static_cast<hsize_t>(global_image);
        for (const auto &mapping: layout_.vds_mappings) {
            if (!CoversFirstChannel(mapping) || !mapping.ContainsVirtualImage(image))
                continue;
            return {OpenCached(mapping.filename), static_cast<uint32_t>(mapping.SourceImage(image)),
                    mapping.filename, mapping.dataset, SourceChannel(mapping).value_or(0)};
        }
        throw JFJochException(JFJochExceptionCategory::HDF5,
                              "Image not covered by /entry/data/data VDS mappings");
    }

    // Contiguous / integrated: pixels live in the master file at the global index.
    if (!layout_.master_file)
        throw JFJochException(JFJochExceptionCategory::HDF5, "Master file not loaded");
    return {layout_.master_file, static_cast<uint32_t>(global_image), layout_.master_filename};
}

std::vector<std::string> HDF5ImageLocator::DataFiles() const {
    std::vector<std::string> ret;
    if (layout_.format == FileWriterFormat::NXmxLegacy) {
        for (const auto &f: layout_.legacy_files)
            ret.push_back(f.path);
    } else if (layout_.format == FileWriterFormat::NXmxVDS
               && layout_.data_layout == HDF5DataSetLayout::VIRTUAL) {
        for (const auto &mapping: layout_.vds_mappings)
            if (CoversFirstChannel(mapping)
                && std::find(ret.begin(), ret.end(), mapping.filename) == ret.end())
                ret.push_back(mapping.filename);
    } else if (!layout_.master_filename.empty())
        ret.push_back(layout_.master_filename);
    return ret;
}

std::vector<HDF5DataSourceMessage> HDF5ImageLocator::GetSourceMapping(uint64_t first_image,
                                                                      std::optional<uint64_t> image_count,
                                                                      uint64_t total_images,
                                                                      uint64_t stride) const {
    if (stride == 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Image stride cannot be zero");

    if (first_image > total_images)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "First image outside dataset range");

    // With a stride the caller's count is the number of OUTPUT images, which spans
    // (count - 1) * stride + 1 source images - not count of them.
    const uint64_t requested_count =
        image_count.value_or((total_images - first_image + stride - 1) / stride);
    if (requested_count > 0 && first_image + (requested_count - 1) * stride >= total_images)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Requested image range outside dataset range");

    std::vector<HDF5DataSourceMessage> ret;
    if (requested_count == 0)
        return ret;

    // Integrated / contiguous source: link directly to the original master file.
    if (layout_.format == FileWriterFormat::NXmxVDS && layout_.data_layout != HDF5DataSetLayout::VIRTUAL) {
        // One extendable run when the images are consecutive; strided output needs one mapping per image,
        // which AppendOrExtendSourceMapping leaves separate because the source indices are not adjacent.
        for (uint64_t local_image = 0; local_image < requested_count; ++local_image)
            AppendOrExtendSourceMapping(ret, layout_.master_filename, "/entry/data/data",
                                        first_image + local_image * stride, local_image, 1);
        return ret;
    }

    // VDS source: expand VDS mappings to original source files, not to the VDS master.
    if (layout_.format == FileWriterFormat::NXmxVDS && layout_.data_layout == HDF5DataSetLayout::VIRTUAL) {
        for (uint64_t local_image = 0; local_image < requested_count; ++local_image) {
            const hsize_t virtual_image = first_image + local_image * stride;

            bool found = false;
            for (const auto &mapping: layout_.vds_mappings) {
                if (!CoversFirstChannel(mapping) || !mapping.ContainsVirtualImage(virtual_image))
                    continue;

                const uint64_t source_image = mapping.SourceImage(virtual_image);
                const std::string dataset = mapping.dataset.empty() ? "/entry/data/data" : mapping.dataset;

                AppendOrExtendSourceMapping(ret, mapping.filename, dataset, source_image, local_image, 1,
                                            SourceChannel(mapping));
                found = true;
                break;
            }

            if (!found)
                throw JFJochException(JFJochExceptionCategory::HDF5,
                                      "Image not covered by /entry/data/data VDS mappings");
        }

        return ret;
    }

    // Legacy source: link directly to the linked data files.
    if (layout_.format == FileWriterFormat::NXmxLegacy) {
        if (layout_.images_per_file == 0)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Cannot generate HDF5 source mapping: images_per_file is zero");

        for (uint64_t local_image = 0; local_image < requested_count; ++local_image) {
            const uint64_t source_global_image = first_image + local_image * stride;
            const uint64_t file_id = source_global_image / layout_.images_per_file;
            const uint64_t source_image = source_global_image % layout_.images_per_file;

            if (file_id >= layout_.legacy_files.size())
                throw JFJochException(JFJochExceptionCategory::HDF5,
                                      "Legacy image source file missing");

            const auto &data_file = layout_.legacy_files.at(file_id);
            AppendOrExtendSourceMapping(ret, data_file.path, data_file.dataset,
                                        source_image, local_image, 1);
        }

        return ret;
    }

    throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                          "Unsupported HDF5 file layout for source mapping");
}