Jungfraujoch/writer/HDF5Objects.h

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

#pragma once

#include <hdf5.h>
#include <memory>
#include <string>
#include <vector>
#include <mutex>
#include <optional>

#include "../common/JFJochException.h"
#include "../compression/CompressionAlgorithmEnum.h"
#include "../common/CompressedImage.h"

extern std::mutex hdf5_mutex;

class HDF5DataSet;

struct HDF5VirtualDatasetMapping {
    std::string filename;
    std::string dataset;

    std::vector<hsize_t> virtual_start;
    std::vector<hsize_t> virtual_stride;
    std::vector<hsize_t> virtual_count;
    std::vector<hsize_t> virtual_block;

    std::vector<hsize_t> source_start;
    std::vector<hsize_t> source_stride;
    std::vector<hsize_t> source_count;
    std::vector<hsize_t> source_block;

    bool ContainsVirtualImage(hsize_t image_number) const;
    hsize_t SourceImage(hsize_t image_number) const;
};

class HDF5Id {
protected:
    hid_t id;
    HDF5Id() = default;
    HDF5Id(HDF5Id &&other) noexcept;
    HDF5Id(const HDF5Id& other);
public:
    // Assignment would require closing ID's first and cannot be generalized easily
    HDF5Id& operator=(HDF5Id &&other) noexcept = delete;
    HDF5Id& operator=(const HDF5Id &other) noexcept = delete;
    hid_t GetID() const;
};

class HDF5DataSpace : public HDF5Id {
    uint8_t ndims = 1;
public:
    struct ScalarTag {};
    static constexpr ScalarTag Scalar{};

    explicit HDF5DataSpace(ScalarTag);

    explicit HDF5DataSpace(const std::vector<hsize_t>& dims = {1}, const std::vector<hsize_t> &max_dims = {});
    explicit HDF5DataSpace(const HDF5DataSet& data_set);
    uint8_t GetNumOfDimensions() const;
    std::vector<hsize_t> GetDimensions() const;
    void SelectHyperslab(const std::vector<hsize_t>& start, const std::vector<hsize_t>& size);
    void SelectHyperslabWithStride(const std::vector<hsize_t>& start, const std::vector<hsize_t>& size,
                                   const std::vector<hsize_t>& stride);
    ~HDF5DataSpace();
};

class HDF5DataType : public HDF5Id {
public:
    HDF5DataType(uint64_t size_in_bytes, bool is_signed);
    explicit HDF5DataType(double val);
    explicit HDF5DataType(float val);
    explicit HDF5DataType(bool val);
    explicit HDF5DataType(int32_t val);
    explicit HDF5DataType(uint32_t val);
    explicit HDF5DataType(int64_t val);
    explicit HDF5DataType(uint64_t val);
    explicit HDF5DataType(int16_t val);
    explicit HDF5DataType(uint16_t val);
    explicit HDF5DataType(int8_t val);
    explicit HDF5DataType(uint8_t val);
    explicit HDF5DataType(const std::string &val);
    explicit HDF5DataType(const char *val);
    explicit HDF5DataType(const HDF5DataSet& data_set);
    explicit HDF5DataType(CompressedImageMode mode);
    size_t GetElemSize() const;
    bool IsSigned() const;
    bool IsInteger() const;
    bool IsFloat() const;
    ~HDF5DataType();
};

enum class HDF5DataSetLayout {CONTIGUOUS, CHUNKED, VIRTUAL};

