Jungfraujoch/writer/HDF5DataFile.cpp

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

#include <filesystem>
#include <iostream>

#include "HDF5DataFile.h"
#include "../compression/JFJochCompressor.h"

#include "HDF5DataFilePluginAzInt.h"
#include "HDF5DataFilePluginMX.h"
#include "HDF5DataFilePluginXFEL.h"
#include "HDF5DataFilePluginDetector.h"
#include "HDF5DataFilePluginROI.h"
#include "HDF5DataFilePluginPerformance.h"
#include "HDF5DataFilePluginImageStats.h"
#include "HDF5DataFilePluginReflection.h"
#include "../include/spdlog/fmt/fmt.h"
#include "HDF5NXmx.h"
#include "../common/time_utc.h"

HDF5DataFile::HDF5DataFile(const StartMessage &msg, uint64_t file_number, const std::string &filename) :
filename(filename),
file_number(file_number),
write_images(msg.write_images.value_or(true)),
declared_bit_depth(msg.bit_depth_image),
declared_pixel_signed(msg.pixel_signed) {
    if (msg.overwrite.has_value())
        overwrite = msg.overwrite.value();

    xpixel = 0;
    ypixel = 0;
    max_image_number = 0;
    nimages = 0;

    if (msg.file_format == FileWriterFormat::NXmxIntegrated) {
        image_low = 0;
        images_per_file = msg.number_of_images;
    } else {
        image_low = file_number * msg.images_per_file;
        images_per_file = msg.images_per_file;
    }

    timestamp.reserve(images_per_file);
    exptime.reserve(images_per_file);
    number.reserve(images_per_file);

    uint64_t tmp_suffix;
    try {
        if (!msg.arm_date.empty())
            tmp_suffix = parse_UTC_to_ms(msg.arm_date);
    } catch (...) {
        tmp_suffix = std::chrono::system_clock::now().time_since_epoch().count();
    }
    tmp_filename = fmt::format("{}.{:08x}.tmp", filename, tmp_suffix);
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginROI>());
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginDetector>(msg));
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginAzInt>(msg));
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginXFEL>());
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginMX>(msg));
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginImageStats>());
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginReflection>());
    plugins.emplace_back(std::make_unique<HDF5DataFilePluginPerformance>());
}

std::optional<HDF5DataFileStatistics> HDF5DataFile::Close() {
    if (!data_file)
        return {};

    HDF5Group group_exp(*data_file, "/entry/detector");
    group_exp.NXClass("NXcollection");

    group_exp.SaveVector("timestamp", timestamp);
    group_exp.SaveVector("exptime", exptime);
    group_exp.SaveVector("number", number);

    for (auto &p: plugins)
        p->WriteFinal(*data_file);

    if (data_set) {
        data_set->SetExtent({max_image_number + 1, ypixel, xpixel});
        data_set
                ->Attr("image_nr_low", (int32_t) (image_low + 1))
                .Attr("image_nr_high", (int32_t) (image_low + 1 + max_image_number));
        data_set->Close();
        data_set.reset();
    }

    if (manage_file ) {
        data_file->Close();
        data_file.reset();

        if (std::filesystem::exists(filename) && !overwrite)
            throw JFJochException(JFJochExceptionCategory::FileWriteError, "File already exists");
        std::error_code ec;
        std::filesystem::rename(tmp_filename, filename, ec);
        if (ec)
            throw JFJochException(JFJochExceptionCategory::FileWriteError,
                                  "Cannot rename temporary HDF5 file " + tmp_filename +
                                  " to " + filename + ": " + ec.message());
    } else {
        data_file.reset();
    }

    closed = true;

    HDF5DataFileStatistics ret;
    ret.max_image_number = max_image_number;
    ret.total_images = nimages;
    ret.filename = filename;
    ret.file_number = file_number + 1;

    return ret;
}

HDF5DataFile::~HDF5DataFile() {
    if (data_file) {
        try {
            data_set.reset();
            data_file.reset();
            if (manage_file) {
                std::error_code ec;
                std::filesystem::remove(tmp_filename, ec);
            }
        } catch (const std::exception &e) {
            std::cerr << "HDF5DataFile::~HDF5DataFile: " << e.what() << std::endl;
        } catch (...) {
            std::cerr << "HDF5DataFile::~HDF5DataFile: Unknown error " << std::endl;
        }
    }
}

