mesytec-mnode/src/mana_plugin_root_rntuple_writer.cc

#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;
}