class HDF5Dcpl : public HDF5Id {
    uint8_t ndim = 0;
    HDF5DataSetLayout layout = HDF5DataSetLayout::CONTIGUOUS;
public:
    HDF5Dcpl();
    HDF5Dcpl(const HDF5DataSet& data_set);
    ~HDF5Dcpl();
    void SetFillValue32(int32_t val);
    void SetFillValue16(int16_t val);
    void SetFillValue8(int8_t val);
    void SetFillValueU32(uint32_t val);
    void SetFillValueU16(uint16_t val);
    void SetFillValueU8(uint8_t val);
    void SetChunking(const std::vector<hsize_t> & dims);
    std::vector<hsize_t> GetChunking();
    CompressionAlgorithm GetCompression();
    void SetCompression(CompressionAlgorithm algorithm, size_t block_size);
    void SetVirtual(const std::string& filename, const std::string& dataset, const HDF5DataSpace& src_dataspace,
                    const HDF5DataSpace& virtual_dataspace);
    HDF5DataSetLayout GetLayout() const;
    std::vector<HDF5VirtualDatasetMapping> GetVirtualMappings() const;
    uint8_t GetNumOfDimensions() const;
};

class HDF5Fapl : public HDF5Id {
public:
    HDF5Fapl();
    ~HDF5Fapl();
    void SetVersionTo1p10orNewer();
    void SetCloseStrong();
    void SetPoisonSec2();
};

class HDF5Object : public HDF5Id {
public:
    HDF5Object& Attr(const std::string& name, double val);
    HDF5Object& Attr(const std::string& name, const std::string& val);
    HDF5Object& Attr(const std::string& name, int32_t val);
    HDF5Object& Attr(const std::string& name, uint32_t val);
    HDF5Object& Attr(const std::string& name, int64_t val);
    HDF5Object& Attr(const std::string& name, uint64_t val);
    HDF5Object& Attr(const std::string& name, const std::vector<double> &val);

    std::vector<double> ReadAttrVec(const std::string& name);
    bool AttrExists(const std::string &name);
    std::string ReadAttrStr(const std::string &name);
    double ReadAttrDouble(const std::string &name);
    int64_t ReadAttrInt(const std::string &name);

    HDF5Object& Units(const std::string& val);
    HDF5Object& NXClass(const std::string& val);
    HDF5Object& Transformation(const std::string& units, const std::string& depends_on,
                               const std::string& equipment, const std::string& equipment_component,
                               const std::string& transformation_type, const std::vector<double> &vec);
    HDF5Object& Transformation(const std::string& units, const std::string& depends_on,
                               const std::string& equipment, const std::string& equipment_component,
                               const std::string& transformation_type, const std::vector<double>& vector,
                               const std::vector<double> &offset, const std::string& offset_units);
    HDF5Object& ExternalLink(std::string filename, const std::string &external_dataset, const std::string &local_name);
    HDF5Object& HardLink(const std::string &source, const std::string &dest);
    std::unique_ptr<HDF5DataSet> SaveScalar(const std::string &name, const char *val);
    std::unique_ptr<HDF5DataSet> SaveScalar(const std::string &name, const std::string& val);
    std::unique_ptr<HDF5DataSet> SaveVector(const std::string &name, const std::vector<std::string> &val);

    template <class T> std::unique_ptr<HDF5DataSet> SaveScalar(const std::string &name, T val);
    template <class T> std::unique_ptr<HDF5DataSet> SaveVector(const std::string &name, const std::vector<T> &val,
                                                               std::vector<hsize_t> dim = {}, CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION);

    bool GetBool(const std::string &name) const;
    int64_t GetInt(const std::string& name) const;
    std::optional<int64_t> GetOptInt(const std::string& name) const;
    std::optional<float> GetOptFloat(const std::string& name) const;
    std::optional<bool> GetOptBool(const std::string& name) const;

    float GetFloat(const std::string& name) const;
    std::string GetString(const std::string& name, const std::string& def = "") const;

    template <class T> std::optional<T> ReadElement(const std::string& name, size_t n) const;

