From e784590a4e6409dd98b2348d1df93bc0cd761753 Mon Sep 17 00:00:00 2001
From: Andrej Babic <andrej.babic@psi.ch>
Date: Sun, 19 Apr 2020 19:33:36 +0200
Subject: [PATCH] Improved FastH5Writer

---
 core-buffer/include/FastH5Writer.hpp |  4 +-
 core-buffer/src/FastH5Writer.cpp     | 85 +++++++++++++++++++++++++---
 2 files changed, 81 insertions(+), 8 deletions(-)

diff --git a/core-buffer/include/FastH5Writer.hpp b/core-buffer/include/FastH5Writer.hpp
index c882057..4e7caad 100644
--- a/core-buffer/include/FastH5Writer.hpp
+++ b/core-buffer/include/FastH5Writer.hpp
@@ -30,9 +30,11 @@ class FastH5Writer {
 
     std::unordered_map<std::string, H5::PredType> scalar_metadata_;
 
-    void create_datasets();
+    void create_file(const std::string& filename);
     void close_file();
 
+    void flush_metadata();
+
 public:
     FastH5Writer(
             const size_t n_frames_per_file,
diff --git a/core-buffer/src/FastH5Writer.cpp b/core-buffer/src/FastH5Writer.cpp
index 0deb738..154b91d 100644
--- a/core-buffer/src/FastH5Writer.cpp
+++ b/core-buffer/src/FastH5Writer.cpp
@@ -14,10 +14,14 @@ using namespace std;
 FastH5Writer::FastH5Writer(
         const size_t n_frames_per_file,
         const uint16_t y_frame_size,
-        const uint16_t x_frame_size) :
+        const uint16_t x_frame_size,
+        const string& device_name,
+        const string& root_folder) :
             n_frames_per_file_(n_frames_per_file),
             y_frame_size_(y_frame_size),
             x_frame_size_(y_frame_size),
+            device_name_(device_name),
+            root_folder_(root_folder),
             frame_bytes_size_(2 * y_frame_size * y_frame_size),
             current_output_filename_(""),
             current_output_file_(),
@@ -29,8 +33,10 @@ FastH5Writer::FastH5Writer(
 //    auto file = H5::H5File(target_filename.c_str(), H5F_ACC_TRUNC);
 }
 
-void FastH5Writer::create_datasets()
+void FastH5Writer::create_file(const string& filename)
 {
+    auto file = H5::H5File(filename.c_str(), H5F_ACC_TRUNC);
+
     hsize_t dataset_dimension[3]  =
             {n_frames_per_file_, y_frame_size_, x_frame_size_};
     hsize_t max_dataset_dimension[3] =
@@ -43,17 +49,67 @@ void FastH5Writer::create_datasets()
     H5::DSetCreatPropList dataset_properties;
     dataset_properties.setChunk(3, dataset_chunking);
 
-    H5::AtomType dataset_data_type(H5::PredType::NATIVE_UINT16);
-    dataset_data_type.setOrder(H5T_ORDER_LE);
-
     current_image_dataset_ = current_output_file_.createDataSet(
-            "image", dataset_data_type, dataspace, dataset_properties);
+            "image",
+            H5::PredType::NATIVE_UINT16,
+            dataspace,
+            dataset_properties);
+
+    for (auto& metadata:scalar_metadata_) {
+        auto dataset_name = metadata.first;
+        auto dataset_type = metadata.second;
+
+        hsize_t dataset_dimension[2] = {n_frames_per_file_, 1};
+        H5::DataSpace dataspace(2, dataset_dimension);
+        auto dataset = current_output_file_.createDataSet(
+                dataset_name,
+                dataset_type,
+                dataspace);
+        datasets_.insert({dataset_name, dataset});
+
+        size_t n_buffer_bytes =
+                dataset.getDataType().getSize() * n_frames_per_file_;
+        buffers_.insert({dataset_name, make_shared<char[]>(n_buffer_bytes)});
+    }
+}
+
+void FastH5Writer::close_file()
+{
+    current_output_filename_ = "";
+    current_output_file_.close();
+    current_image_dataset_.close();
+    current_pulse_id_ = 0;
+    current_frame_index_ = 0;
+
+    for (auto& dataset:datasets_) {
+        dataset.second.close();
+    }
+    datasets_.clear();
+
+    buffers_.clear();
 }
 
 void FastH5Writer::set_pulse_id(const uint64_t pulse_id)
 {
     current_pulse_id_ = pulse_id;
     current_frame_index_ = BufferUtils::get_file_frame_index(pulse_id);
+
+    auto filename = BufferUtils::get_filename(
+            root_folder_, device_name_, pulse_id);
+
+    if (filename != current_output_filename_){
+        if (current_output_file_.getId() != -1) {
+            flush_metadata();
+            close_file();
+        }
+    }
+
+    create_file(filename);
+}
+
+void FastH5Writer::flush_metadata()
+{
+    // TODO: Actually flush this metadata.
 }
 
 void FastH5Writer::write_data(const char *buffer)
@@ -83,7 +139,22 @@ void FastH5Writer::write_data(const char *buffer)
 }
 
 template <>
-void FastH5Writer::add_metadata<uint64_t>(const std::string& metadata_name)
+void FastH5Writer::add_scalar_metadata<uint64_t>(
+        const std::string& metadata_name)
 {
+    scalar_metadata_.insert({metadata_name, H5::PredType::NATIVE_UINT64});
+}
 
+template <>
+void FastH5Writer::add_scalar_metadata<uint32_t>(
+        const std::string& metadata_name)
+{
+    scalar_metadata_.insert({metadata_name, H5::PredType::NATIVE_UINT32});
+}
+
+template <>
+void FastH5Writer::add_scalar_metadata<uint16_t>(
+        const std::string& metadata_name)
+{
+    scalar_metadata_.insert({metadata_name, H5::PredType::NATIVE_UINT16});
 }
\ No newline at end of file