void HDF5DataFile::CreateFile(const DataMessage& msg, std::shared_ptr<HDF5File> in_data_file) {
    data_file = in_data_file;

    HDF5Group(*data_file, "/entry").NXClass("NXentry");
    if (write_images) {
        HDF5Dcpl dcpl;

        // The only point where the two independent statements of the pixel format meet: the images
        // carry their own type (a CBOR tag per image, which is what the data files are written with)
        // while the master is typed from the START message. Nothing else compares them, so a stream
        // whose header contradicts its images produced data files of one type under a master
        // declaring another - and with NXmxVDS, HDF5 then converts silently on every read.
        if (msg.image.GetByteDepth() * 8 != declared_bit_depth
            || msg.image.IsSigned() != declared_pixel_signed)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "Image is " + std::to_string(msg.image.GetByteDepth() * 8)
                                  + "-bit " + (msg.image.IsSigned() ? "signed" : "unsigned")
                                  + ", but the start message declares "
                                  + std::to_string(declared_bit_depth) + "-bit "
                                  + (declared_pixel_signed ? "signed" : "unsigned")
                                  + " - the master file would not describe the data it links to");

        HDF5DataType data_type(msg.image.GetMode());

        xpixel = msg.image.GetWidth();
        ypixel = msg.image.GetHeight();

        dcpl.SetCompression(msg.image.GetCompressionAlgorithm(),
                            JFJochBitShuffleCompressor::BlockSize(msg.image.GetCompressionAlgorithm(),
                                                                  msg.image.GetByteDepth()));
        dcpl.SetChunking( {1, ypixel, xpixel});

        H5Pset_fill_time(dcpl.GetID(), H5D_FILL_TIME_NEVER);
        H5Pset_alloc_time(dcpl.GetID(), H5D_ALLOC_TIME_INCR);

        switch (msg.image.GetMode()) {
            case CompressedImageMode::Int8:
                dcpl.SetFillValue8(INT8_MIN);
                break;
            case CompressedImageMode::Int16:
                dcpl.SetFillValue16(INT16_MIN);
                break;
            case CompressedImageMode::Int32:
                dcpl.SetFillValue32(INT32_MIN);
                break;
            default:
                break;
        }

        HDF5Group(*data_file, "/entry/data").NXClass("NXdata");
        HDF5DataSpace data_space({1, ypixel, xpixel}, {H5S_UNLIMITED, ypixel, xpixel});
        data_set = std::make_unique<HDF5DataSet>(*data_file, "/entry/data/data", data_type, data_space, dcpl);
        data_set->SetExtent({images_per_file, ypixel, xpixel});
    }

    for (auto &p: plugins)
        p->OpenFile(*data_file, msg, images_per_file);
}

void HDF5DataFile::Write(const DataMessage &msg, uint64_t image_number) {
    if (closed)
        throw JFJochException(JFJochExceptionCategory::FileWriteError,
            "Trying to write to already closed file");
    if (image_number >= images_per_file)
        throw JFJochException(JFJochExceptionCategory::FileWriteError,
            "Image number out of bounds");

    if (!data_file) {
        manage_file = true;
        CreateFile(msg, std::make_shared<HDF5File>(tmp_filename));
    }

    if (new_file || (static_cast<int64_t>(image_number) > max_image_number)) {
        max_image_number = image_number;
        timestamp.resize(max_image_number + 1);
        exptime.resize(max_image_number + 1);
        number.resize(max_image_number + 1);
        new_file = false;
    }

    nimages++;
    if (data_set)
        data_set->WriteDirectChunk(msg.image.GetCompressed(), msg.image.GetCompressedSize(), {image_number, 0, 0});

    for (auto &p: plugins)
        p->Write(msg, image_number);

    timestamp[image_number] = msg.timestamp;
    exptime[image_number] = msg.exptime;
    number[image_number] = (msg.original_number) ? msg.original_number.value() : msg.number;
}

size_t HDF5DataFile::GetNumImages() const {
    return nimages;
}