    template <class T> std::vector<T> ReadVector(const std::string &name);
    template <class T> std::vector<T> ReadOptVector(const std::string &name);
    template <class T> std::vector<T> ReadVector(const std::string &name,
                                                 const std::vector<hsize_t>& start,
                                                 const std::vector<hsize_t>& size);
    template <class T> std::vector<T> ReadOptVector(const std::string &name,
                                                    const std::vector<hsize_t>& start,
                                                    const std::vector<hsize_t>& size);
    bool Exists(const std::string& name) const;
    bool IsExternalLink(const std::string& name) const;
    std::string GetLinkedFileName(const std::string& name) const;
    std::vector<std::string> FindLeafs(const std::string &name) const;
    std::vector<hsize_t> GetDimension(const std::string &name);
};

class HDF5Group : public HDF5Object {
public:
    HDF5Group(const HDF5Object& parent, const std::string& name);
    HDF5Group(const HDF5Object& parent, const char *name);
    ~HDF5Group();
};

class HDF5File : public HDF5Object {
public:
    explicit HDF5File(const std::string& filename, bool v1_10 = false);
    ~HDF5File();
    void Delete(const std::string& path);
    void Close();
};

class HDF5ReadOnlyFile : public HDF5Object {
public:
    explicit HDF5ReadOnlyFile(const std::string& filename);
    ~HDF5ReadOnlyFile();
};

class HDF5DataSet : public HDF5Object {
    uint8_t ndim;
    const std::string dataset_name;
public:
    HDF5DataSet(const HDF5Object& parent, const std::string& name); // Open existing dataset
    HDF5DataSet(const HDF5Object& parent, const std::string& name, const HDF5DataType &data_type,
                const HDF5DataSpace &data_space, const HDF5Dcpl &dcpl);
    HDF5DataSet(const HDF5Object &parent, const std::string &name, const HDF5DataType &data_type,
                const HDF5DataSpace &data_space);
    ~HDF5DataSet();
    HDF5DataSet& Write(const HDF5DataType &data_type, const void *val);

