From 1275008864b2f01db28e84de19167d46da8ad74e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florian=20L=C3=BCke?= <f.lueke@mesytec.com>
Date: Sun, 29 Dec 2024 18:58:41 +0100
Subject: [PATCH] begin work on a root rntuple writer

---
 src/CMakeLists.txt                     |   4 +
 src/mana_plugin_root_rntuple_writer.cc | 126 +++++++++++++++++++++++++
 2 files changed, 130 insertions(+)
 create mode 100644 src/mana_plugin_root_rntuple_writer.cc

diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 4c084d7..ff19ec2 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -47,8 +47,12 @@ find_package(ROOT COMPONENTS Hist)
 if (ROOT_FOUND)
     message("-- Using ROOT installation from ${ROOT_USE_FILE}")
     include(${ROOT_USE_FILE})
+
     add_library(mana-plugin-root-histogram SHARED mana_plugin_root_histogram.cc)
     target_link_libraries(mana-plugin-root-histogram PRIVATE mana mesytec-mnode rxi-logc ${ROOT_LIBRARIES})
+
+    add_library(mana-plugin-root-rntuple-writer SHARED mana_plugin_root_rntuple_writer.cc)
+    target_link_libraries(mana-plugin-root-rntuple-writer PRIVATE mana mesytec-mnode rxi-logc ${ROOT_LIBRARIES})
 endif()
 
 if (pybind11_FOUND)
diff --git a/src/mana_plugin_root_rntuple_writer.cc b/src/mana_plugin_root_rntuple_writer.cc
new file mode 100644
index 0000000..6e56ace
--- /dev/null
+++ b/src/mana_plugin_root_rntuple_writer.cc
@@ -0,0 +1,126 @@
+#include <ROOT/RNTuple.hxx>
+#include <ROOT/RNTupleModel.hxx>
+#include <ROOT/RNTupleWriter.hxx>
+#include <ROOT/RVec.hxx>
+#include <TFile.h>
+#include <filesystem>
+#include <regex>
+#include <spdlog/spdlog.h>
+#include <vector>
+
+#include "internal/mana_lib.hpp"
+#include "internal/mana_root.hpp"
+#include "internal/rxi/log.hpp"
+
+using namespace mesytec::mnode;
+using namespace mesytec::mnode::mana;
+
+#if 0
+using RNTuple = ROOT::Experimental::RNTuple;
+#else
+using RNTuple = ROOT::RNTuple;
+#endif
+using RNTupleModel = ROOT::Experimental::RNTupleModel;
+using RNTupleWriter = ROOT::Experimental::RNTupleWriter;
+
+struct Context
+{
+    std::unique_ptr<TFile> outputFile;
+    std::unique_ptr<RNTupleModel> model;
+    std::vector<std::vector<std::shared_ptr<ROOT::RVecF>>> eventRVecs;
+};
+
+static Context *g_ctx = nullptr;
+
+MANA_DEFINE_PLUGIN_INIT(init)
+{
+    if (g_ctx)
+    {
+        log_warn("init() called multiple times. This plugin is a singleton!");
+        return g_ctx;
+    }
+    log_set_level(LOG_INFO);
+    log_debug("init");
+    return g_ctx = new Context;
+}
+
+MANA_DEFINE_PLUGIN_SHUTDOWN(shutdown)
+{
+    log_debug("shutdown");
+    if (context != g_ctx)
+    {
+        log_warn("shutdown() called with invalid context");
+        return;
+    }
+    delete g_ctx;
+    g_ctx = nullptr;
+}
+
+inline void make_things(const nlohmann::json &jRun)
+{
+    std::vector<std::vector<std::shared_ptr<ROOT::RVecF>>> result;
+    for (const auto &jEvent: jRun["events"])
+    {
+        auto dir = gDirectory->mkdir(jEvent["name"].get<std::string>().c_str(), "", true);
+        auto model = RNTupleModel::Create();
+        std::vector<std::shared_ptr<ROOT::RVecF>> rvecs;
+        for (const auto &jOutput: jEvent["outputs"])
+        {
+            auto rvec = model->MakeField<ROOT::RVecF>(jOutput["name"].get<std::string>());
+            rvec->resize(jOutput["size"].get<size_t>());
+            rvecs.emplace_back(std::move(rvec));
+        }
+        result.emplace_back(std::move(rvecs));
+    }
+}
+
+MANA_DEFINE_PLUGIN_BEGIN_RUN(begin_run)
+{
+    log_debug("begin_run: context=%p, descriptor_json=%s", context, descriptor_json);
+    auto jRun = nlohmann::json::parse(descriptor_json);
+    std::filesystem::path rp(jRun["name"].get<std::string>());
+    auto filename =
+        fmt::format("{}_rntuple_events.root", rp.filename().replace_extension().string());
+    auto ctx = reinterpret_cast<Context *>(context);
+    ctx->outputFile = std::make_unique<TFile>(filename.c_str(), "RECREATE");
+    ctx->outputFile->cd();
+
+    make_things(jRun);
+
+    ctx->model = ROOT::Experimental::RNTupleModel::Create();
+    auto field = ctx->model->MakeField<float>("value");
+
+    log_info("writing rntuple into: %s", filename.c_str());
+}
+
+MANA_DEFINE_PLUGIN_END_RUN(end_run)
+{
+    log_debug("end: context=%p, descriptor_json=%s", context, descriptor_json);
+    auto ctx = reinterpret_cast<Context *>(context);
+    ctx->outputFile->Write();
+    *ctx = {};
+}
+
+MANA_DEFINE_PLUGIN_EVENT_DATA(process_event)
+{
+    log_trace("event: ctx=%p, eventIndex=%d, arrayCount=%zu, totalBytes=%zu", context, eventIndex,
+              arrayCount, totalBytes);
+    auto ctx = reinterpret_cast<Context *>(context);
+}
+
+MANA_DEFINE_PLUGIN_SYSTEM_EVENT(process_system_event)
+{
+    log_trace("system_event: ctx=%p, size=%zu", context, size);
+}
+
+extern "C" MANA_C_SINK_PLUGIN()
+{
+    mana_sink_plugin_t plugin;
+    plugin.init = init;
+    plugin.shutdown = shutdown;
+    plugin.begin_run = begin_run;
+    plugin.end_run = end_run;
+    plugin.process_event = process_event;
+    plugin.process_system_event = process_system_event;
+    return plugin;
+}