    template <class T>
    HDF5DataSet& WriteVec(const std::vector<T> &v,
                          const std::vector<hsize_t> &start,
                          const std::vector<hsize_t> &size) {
        HDF5DataSpace mem_space({v.size()});
        HDF5DataSpace file_space(*this);

        file_space.SelectHyperslab(start, size);
        if (H5Dwrite(id,
                     HDF5DataType(v[0]).GetID(),
                     mem_space.GetID(),
                     file_space.GetID(),
                     H5P_DEFAULT,
                     v.data()) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Vector dataset write unsuccessful");
        return *this;
    }

    template <class T>
    HDF5DataSet& WriteScalar(const T &val,
                             const std::vector<hsize_t> &start) {
        HDF5DataSpace mem_space({1});
        HDF5DataSpace file_space(*this);
        file_space.SelectHyperslab(start, {1});
        if (H5Dwrite(id,
                     HDF5DataType(val).GetID(),
                     mem_space.GetID(),
                     file_space.GetID(),
                     H5P_DEFAULT,
                     &val) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Vector dataset write unsuccessful");
        return *this;
    }

    HDF5DataSet& WriteDirectChunk(const void *val, hsize_t data_size, const std::vector<hsize_t>& offset);

    // Templated on the allocator so a caller can hand over a buffer that does not zero what it is
    // about to overwrite (RawByteBuffer).
    template<class Alloc> void ReadDirectChunk(std::vector<uint8_t, Alloc> &val, const std::vector<hsize_t>& offset) {
        if (offset.size() != ndim)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Inconsistent dimension settings");

        hsize_t chunk_bytes;
        if (H5Dget_chunk_storage_size(id, offset.data(), &chunk_bytes) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Error checking chunk size");

        val.resize(chunk_bytes);
        uint32_t filters;
        size_t buf_size = val.size();
        if (H5Dread_chunk(id, H5P_DEFAULT, offset.data(), &filters, val.data(), &buf_size))
            throw JFJochException(JFJochExceptionCategory::HDF5,
                                  "Error reading dataset (with direct chunks) to HDF5 file");
    }

    HDF5DataSet& Flush();
    void SetExtent(const std::vector<hsize_t>& dims);
    template<class T> T ReadScalar() const {
        HDF5DataSpace mem_space({1});
        HDF5DataSpace file_space(*this);
        if (file_space.GetNumOfDimensions() != 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Dataset tries to read scalar from vector dataset");;

        T output{};
        if (H5Dread(id, HDF5DataType(output).GetID(), mem_space.GetID(), H5S_ALL, H5P_DEFAULT, &output) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, dataset_name + ": read unsuccessful");
        return output;
    }

    template<class Alloc>
    void ReadVectorToU8(std::vector<uint8_t, Alloc> &v, const std::vector<hsize_t>& slab_start, const std::vector<hsize_t>& slab_size) const {
        HDF5DataType data_type(*this);
        HDF5DataSpace mem_space(slab_size);
        HDF5DataSpace file_space(*this);
        file_space.SelectHyperslab(slab_start, slab_size);

        size_t out_size = 1;
        out_size *= data_type.GetElemSize();
        for (unsigned long i : slab_size)
            out_size *= i;
        if (out_size == 0)
            throw JFJochException(JFJochExceptionCategory::HDF5,"Output size cannot be zero");

        v.resize(out_size);

        if (H5Dread(id, data_type.GetID(), mem_space.GetID(), file_space.GetID(), H5P_DEFAULT, v.data()) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Multidimensional vector dataset read unsuccessful");
    }

    template<class T> void ReadVector(std::vector<T> &v) const {
        HDF5DataSpace file_space(*this);

        if (file_space.GetNumOfDimensions() == 0) {
            // Handle trivial case of scalar
            v.resize(1);
            if (H5Dread(id, HDF5DataType(v[0]).GetID(), H5S_ALL, H5S_ALL, H5P_DEFAULT, v.data()) < 0)
                throw JFJochException(JFJochExceptionCategory::HDF5, dataset_name + ": read unsuccessful");
            return;
        }

        if (file_space.GetNumOfDimensions() != 1)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Dataset tries to read multi-dimensional value into 1D vector ");
        v.resize(file_space.GetDimensions()[0]);

        if (H5Dread(id, HDF5DataType(v[0]).GetID(), H5S_ALL, H5S_ALL, H5P_DEFAULT, v.data()) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "1D vector dataset read unsuccessful");
    }

    template<class T> void ReadVector(std::vector<T> &v, const std::vector<hsize_t>& start, const std::vector<hsize_t>& size) const {
        HDF5DataSpace mem_space({v.size()});
        HDF5DataSpace file_space(*this);
        file_space.SelectHyperslab(start, size);

        if (H5Dread(id, HDF5DataType(v[0]).GetID(), mem_space.GetID(), file_space.GetID(), H5P_DEFAULT, v.data()) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Multidimensional vector dataset read unsuccessful");
    }

    template<class T> void ReadVector(std::vector<T> &v, const std::vector<hsize_t>& start,
                                      const std::vector<hsize_t>& size, const std::vector<hsize_t>& stride) const {
        HDF5DataSpace mem_space({v.size()});
        HDF5DataSpace file_space(*this);
        file_space.SelectHyperslabWithStride(start, size, stride);

        if (H5Dread(id, HDF5DataType(v[0]).GetID(), mem_space.GetID(), file_space.GetID(), H5P_DEFAULT, v.data()) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "Multidimensional vector dataset read unsuccessful");
    }

    template<class T>
    std::optional<T> ReadElement(size_t n) const;

    std::string ReadString() const;
    void Close();
};

inline std::unique_ptr<HDF5DataSet> SaveScalar(const HDF5Object& parent, const std::string &name, const char* val) {
    HDF5DataType data_type(val);
    HDF5DataSpace data_space(HDF5DataSpace::Scalar);
    auto dataset = std::make_unique<HDF5DataSet>(parent, name, data_type, data_space);
    dataset->Write(data_type, val);
    return dataset;
}

inline std::unique_ptr<HDF5DataSet> SaveScalar(const HDF5Object& parent, const std::string &name, const std::string& val) {
    return SaveScalar(parent, name, val.c_str());
}

template <class T> std::unique_ptr<HDF5DataSet> SaveScalar(const HDF5Object& parent, const std::string &name, T val) {
    HDF5DataType data_type(val);
    HDF5DataSpace data_space(HDF5DataSpace::Scalar);
    auto dataset = std::make_unique<HDF5DataSet>(parent, name, data_type, data_space);
    dataset->Write(data_type, &val);
    return dataset;
}

template <class T> std::unique_ptr<HDF5DataSet> HDF5Object::SaveScalar(const std::string &name, T val) {
    HDF5DataType data_type(val);
    HDF5DataSpace data_space(HDF5DataSpace::Scalar);
    auto dataset = std::make_unique<HDF5DataSet>(*this, name, data_type, data_space);
    dataset->Write(data_type, &val);
    return dataset;
}

template <class T> std::vector<T> HDF5Object::ReadVector(const std::string &name) {
    std::vector<T> tmp;
    HDF5DataSet dataset(*this, name);
    dataset.ReadVector(tmp);
    return tmp;
}

template <class T> std::vector<T> HDF5Object::ReadOptVector(const std::string &name) {
    std::vector<T> tmp;
    if (Exists(name)) {
        HDF5DataSet dataset(*this, name);
        dataset.ReadVector(tmp);
    }
    return tmp;
}

template<class T>
std::vector<T> HDF5Object::ReadVector(const std::string &name,
                                      const std::vector<hsize_t> &start,
                                      const std::vector<hsize_t> &size) {
    std::vector<T> tmp;

    if (start.empty() || (start.size() != size.size()))
        throw JFJochException(JFJochExceptionCategory::HDF5, "Dimension error");
    size_t v_size = size[0];
    for (int i = 1; i < size.size(); i++)
        v_size *= size[i];
    tmp.resize(v_size);

    HDF5DataSet dataset(*this, name);
    dataset.ReadVector(tmp, start, size);

    return tmp;
}

template<class T>
std::vector<T> HDF5Object::ReadOptVector(const std::string &name,
                                         const std::vector<hsize_t> &start,
                                         const std::vector<hsize_t> &size) {
    std::vector<T> tmp;
    if (Exists(name))
        tmp = ReadVector<T>(name, start, size);
    return tmp;
}

template<class T>
std::optional<T> HDF5Object::ReadElement(const std::string &name, size_t n) const {
    if (Exists(name)) {
        const HDF5DataSet dataset(*this, name);
        return dataset.ReadElement<T>(n);
    }
    return std::nullopt;
}

template<class T>
std::optional<T> HDF5DataSet::ReadElement(size_t n) const {
    HDF5DataSpace file_space(*this);
    auto dims = file_space.GetDimensions();
    // Allow scalar (0-D) to be read with n==0 as a convenience.
    if (file_space.GetNumOfDimensions() == 0) {
        if (n != 0) return std::nullopt;
        return ReadScalar<T>();
    }
    if (file_space.GetNumOfDimensions() != 1)
        throw JFJochException(JFJochExceptionCategory::HDF5, "ReadElement requires a 1D dataset");
    if (n >= dims[0])
        return std::nullopt;

    // Select a single-element hyperslab at position n.
    file_space.SelectHyperslab({static_cast<hsize_t>(n)}, {1});
    HDF5DataSpace mem_space({1});
    T out{};
    if (H5Dread(id, HDF5DataType(out).GetID(), mem_space.GetID(), file_space.GetID(), H5P_DEFAULT, &out) < 0)
        throw JFJochException(JFJochExceptionCategory::HDF5, "ReadElement unsuccessful");
    return out;
}

template<>
inline std::optional<std::string> HDF5DataSet::ReadElement<std::string>(size_t n) const {
    HDF5DataSpace file_space(*this);
    auto dims = file_space.GetDimensions();

    if (file_space.GetNumOfDimensions() == 0) {
        if (n != 0) return std::nullopt;
        return ReadString();
    }

    if (file_space.GetNumOfDimensions() != 1)
        throw JFJochException(JFJochExceptionCategory::HDF5, "ReadElement<string> requires a 1D dataset");

    if (n >= dims[0])
        return std::nullopt;

    file_space.SelectHyperslab({static_cast<hsize_t>(n)}, {1});

    HDF5DataType file_type(*this);
    if (H5Tget_class(file_type.GetID()) != H5T_STRING)
        throw JFJochException(JFJochExceptionCategory::HDF5, "ReadElement<string>: dataset is not a string type");

    // Variable-length string
    if (H5Tis_variable_str(file_type.GetID()) > 0) {
        HDF5DataSpace mem_space({1});
        char *tmp = nullptr;

        if (H5Dread(id, file_type.GetID(), mem_space.GetID(), file_space.GetID(), H5P_DEFAULT, &tmp) < 0)
            throw JFJochException(JFJochExceptionCategory::HDF5, "ReadElement<string> unsuccessful");

        std::string out = (tmp != nullptr) ? std::string(tmp) : std::string{};
        if (tmp != nullptr)
            H5free_memory(tmp);
        return out;
    }

    // Fixed-length string
    const size_t elem_size = file_type.GetElemSize();
    if (elem_size == 0)
        return std::string{};

    std::vector<char> buffer(elem_size + 1, '\0');
    HDF5DataSpace mem_space({1});

    if (H5Dread(id, file_type.GetID(), mem_space.GetID(), file_space.GetID(), H5P_DEFAULT, buffer.data()) < 0)
        throw JFJochException(JFJochExceptionCategory::HDF5, "ReadElement<string> unsuccessful");

    return std::string(buffer.data());
}

template <class T>
std::unique_ptr<HDF5DataSet> SaveVector(const HDF5Object& parent, const std::string &name, const std::vector<T> &val,
                                                           std::vector<hsize_t> dims = {},
                                                           CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION) {
    if (dims.empty()) dims = {val.size()};

    if (val.empty())
        throw JFJochException(JFJochExceptionCategory::HDF5, "Cannot write empty vector");

    HDF5DataType data_type(val[0]);
    HDF5Dcpl dcpl;
    if (algorithm != CompressionAlgorithm::NO_COMPRESSION) {
        dcpl.SetCompression(algorithm, 0);
        dcpl.SetChunking(dims);
    }

    HDF5DataSpace data_space(dims);
    auto dataset = std::make_unique<HDF5DataSet>(parent, name, data_type, data_space, dcpl);
    dataset->Write(data_type, val.data());
    return dataset;
}

template <class T>
std::unique_ptr<HDF5DataSet> HDF5Object::SaveVector(const std::string &name, const std::vector<T> &val,
                                                                       std::vector<hsize_t> dims, CompressionAlgorithm algorithm) {
    if (val.empty())
        throw JFJochException(JFJochExceptionCategory::HDF5, "Cannot write empty vector");

    if (dims.empty()) dims = {val.size()};

    HDF5DataType data_type(val[0]);
    HDF5Dcpl dcpl;

    if (algorithm != CompressionAlgorithm::NO_COMPRESSION) {
        dcpl.SetCompression(algorithm, 0);
        dcpl.SetChunking(dims);
    }

    HDF5DataSpace data_space(dims);
    auto dataset = std::make_unique<HDF5DataSet>(*this, name, data_type, data_space, dcpl);
    dataset->Write(data_type, val.data());
    return dataset;
}

inline std::string hdf5_version() {
    unsigned majnum, minnum, relnum;
    H5get_libversion(&majnum, &minnum, &relnum);
    return "hdf5-" + std::to_string(majnum) + "." + std::to_string(minnum) + "." + std::to_string(relnum);
}

void RegisterHDF5Filter();
std::string ExtractFilename(const std::string& str);