NGnTree.cxx

#include <chrono>
#include <cstddef>
#include <ctime>
#include <numbers>
#include <set>
#include <string>
#include <vector>
#include "TAxis.h"
#include <TDirectory.h>
#include <TObject.h>
#include <TList.h>
#include <TROOT.h>
#include <THnSparse.h>
#include <TH1.h>
#include <TCanvas.h>
#include <TSystem.h>
#include <TMap.h>
#include <TObjString.h>
#include <TTree.h>
#include <TBufferJSON.h>
#include <sys/poll.h>
#include <zmq.h>
#include "NParameters.h"
#include "NStorageTree.h"
#include "NBinning.h"
#include "NBinningDef.h"
#include "NDimensionalExecutor.h"
#include "NDimensionalIpcRunner.h"
#include "NGnThreadData.h"
#include "NLogger.h"
#include "NTreeBranch.h"
#include "NUtils.h"
#include "NStorageTree.h"
#include "NGnNavigator.h"
#include "NGnTree.h"
 
ClassImp(Ndmspc::NGnTree);
 
namespace Ndmspc {
 
std::string NGnTree::BuildObjectPath(const json & cfg, const json & objCfg, const NBinningPoint * point)
{
  std::string objPath = "";
  if (objCfg.contains("prefix") && objCfg["prefix"].is_string()) {
    objPath = objCfg["prefix"].get<std::string>();
  }
 
  std::string axisObjectDefaultFormat =
      cfg["axisObjectDefaultFormat"].is_string() ? cfg["axisObjectDefaultFormat"].get<std::string>() : "%.2f_%.2f";
  std::string axisDefaultSeparator =
      cfg["axisDefaultSeparator"].is_string() ? cfg["axisDefaultSeparator"].get<std::string>() : "/";
 
  std::string lastSep;
  for (auto & axisEntry : cfg["axes"]) {
    std::string axisName;
    std::string mode;
    std::string format;
 
    if (axisEntry.is_string()) {
      axisName = axisEntry.get<std::string>();
      if (axisObjectDefaultFormat.empty()) {
        mode = "bin";
      }
      else {
        mode   = "minmax";
        format = axisObjectDefaultFormat;
      }
    }
    else if (axisEntry.is_object()) {
      if (axisEntry.contains("name") && axisEntry["name"].is_string()) {
        axisName = axisEntry["name"].get<std::string>();
      }
      else {
        continue;
      }
      if (axisEntry.contains("mode") && axisEntry["mode"].is_string()) {
        mode = axisEntry["mode"].get<std::string>();
      }
      if (axisEntry.contains("format") && axisEntry["format"].is_string()) {
        format = axisEntry["format"].get<std::string>();
      }
    }
    else {
      continue;
    }
 
    if (mode.empty()) {
      if (axisObjectDefaultFormat.empty())
        mode = "bin";
      else
        mode = "minmax";
    }
    if (format.empty()) {
      if (mode == "minmax")
        format = axisObjectDefaultFormat.empty() ? "%.2f_%.2f" : axisObjectDefaultFormat;
      else if (mode == "bin")
        format = "%d";
      else
        format = "%.2f";
    }
 
    if (mode == "minmax") {
      double min = point->GetBinMin(axisName);
      double max = point->GetBinMax(axisName);
      objPath += TString::Format(format.c_str(), min, max).Data();
    }
    else if (mode == "min") {
      double min = point->GetBinMin(axisName);
      objPath += TString::Format(format.c_str(), min).Data();
    }
    else if (mode == "max") {
      double max = point->GetBinMax(axisName);
      objPath += TString::Format(format.c_str(), max).Data();
    }
    else if (mode == "center") {
      double c = point->GetBinCenter(axisName);
      objPath += TString::Format(format.c_str(), c).Data();
    }
    else if (mode == "label") {
      std::string lbl = point->GetBinLabel(axisName);
      objPath += lbl;
    }
    else if (mode == "bin") {
      objPath += std::to_string(point->GetBin(axisName));
    }
    else {
      objPath += std::to_string(point->GetBin(axisName));
    }
 
    std::string sep = axisDefaultSeparator;
    if (axisEntry.is_object() && axisEntry.contains("sufix") && axisEntry["sufix"].is_string()) {
      sep = axisEntry["sufix"].get<std::string>();
    }
    objPath += sep;
    lastSep = sep;
  }
 
  if (!lastSep.empty() && objPath.size() >= lastSep.size()) {
    objPath = objPath.substr(0, objPath.size() - lastSep.size());
  }
  if (objCfg.contains("sufix") && objCfg["sufix"].is_string()) {
    objPath += objCfg["sufix"].get<std::string>();
  }
 
  return objPath;
}
NGnTree::NGnTree() : TObject() {}
 
NGnTree::NGnTree(std::vector<TAxis *> axes, std::string filename, std::string treename) : TObject(), fInput(nullptr)
{
  if (axes.empty()) {
    NLogError("NGnTree::NGnTree: No axes provided, binning is nullptr.");
    // fBinning = new NBinning();
    MakeZombie();
    return;
  }
  fBinning     = new NBinning(axes);
  fTreeStorage = new NStorageTree(fBinning);
  fTreeStorage->InitTree(filename, treename);
  fNavigator = new NGnNavigator();
  fNavigator->SetGnTree(this);
}
NGnTree::NGnTree(TObjArray * axes, std::string filename, std::string treename) : TObject(), fInput(nullptr)
{
 
  if (axes == nullptr) {
    NLogError("NGnTree::NGnTree: Axes TObjArray is nullptr.");
    MakeZombie();
    return;
  }
 
  if (axes == nullptr && axes->GetEntries() == 0) {
    NLogError("NGnTree::NGnTree: No axes provided, binning is nullptr.");
    MakeZombie();
    return;
  }
  fBinning     = new NBinning(axes);
  fTreeStorage = new NStorageTree(fBinning);
  fTreeStorage->InitTree(filename, treename);
  fNavigator = new NGnNavigator();
  fNavigator->SetGnTree(this);
}
 
NGnTree::NGnTree(NGnTree * ngnt, std::string filename, std::string treename) : TObject(), fInput(nullptr)
{
  if (ngnt == nullptr) {
    NLogError("NGnTree::NGnTree: NGnTree is nullptr.");
    MakeZombie();
    return;
  }
 
  if (ngnt->GetBinning() == nullptr) {
    NLogError("NGnTree::NGnTree: Binning in NGnTree is nullptr.");
    MakeZombie();
    return;
  }
 
  // TODO: Import binning from user
  fBinning     = (NBinning *)ngnt->GetBinning()->Clone();
  fTreeStorage = new NStorageTree(fBinning);
  fTreeStorage->InitTree(filename, treename);
  fNavigator = new NGnNavigator();
  fNavigator->SetGnTree(this);
}
 
NGnTree::NGnTree(NBinning * b, NStorageTree * s)
  : TObject(), fBinning(b), fTreeStorage(s), fInput(nullptr), fOwnsBinning(false), fOwnsTreeStorage(false)
{
  if (fBinning == nullptr) {
    NLogError("NGnTree::NGnTree: Binning is nullptr.");
    MakeZombie();
  }
  if (s == nullptr) {
    fTreeStorage = new NStorageTree(fBinning);
    fTreeStorage->InitTree("", "ngnt");
    fOwnsTreeStorage = true;
  }
 
  if (fTreeStorage == nullptr) {
    NLogError("NGnTree::NGnTree: Storage tree is nullptr.");
    MakeZombie();
  }
 
  // fBinning->Initialize();
  //
  // TODO: Check if this is needed
  fTreeStorage->SetBinning(fBinning);
  fBinning->GetPoint()->SetTreeStorage(fTreeStorage);
  fNavigator = new NGnNavigator();
  fNavigator->SetGnTree(this);
}
NGnTree::NGnTree(THnSparse * hns, std::string parameterAxis, const std::string & outFileName, json cfg)
    : TObject(), fInput(nullptr)
{
 
  std::map<std::string, std::vector<std::vector<int>>> b;
  TObjArray *                                          axes             = new TObjArray();
  int                                                  parameterAxisIdx = -1;
  std::vector<std::string>                             labels;
  for (int i = 0; i < hns->GetNdimensions(); i++) {
    TAxis * axisIn = (TAxis *)hns->GetAxis(i);
    TAxis * axis   = (TAxis *)axisIn->Clone();
 
    // check if parameterAxis matches axis name
    if (parameterAxis.compare(axis->GetName()) == 0) {
      parameterAxisIdx = i;
      TAxis * axis     = hns->GetAxis(parameterAxisIdx);
      for (int bin = 1; bin <= axis->GetNbins(); bin++) {
        // NLogInfo("Axis bin %d label: %s", bin, axis->GetBinLabel(bin));
        labels.push_back(axis->GetBinLabel(bin));
      }
      continue;
    }
 
    // set label
    if (axisIn->IsAlphanumeric()) {
      NLogTrace("Setting axis '%s' labels from input THnSparse", axis->GetName());
      for (int bin = 1; bin <= axisIn->GetNbins(); bin++) {
        std::string label = axisIn->GetBinLabel(bin);
        if (!labels.empty()) axis->SetBinLabel(bin, axisIn->GetBinLabel(bin));
      }
    }
 
    axes->Add(axis);
    b[axis->GetName()] = {{1}};
  }
 
  // json cfg;
  cfg["_parameterAxis"] = parameterAxisIdx;
  cfg["_labels"]        = labels;
 
  // return nullptr;
  NLogDebug("Importing THnSparse as NGnTree with parameter axis '%s' (index %d) ...", parameterAxis.c_str(),
            parameterAxisIdx);
  NGnTree * ngnt = new NGnTree(axes, outFileName);
  NLogDebug("Created NGnTree for THnSparse import ...");
  if (ngnt->IsZombie()) {
    NLogError("NGnTree::Import: Failed to create NGnTree !!!");
    MakeZombie();
    return;
  }
  // ngnt->GetStorageTree()->GetTree()->GetUserInfo()->Add(hns->Clone());
  // Get env variable for tmp dir $NDMSPC_TMP_DIR
  const char * tmpDirStr = gSystem->Getenv("NDMSPC_TMP_DIR");
 
  std::string tmpDir;
 
  if (!tmpDirStr || tmpDir.empty()) {
    tmpDir = "/tmp";
  }
  std::string tmpFilename = tmpDir + "/ngnt_imported_input" + std::to_string(gSystem->GetPid()) + ".root";
  NGnTree *   ngntIn      = new NGnTree(axes, tmpFilename);
  if (ngntIn->IsZombie()) {
    NLogError("NGnTree::Import: Failed to create NGnTree for input !!!");
    SafeDelete(ngnt);
    MakeZombie();
    return;
  }
  // ngntIn->GetStorageTree()->GetTree()->GetUserInfo()->Add(hns->Clone());
  ngntIn->GetOutput("default")->Add(hns->Clone("test"));
  ngntIn->Close(true);
 
  // return;
  // delete ngntIn;
 
  ngnt->SetInput(NGnTree::Open(tmpFilename)); // Set input to self
 
  ngnt->GetInput()->Print();
 
  ngnt->GetBinning()->AddBinningDefinition("default", b);
  ngnt->InitParameters(cfg["_labels"].get<std::vector<std::string>>());
 
  Ndmspc::NGnProcessFuncPtr processFunc = [](Ndmspc::NBinningPoint * point, TList * /*output*/, TList * outputPoint,
                                             int /*threadId*/) {
    // NLogInfo("Thread ID: %d", threadId);
    TH1::AddDirectory(kFALSE); // Prevent histograms from being associated with the current directory
    // point->Print();
    json cfg = point->GetCfg();
 
    NGnTree * ngntIn = point->GetInput();
    if (!ngntIn) {
      NLogError("NGnTree::Import: Input NGnTree is nullptr !!!");
      return;
    }
    // ngntIn->Print();
 
    THnSparse * hns = (THnSparse *)ngntIn->GetOutput("default")->At(0);
    if (hns == nullptr) {
      NLogError("NGnTree::Import: THnSparse 'hns' not found in storage tree !!!");
      return;
    }
 
    int                           axisIdx = cfg["_parameterAxis"].get<int>();
    std::vector<std::vector<int>> ranges;
    // set ranges from point storage coords
    int iAxis = 0;
    for (int i = 0; i < hns->GetNdimensions(); i++) {
      if (i == axisIdx) continue; // skip parameter axis
      int coord = point->GetStorageCoords()[iAxis++];
      ranges.push_back({i, coord, coord});
      // NLogInfo("Setting axis %d range to [%d, %d]", i, coord, coord);
    }
 
    NUtils::SetAxisRanges(hns, ranges);
    TH1 * h = hns->Projection(axisIdx, "O");
    if (!h) {
      NLogError("NGnTree::Import: Projection of THnSparse failed !!!");
      return;
    }
    if (h->GetEntries() > 0) {
      NParameters * params = point->GetParameters();
      if (params) {
        for (int bin = 1; bin <= h->GetNbinsX(); bin++) {
          params->SetParameter(bin, h->GetBinContent(bin), h->GetBinError(bin));
        }
      }
      // outputPoint->Add(hParams);
      outputPoint->Add(h);
 
      std::string filename = cfg["filename"].get<std::string>();
      TFile *     f        = (TFile *)point->GetTempObject("file");
      if (!f || filename.compare(f->GetName()) != 0) {
        if (f) {
          NLogDebug("NGnTree::Import: Closing previously opened file '%s' ...", f->GetName());
          f->Close();
        }
        NLogDebug("NGnTree::Import: Opening file '%s' ...", filename.c_str());
        f = TFile::Open(filename.c_str());
        if (!f || f->IsZombie()) {
          NLogError("NGnTree::Import: Cannot open file '%s' !!!", filename.c_str());
          return;
        }
        point->SetTempObject("file", f);
      }
 
      std::string axisObjectDefaultFormat =
          cfg["axisObjectDefaultFormat"].is_string() ? cfg["axisObjectDefaultFormat"].get<std::string>() : "%.2f_%.2f";
      std::string axisDefaultSeparator =
          cfg["axisDefaultSeparator"].is_string() ? cfg["axisDefaultSeparator"].get<std::string>() : "/";
      bool dryrun = false;
      if (cfg.contains("dryrun") && cfg["dryrun"].is_boolean()) {
        dryrun = cfg["dryrun"].get<bool>();
      }
 
      if (dryrun) {
        NLogInfo("NGnTree::Import (dryrun): '%s' ...", point->GetString().c_str());
      }
 
      // loop over all object in cfg["objects"]
      for (auto & [objName, objCfg] : cfg["objects"].items()) {
        std::string objPath = NGnTree::BuildObjectPath(cfg, objCfg, point);
 
        if (dryrun) {
          NLogInfo("NGnTree::Import (dryrun): would retrieve object '%s'", objPath.c_str());
          continue;
        }
 
        if (point->GetEntryNumber() == 0) {
          NLogInfo("NGnTree::Import: Retrieving object '%s' from file '%s' ...", objPath.c_str(),
                   cfg["filename"].get<std::string>().c_str());
        }
        TObject * obj = f->Get(objPath.c_str());
        if (!obj) {
          if (point->GetEntryNumber() == 0) {
            NLogWarning("NGnTree::Import: Cannot get object '%s' from file '%s' !!!", objPath.c_str(),
                        cfg["filename"].get<std::string>().c_str());
          }
          continue;
        }
 
        if (obj->InheritsFrom(TCanvas::Class())) {
          TCanvas * cObj = (TCanvas *)obj;
          cObj->SetName(objName.c_str());
        }
        outputPoint->Add(obj->Clone(objName.c_str()));
      }
      // f->Close();
    }
  };
 
  // NUtils::SetAxisRanges(, std::vector<std::vector<int>> ranges)
  ngnt->Process(processFunc, cfg);
  ngnt->Close(true);
  // Remove tmp file
  gSystem->Exec(TString::Format("rm -f %s", tmpFilename.c_str()));
}
 
NGnTree::~NGnTree()
{
 
  if (fOwnsBinning) {
    SafeDelete(fBinning);
  }
  if (fOwnsTreeStorage) {
    SafeDelete(fTreeStorage);
  }
  SafeDelete(fNavigator);
  SafeDelete(fParameters);
}
void NGnTree::Print(Option_t * option) const
{
 
  TString opt = option;
 
  // Print list of axes
  NLogInfo("NGnTree::Print: Printing NGnTree object [ALL] ...");
  if (fBinning) {
    fBinning->Print(option);
  }
  else {
    NLogError("Binning is not initialized in NGnTree !!!");
  }
  if (fTreeStorage) {
    fTreeStorage->Print(option);
  }
  else {
    NLogError("Storage tree is not initialized in NGnTree !!!");
  }
}
 
void NGnTree::Draw(Option_t * /*option*/)
{
 
  NLogInfo("NGnTree::Draw: Drawing NGnTree object [not implemented yet]...");
}
 
bool NGnTree::Process(NGnProcessFuncPtr func, const json & cfg, std::string binningName, NGnBeginFuncPtr beginFunc,
                      NGnEndFuncPtr endFunc)
{
 
  if (!fBinning) {
    NLogError("Binning is not initialized in NGnTree !!!");
    return false;
  }
 
  NBinning * binningIn = (NBinning *)fBinning->Clone();
 
  std::vector<std::string> defNames = fBinning->GetDefinitionNames();
  if (!binningName.empty()) {
    // Check if binning definitions exist
    if (std::find(defNames.begin(), defNames.end(), binningName) == defNames.end()) {
      NLogError("Binning definition '%s' not found in NGnTree !!!", binningName.c_str());
      return false;
    }
    defNames.clear();
    defNames.push_back(binningName);
  }
 
  fBinning->Reset();
  fBinning->SetCfg(cfg); // Set configuration to binning point
  bool rc = Process(func, defNames, cfg, binningIn, beginFunc, endFunc);
  if (!rc) {
    NLogError("NGnTree::Process: Processing failed !!!");
    return false;
  }
  // bool rc = false;
  return true;
}
 
bool NGnTree::Process(NGnProcessFuncPtr func, const std::vector<std::string> & defNames, const json & cfg,
                      NBinning * binningIn, NGnBeginFuncPtr beginFunc, NGnEndFuncPtr endFunc)
{
 
  NLogInfo("NGnTree::Process: Starting processing with %zu definitions ...", defNames.size());
  bool batch = gROOT->IsBatch();
  gROOT->SetBatch(kTRUE);
  TH1::AddDirectory(kFALSE);
 
  // --- Worker mode: run as a remote TCP worker ---
  if (const char * workerEndpoint = gSystem->Getenv("NDMSPC_WORKER_ENDPOINT")) {
    size_t workerIndex = 0;
    if (const char * envIdx = gSystem->Getenv("NDMSPC_WORKER_INDEX")) {
      try { workerIndex = static_cast<size_t>(std::stoul(envIdx)); } catch (...) {}
    }
    NLogInfo("NGnTree::Process: Worker mode — connecting to %s as worker %zu", workerEndpoint, workerIndex);
 
    const char * tmpDirEnv  = gSystem->Getenv("NDMSPC_TMP_DIR");
    std::string  workerBase = tmpDirEnv ? tmpDirEnv : "/tmp";
    // jobDir and treeName will be received from master via INIT
    // For now init NGnThreadData with a placeholder filename; Init() will be called with real paths
    Ndmspc::NGnThreadData workerData;
 
    void * ctx    = zmq_ctx_new();
    void * dealer = zmq_socket(ctx, ZMQ_DEALER);
    const std::string identity = Ndmspc::NDimensionalIpcRunner::BuildWorkerIdentity(workerIndex);
    zmq_setsockopt(dealer, ZMQ_IDENTITY, identity.data(), identity.size());
    int timeoutMs = 1000;
    zmq_setsockopt(dealer, ZMQ_RCVTIMEO, &timeoutMs, sizeof(timeoutMs));
    zmq_connect(dealer, workerEndpoint);
    Ndmspc::NDimensionalIpcRunner::SendFrames(dealer, {"READY"});
 
    // Wait for INIT
    const auto initDeadline = std::chrono::steady_clock::now() + std::chrono::seconds(30);
    bool initOk = false;
    while (!initOk) {
      std::vector<std::string> frames;
      if (!Ndmspc::NDimensionalIpcRunner::ReceiveFrames(dealer, frames)) {
        if (errno == EAGAIN || errno == EWOULDBLOCK) {
          if (std::chrono::steady_clock::now() > initDeadline) break;
          continue;
        }
        break;
      }
      // INIT frames: "INIT", workerIdx, sessionId, resultsDir, treeName[, tmpDir, tmpResultsDir]
      if (frames.size() >= 1 && frames[0] == "STOP") {
        NLogPrint("NGnTree::Process: Worker received STOP before INIT — session already finished, exiting.");
        zmq_close(dealer);
        zmq_ctx_term(ctx);
        gROOT->SetBatch(batch);
        return true;
      }
      if (frames.size() >= 5 && frames[0] == "INIT") {
        workerIndex                       = static_cast<size_t>(std::stoul(frames[1]));
        const std::string & sessionId     = frames[2];
        const std::string & initResultsDir = frames[3];
        const std::string & initTreeName  = frames[4];
 
        // Apply env vars sent by supervisor — these override the worker's inherited environment
        if (frames.size() >= 7) {
          if (!frames[5].empty())
            gSystem->Setenv("NDMSPC_TMP_DIR", frames[5].c_str());
          if (!frames[6].empty())
            gSystem->Setenv("NDMSPC_TMP_RESULTS_DIR", frames[6].c_str());
        }
        // Fallback: if NDMSPC_TMP_RESULTS_DIR is still unset/empty, use NDMSPC_TMP_DIR
        if (!gSystem->Getenv("NDMSPC_TMP_RESULTS_DIR") || gSystem->Getenv("NDMSPC_TMP_RESULTS_DIR")[0] == '\0') {
          const char * tmpDirEnv = gSystem->Getenv("NDMSPC_TMP_DIR");
          if (tmpDirEnv && tmpDirEnv[0] != '\0')
            gSystem->Setenv("NDMSPC_TMP_RESULTS_DIR", tmpDirEnv);
        }
 
        // Local work file — always on this machine's NDMSPC_TMP_DIR
        const char *      localTmpEnv = gSystem->Getenv("NDMSPC_TMP_DIR");
        const std::string localBase   = localTmpEnv ? localTmpEnv : "/tmp";
        const std::string localFile   = localBase + "/.ndmspc/tmp/" + sessionId + "/" +
                                        std::to_string(workerIndex) + "/" + fTreeStorage->GetPostfix();
 
        // Results file — on shared FS; supervisor reads from here to merge
        const std::string resultsFile = initResultsDir + "/" + std::to_string(workerIndex) + "/" +
                                        fTreeStorage->GetPostfix();
 
        bool rc = workerData.Init(workerIndex, func, beginFunc, endFunc, this, binningIn, fInput, localFile, initTreeName);
        if (!rc) {
          NLogError("NGnTree::Process: Worker failed to initialize NGnThreadData");
          zmq_close(dealer);
          zmq_ctx_term(ctx);
          return false;
        }
        // If results dir differs from local dir, tell TaskLoop to copy after Close(true)
        if (resultsFile != localFile) {
          workerData.SetResultsFilename(resultsFile);
        }
        workerData.SetCfg(cfg);
        Ndmspc::NDimensionalIpcRunner::SendFrames(dealer, {"ACK"});
        initOk = true;
      }
    }
    if (!initOk) {
      NLogError("NGnTree::Process: Worker did not receive INIT from supervisor");
      zmq_close(dealer);
      zmq_ctx_term(ctx);
      return false;
    }
 
    Ndmspc::NDimensionalIpcRunner::TaskLoop(dealer, workerIndex, &workerData);
    zmq_close(dealer);
    zmq_ctx_term(ctx);
    gROOT->SetBatch(batch);
    return true;
  }
  // --- End worker mode ---
 
  NUtils::EnableMT();
 
  int nThreads = ROOT::GetThreadPoolSize(); // Get the number of threads to use
  if (nThreads < 1) nThreads = 1;
 
  std::string executionMode = "thread";
  const char * envMode      = gSystem->Getenv("NDMSPC_EXECUTION_MODE");
  const bool   modeExplicit = (envMode && envMode[0] != '\0');
  if (modeExplicit) {
    executionMode = envMode;
  }
 
  std::string normalizedMode = executionMode;
  std::transform(normalizedMode.begin(), normalizedMode.end(), normalizedMode.begin(),
                 [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
  if (normalizedMode == "process") normalizedMode = "ipc";
 
  bool   useProcessIpc = (normalizedMode == "ipc" || normalizedMode == "tcp");
  bool   useTcp        = (normalizedMode == "tcp");
  size_t nProcesses    = static_cast<size_t>(nThreads);
  bool   ndmspcNProcExplicit = false;
 
  if (const char * envNdmspcNProc = gSystem->Getenv("NDMSPC_MAX_PROCESSES")) {
    ndmspcNProcExplicit = true;
    try {
      nProcesses = std::max<size_t>(1, static_cast<size_t>(std::stoll(envNdmspcNProc)));
    }
    catch (...) {
      NLogWarning("NGnTree::Process: Invalid NDMSPC_MAX_PROCESSES='%s', using default=%zu", envNdmspcNProc,
                  nProcesses);
    }
  }
  else if (const char * envNProc = gSystem->Getenv("ROOT_MAX_THREADS")) {
    // Backward-compatible fallback when NDMSPC_MAX_PROCESSES is not set.
    try {
      nProcesses = std::max<size_t>(1, static_cast<size_t>(std::stoll(envNProc)));
    }
    catch (...) {
      NLogWarning("NGnTree::Process: Invalid ROOT_MAX_THREADS='%s', using default=%zu", envNProc, nProcesses);
    }
  }
 
  // Keep explicit NDMSPC_EXECUTION_MODE settings authoritative.
  // If mode is not explicitly set, default to local IPC for multi-process runs.
  if (modeExplicit) {
    if (normalizedMode == "thread") {
      useProcessIpc = false;
      useTcp        = false;
    }
    else if (normalizedMode == "tcp") {
      useProcessIpc = true;
      useTcp        = true;
    }
    else if (normalizedMode == "ipc") {
      useProcessIpc = true;
      useTcp        = false;
    }
    else {
      NLogWarning("NGnTree::Process: Unknown NDMSPC_EXECUTION_MODE='%s', falling back to auto mode selection.",
                  executionMode.c_str());
      useProcessIpc = (nProcesses > 1);
      useTcp        = false;
    }
  }
  else if (nProcesses > 1) {
    useProcessIpc   = true;
    useTcp          = false;
    executionMode   = "ipc";
    normalizedMode  = "ipc";
  }
 
  if (ndmspcNProcExplicit && normalizedMode == "thread" && nProcesses > 1) {
    NLogWarning("NGnTree::Process: NDMSPC_MAX_PROCESSES=%zu is set, but NDMSPC_EXECUTION_MODE=thread disables IPC.",
                nProcesses);
  }
 
  const size_t workerObjectCount = useProcessIpc ? std::max(static_cast<size_t>(nThreads), nProcesses)
                                                 : static_cast<size_t>(nThreads);
  std::vector<Ndmspc::NGnThreadData> threadDataVector(workerObjectCount);
 
  NLogInfo("NGnTree::Process: executionMode='%s', useProcessIpc=%d, ROOT threads=%d, ipcProcesses=%zu, workerObjects=%zu",
           executionMode.c_str(), useProcessIpc ? 1 : 0, nThreads, nProcesses, workerObjectCount);
 
  const char * tmpDirEnv = gSystem->Getenv("NDMSPC_TMP_DIR");
  std::string  tmpDir;
  if (tmpDirEnv && tmpDirEnv[0] != '\0') {
    tmpDir = tmpDirEnv;
  } else {
    TString tmpDirPrefix = fTreeStorage->GetPrefix();
    // Use storage prefix only if it is a local path (not a remote URL)
    if (!(tmpDirPrefix.BeginsWith("root://") || tmpDirPrefix.BeginsWith("http://") ||
          tmpDirPrefix.BeginsWith("https://"))) {
      tmpDir = tmpDirPrefix.Data();
    }
    if (tmpDir.empty()) tmpDir = "/tmp";
  }
 
  std::string jobDir = tmpDir + "/.ndmspc/tmp/" + std::to_string(gSystem->GetPid());
 
  // Results dir: when NDMSPC_TMP_RESULTS_DIR equals NDMSPC_TMP_DIR (or is unset),
  // reuse jobDir so that localTmpFile == resultsFilename — no copy or delete needed.
  const char * resultsDirEnv = gSystem->Getenv("NDMSPC_TMP_RESULTS_DIR");
  const bool   sameDir       = !resultsDirEnv || std::string(resultsDirEnv) == tmpDir;
  std::string  resultsDir    = sameDir ? jobDir
                                       : (std::string(resultsDirEnv) + "/" +
                                          std::to_string(gSystem->GetPid()));
 
  std::string filePrefix = jobDir;
  for (size_t i = 0; i < threadDataVector.size(); ++i) {
    std::string filename = filePrefix + "/" + std::to_string(i) + "/" + fTreeStorage->GetPostfix();
    bool        rc       = threadDataVector[i].Init(i, func, beginFunc, endFunc, this, binningIn, fInput, filename,
                                                    fTreeStorage->GetTree()->GetName());
    if (!rc) {
      NLogError("Failed to initialize thread data %zu, exiting ...", i);
      return false;
    }
    threadDataVector[i].SetCfg(cfg); // Set configuration to binning point
    if (useTcp) {
      // Tell the merge step where workers will deposit their finished files.
      // When resultsDir == jobDir (NDMSPC_TMP_RESULTS_DIR unset) the paths are
      // identical so no copy or delete is needed — handled in TaskLoop.
      std::string resultsFile = resultsDir + "/" + std::to_string(i) + "/" + fTreeStorage->GetPostfix();
      threadDataVector[i].SetResultsFilename(resultsFile);
    }
  }
  size_t processedEntries = 0;
  size_t totalEntries     = 0;
  auto   start_par        = std::chrono::high_resolution_clock::now();
  auto   start_par_job    = std::chrono::high_resolution_clock::now();
  auto   task             = [&](const std::vector<int> & coords, Ndmspc::NGnThreadData & thread_obj) {
    // NLogWarning("Processing coordinates %s in thread %zu", NUtils::GetCoordsString(coords).c_str(),
    //                       thread_obj.GetAssignedIndex());
    // thread_obj.Print();
    thread_obj.Process(coords);
    processedEntries++;
    if (!NLogger::GetConsoleOutput()) {
      size_t nRunning = (totalEntries - processedEntries >= threadDataVector.size()) ? threadDataVector.size()
                                                                                     : totalEntries - processedEntries;
      NUtils::ProgressBar(processedEntries, totalEntries, start_par, TString::Format("R%4zu", nRunning).Data());
    }
  };
 
  size_t iDef   = 0;
  int    sumIds = 0;
 
  std::vector<Ndmspc::NThreadData *>            processWorkers;
  std::unique_ptr<Ndmspc::NDimensionalExecutor> ipcExecutor;
  if (useProcessIpc) {
    processWorkers.reserve(threadDataVector.size());
    for (size_t i = 0; i < threadDataVector.size(); ++i) {
      processWorkers.push_back(&threadDataVector[i]);
    }
    ipcExecutor = std::make_unique<Ndmspc::NDimensionalExecutor>(std::vector<int>{0}, std::vector<int>{0});
    if (useTcp) {
      const char * tcpPort    = gSystem->Getenv("NDMSPC_TCP_PORT");
      std::string  tcpEndpoint = std::string("tcp://0.0.0.0:") + (tcpPort ? tcpPort : "5555");
      const char * resultsDirBase = gSystem->Getenv("NDMSPC_TMP_RESULTS_DIR");
      // Auto-detect the macro to send to workers: explicit SetWorkerMacro() takes
      // priority; otherwise fall back to NDMSPC_MACRO set by ndmspc-run.
      std::string workerMacro = fWorkerMacroList;
      if (workerMacro.empty()) {
        if (const char * envMacro = gSystem->Getenv("NDMSPC_MACRO")) workerMacro = envMacro;
      }
      ipcExecutor->StartProcessIpc(processWorkers, nProcesses, tcpEndpoint, resultsDir,
                                   fTreeStorage->GetTree()->GetName(), workerMacro, tmpDir,
                                   resultsDirBase ? resultsDirBase : "");
    } else {
      ipcExecutor->StartProcessIpc(processWorkers, nProcesses);
    }
  }
 
  std::map<std::string, std::vector<Long64_t>>
      defIdMapProcessedRemoved; // Map to track which ids belong to which definition
  // for (auto & name : defNames) {
  //   auto binningDef = binningIn->GetDefinition(name);
  //   if (!binningDef) {
  //     NLogError("NGnTree::Process: Binning definition '%s' not found in NGnTree !!!", name.c_str());
  //     return false;
  //   }
  //   for (auto & id : binningDef->GetIds()) {
  //     defIdMap[name].push_back(id);
  //   }
  // }
 
  try {
    for (auto & name : defNames) {
    auto binningDef = binningIn->GetDefinition(name);
    if (!binningDef) {
      NLogError("NGnTree::Process: Binning definition '%s' not found in NGnTree !!!", name.c_str());
      return false;
    }
 
    if (binningDef->GetIds().size() == 0) {
      NLogWarning("NGnTree::Process: Binning definition '%s' has no entries, skipping ...", name.c_str());
      continue;
    }
 
    const std::vector<Long64_t> originalDefinitionIds = binningDef->GetIds();
 
    std::vector<int> mins, maxs;
    mins.push_back(0);
    maxs.push_back(binningDef->GetIds().size() - 1);
    NLogDebug("NGnTree::Process: Processing with binning definition '%s' with %zu entries", name.c_str(),
              binningDef->GetIds().size());
 
    if (NLogger::GetConsoleOutput()) {
      NLogInfo("NGnTree::Process: Processing binning definition '%s' with %d tasks ...", name.c_str(), maxs[0] + 1);
    }
    else {
      Printf("Processing binning definition '%s' with %d tasks ...", name.c_str(), maxs[0] + 1);
    }
    start_par        = std::chrono::high_resolution_clock::now();
    processedEntries = 0;
    totalEntries     = maxs[0] + 1;
    const size_t activeWorkers = useProcessIpc ? std::max<size_t>(1, std::min(nProcesses, processWorkers.size()))
                                               : threadDataVector.size();
    if (!NLogger::GetConsoleOutput())
      NUtils::ProgressBar(processedEntries, totalEntries, start_par,
                          TString::Format("R%4zu", activeWorkers).Data());
 
    // binningIn->SetCurrentDefinitionName(name);
    for (size_t i = 0; i < threadDataVector.size(); ++i) {
      // threadDataVector[i].SetCurrentDefinitionName(name);
      threadDataVector[i].GetHnSparseBase()->GetBinning()->SetCurrentDefinitionName(name);
    }
 
    Ndmspc::NDimensionalExecutor executorMT(mins, maxs);
 
    if (!useProcessIpc) {
      // Disable ROOT's RecursiveRemove during the parallel phase.
      // Without this, concurrent threads' object deletions trigger RecursiveRemove
      // which iterates pad->fPrimitives without per-object locks → TObjLink corruption.
      Bool_t prevMustClean = gROOT->MustClean();
      gROOT->SetMustClean(kFALSE);
 
      // Enable batch mode during the parallel phase so that ROOT drawing calls
      // (e.g. TH1::Fit drawing the fit function via f1->Draw) are no-ops.
      // Without this, concurrent threads each trigger ROOT canvas creation ("c1"),
      // the second canvas creation deletes the first → double-free / heap corruption,
      // which then causes TPad::RecursiveRemove to crash when closing the per-thread TTree.
      Bool_t prevBatch = gROOT->IsBatch();
      gROOT->SetBatch(kTRUE);
 
      executorMT.ExecuteParallel<Ndmspc::NGnThreadData>(task, threadDataVector);
 
      // Restore both flags before flushing deferred deletes, so each object's destructor
      // properly calls gROOT->RecursiveRemove and removes itself from ROOT's global lists.
      // This prevents dangling pointers that would crash the RecursiveRemove cascade
      // triggered by TTree::~TTree when the per-thread storage files are closed below.
      // It is safe here because all worker threads have already finished.
      gROOT->SetMustClean(prevMustClean);
      gROOT->SetBatch(prevBatch);
 
      // Flush deferred deletes single-threaded with MustClean and batch mode restored.
      for (size_t i = 0; i < threadDataVector.size(); ++i) {
        threadDataVector[i].FlushDeferredDeletes();
      }
 
      for (size_t i = 0; i < threadDataVector.size(); ++i) {
        threadDataVector[i].ExecuteEndFunction();
      }
    }
    else {
      ipcExecutor->SetBounds(mins, maxs);
      size_t acked     = ipcExecutor->ExecuteCurrentBoundsProcessIpc(
          name, &originalDefinitionIds, [&, activeWorkers](size_t ackCount, size_t activeWorkersNow) {
            processedEntries = ackCount;
            if (!NLogger::GetConsoleOutput()) {
              size_t nRunning = std::min(activeWorkersNow, activeWorkers);
              NUtils::ProgressBar(processedEntries, totalEntries, start_par,
                                  TString::Format("R%4zu", nRunning).Data());
            }
          });
      processedEntries = acked;
 
      // Child processes update their own worker-object copies. Rebuild parent-side
      // per-worker counters and processed-id vectors deterministically from task assignment.
      const size_t processesToUse = std::max<size_t>(1, std::min(nProcesses, processWorkers.size()));
      for (size_t i = 0; i < threadDataVector.size(); ++i) {
        auto * workerDef = threadDataVector[i].GetHnSparseBase()->GetBinning()->GetDefinition(name);
        if (workerDef) {
          workerDef->GetIds().clear();
        }
      }
 
      for (size_t taskIndex = 0; taskIndex < originalDefinitionIds.size(); ++taskIndex) {
        const size_t workerIndex = taskIndex % processesToUse;
        threadDataVector[workerIndex].SetNProcessed(threadDataVector[workerIndex].GetNProcessed() + 1);
 
        auto * workerDef = threadDataVector[workerIndex].GetHnSparseBase()->GetBinning()->GetDefinition(name);
        if (workerDef) {
          workerDef->GetIds().push_back(originalDefinitionIds[taskIndex]);
        }
      }
 
      if (!NLogger::GetConsoleOutput() && processedEntries < totalEntries) {
        NUtils::ProgressBar(processedEntries, totalEntries, start_par, "R   0");
      }
    }
 
    if (!NLogger::GetConsoleOutput())
      Printf("Finished processing binning definition '%s'. Post-processing results ...", name.c_str());
    // Update hnsbBinningIn with the processed ids
    NLogDebug("NGnTree::Process: [BEGIN] ------------------------------------------------");
    sumIds += binningIn->GetDefinition(name)->GetIds().size();
    binningIn->GetDefinition(name)->GetIds().clear();
    for (size_t i = 0; i < threadDataVector.size(); ++i) {
      NLogDebug("NGnTree::Process: -> Thread %zu processed %lld entries", i, threadDataVector[i].GetNProcessed());
      // threadDataVector[i].GetHnSparseBase()->GetBinning()->GetDefinition(name)->Print();
      binningIn->GetDefinition(name)->GetIds().insert(
          binningIn->GetDefinition(name)->GetIds().end(),
          threadDataVector[i].GetHnSparseBase()->GetBinning()->GetDefinition(name)->GetIds().begin(),
          threadDataVector[i].GetHnSparseBase()->GetBinning()->GetDefinition(name)->GetIds().end());
      sort(binningIn->GetDefinition(name)->GetIds().begin(), binningIn->GetDefinition(name)->GetIds().end());
    }
    // hnsbBinningIn->GetDefinition(name)->Print();
    // remove entries present in hnsbBinningIn from other definitions
    for (size_t i = 0; i < defNames.size(); i++) {
 
      std::string other_name = defNames[i];
      auto        otherDef   = binningIn->GetDefinition(other_name);
      if (i <= iDef) {
        continue;
      }
      if (!otherDef) {
        NLogError("NGnTree::Process: Binning definition '%s' not found in NGnTree !!!", other_name.c_str());
        return false;
      }
      // remove entries that has value less then sumIds
      for (auto it = otherDef->GetIds().begin(); it != otherDef->GetIds().end();) {
        NLogTrace("NGnTree::Process: Checking entry %lld from definition '%s' against sumIds=%d", *it,
                  other_name.c_str(), sumIds);
        if (*it < sumIds) {
          NLogTrace("NGnTree::Process: Removing entry %lld from definition '%s'", *it, other_name.c_str());
          defIdMapProcessedRemoved[other_name].push_back(*it);
          it = otherDef->GetIds().erase(it);
        }
        else {
          ++it;
        }
      }
 
      binningIn->GetDefinition(other_name)->Print();
    }
    // hnsbBinningIn->GetDefinition(name)->Print();
    iDef++;
 
      NLogDebug("NGnTree::Process: [END] ------------------------------------------------");
    }
  }
  catch (const std::exception & ex) {
    if (ipcExecutor) {
      ipcExecutor->FinishProcessIpc(/*abort=*/true);
      ipcExecutor.reset();
    }
 
    TString what(ex.what());
    if (what.Contains("Interrupted by user")) {
      if (gROOT) {
        gROOT->SetInterrupt(kFALSE);
      }
      NLogWarning("NGnTree::Process: Interrupted by user, stopping processing.");
    }
    else {
      NLogError("NGnTree::Process: Processing failed: %s", ex.what());
    }
    return false;
  }
 
  // return true; // For testing, skip merging and post-processing
 
  auto                                      end_par      = std::chrono::high_resolution_clock::now();
  std::chrono::duration<double, std::milli> par_duration = end_par - start_par_job;
 
  if (ipcExecutor) {
    ipcExecutor->FinishProcessIpc();
  }
 
  // For TCP mode, only merge results from workers that actually connected.
  // For IPC/fork and thread modes, all indices are valid.
  const std::set<size_t> registeredWorkers =
    (ipcExecutor && useTcp) ? ipcExecutor->GetRegisteredWorkerIndices() : std::set<size_t>{};
  const bool filterByRegistered = !registeredWorkers.empty();
 
  if (!NLogger::GetConsoleOutput()) {
    Printf("NGnTree::Process: Execution completed and it took %s .",
           NUtils::FormatTime(par_duration.count() / 1000).c_str());
  }
  else {
    NLogInfo("NGnTree::Process: Execution completed and it took %s .",
             NUtils::FormatTime(par_duration.count() / 1000).c_str());
  }
 
  NLogInfo("NGnTree::Process: Post processing %zu results ...", threadDataVector.size());
  for (auto & data : threadDataVector) {
    if (useProcessIpc) {
      NLogTrace("NGnTree::Process: Releasing parent handle for worker %zu file without writing",
                data.GetAssignedIndex());
      // data.GetHnSparseBase()->GetStorageTree()->Close(false);
    }
    else {
      NLogTrace("NGnTree::Process: Closing file from thread %zu with write", data.GetAssignedIndex());
      data.GetHnSparseBase()->GetStorageTree()->Close(true);
    }
  }
 
  NLogDebug("NGnTree::Process: Merging %zu results ...", threadDataVector.size());
  if (!NLogger::GetConsoleOutput()) {
    Printf("NGnTree::Process: [phase] merge start (%zu workers)", threadDataVector.size());
  }
  const auto mergeStart = std::chrono::high_resolution_clock::now();
  TList *                 mergeList  = new TList();
  Ndmspc::NGnThreadData * outputData = new Ndmspc::NGnThreadData();
  outputData->Init(0, func, nullptr, nullptr, this, binningIn);
  outputData->SetCfg(cfg);
  // outputData->Init(0, func, this);
 
  for (auto & data : threadDataVector) {
    if (filterByRegistered && registeredWorkers.find(data.GetAssignedIndex()) == registeredWorkers.end()) {
      NLogInfo("NGnTree::Process: Skipping worker %zu — never connected", data.GetAssignedIndex());
      continue;
    }
    NLogTrace("NGnTree::Process: Adding thread data %zu to merge list ...", data.GetAssignedIndex());
    mergeList->Add(&data);
  }
 
  Long64_t nmerged = outputData->Merge(mergeList);
  const auto mergeEnd = std::chrono::high_resolution_clock::now();
  if (!NLogger::GetConsoleOutput()) {
    const auto mergeSec = std::chrono::duration_cast<std::chrono::duration<double>>(mergeEnd - mergeStart).count();
    Printf("NGnTree::Process: [phase] merge done (%lld outputs, %.2f s)", nmerged, mergeSec);
  }
  if (nmerged <= 0) {
    NLogError("NGnTree::Process: Failed to merge thread data, exiting ...");
    delete mergeList;
    return false;
  }
  NLogInfo("NGnTree::Process: Merged %lld outputs successfully", nmerged);
  // delete all temporary files
  // for (auto & data : threadDataVector) {
  //   std::string filename = data.GetHnSparseBase()->GetStorageTree()->GetFileName();
  //   NLogTrace("NGnTree::Process: Deleting temporary file '%s' ...", filename.c_str());
  //   gSystem->Exec(TString::Format("rm -f %s", filename.c_str()));
  // }
  //
 
  // binningIn= outputData->GetHnSparseBase()->GetBinning();
 
  auto * mergedBinning = outputData->GetHnSparseBase()->GetBinning();
  std::set<Long64_t> mergedContentIds;
  std::vector<std::pair<Long64_t, std::vector<int>>> mergedContentCoords;
 
  // add missing entries to definitions based on defIdMapProcessedRemoved
  for (size_t i = 0; i < defNames.size(); i++) {
    std::string name = defNames[i];
    // auto        def  = binningIn->GetDefinition(name);
    auto def = mergedBinning->GetDefinition(name);
    if (!def) {
      NLogError("NGnTree::Process: Binning definition '%s' not found in NGnTree !!!", name.c_str());
      return false;
    }
    for (auto & [other_name, removedIds] : defIdMapProcessedRemoved) {
      if (other_name.compare(name) != 0) {
        continue;
      }
      for (auto & id : removedIds) {
        if (std::find(def->GetIds().begin(), def->GetIds().end(), id) == def->GetIds().end()) {
          NLogTrace("NGnTree::Process: Adding missing entry %lld to definition '%s'", id, name.c_str());
          def->GetIds().push_back(id);
        }
      }
    }
    sort(def->GetIds().begin(), def->GetIds().end());
    // outputData->GetHnSparseBase()->GetBinning()->GetDefinition(name)->GetIds() = def->GetIds();
 
    // Modify content in binning definitions based on def->GetIds()
    def->GetContent()->Reset();
    for (auto id : def->GetIds()) {
      NBinningPoint point(def->GetBinning());
      def->GetBinning()->GetContent()->GetBinContent(id, point.GetCoords());
      point.RecalculateStorageCoords();
      Long64_t bin = def->GetContent()->GetBin(point.GetStorageCoords());
      NLogTrace("NGnThreadData::Merge: [%s] Adding def_id=%lld to content_bin=%lld", name.c_str(), id, bin);
      def->GetContent()->SetBinContent(bin, id);
 
      if (mergedContentIds.insert(id).second) {
        mergedContentCoords.emplace_back(
            id, NUtils::ArrayToVector(point.GetCoords(), point.GetNDimensionsContent()));
      }
    }
  }
 
  // Rebuild the merged top-level content from final definition ids. In IPC/TCP mode
  // the merge setup may still carry sparse source-bin content; resetting here keeps
  // only the bins that correspond to actual merged tree entries.
  mergedBinning->GetContent()->Reset();
  for (const auto & entry : mergedContentCoords) {
    Long64_t bin = mergedBinning->GetContent()->GetBin(entry.second.data());
    mergedBinning->GetContent()->SetBinContent(bin, entry.first);
  }
 
  // print final binning definitions
  NLogDebug("NGnTree::Process: Final binning definitions after processing:");
  for (auto & name : defNames) {
    // auto binningDef = binningIn->GetDefinition(name);
    auto binningDef = outputData->GetHnSparseBase()->GetBinning()->GetDefinition(name);
    if (!binningDef) {
      NLogError("NGnTree::Process: Binning definition '%s' not found in NGnTree !!!", name.c_str());
      return false;
    }
    binningDef->Print();
  }
 
  fTreeStorage = outputData->GetHnSparseBase()->GetStorageTree();
  fOutputs     = outputData->GetHnSparseBase()->GetOutputs();
  fBinning     = outputData->GetHnSparseBase()->GetBinning(); // Update binning to the merged one
  fParameters  = outputData->GetHnSparseBase()->GetParameters();
 
  if (NLogger::GetConsoleOutput()) {
    NLogInfo("NGnTree::Process: Processing completed successfully. Output was stored in '%s'.",
             fTreeStorage->GetFileName().c_str());
  }
  else {
    Printf("Processing completed successfully. Output was stored in '%s'.", fTreeStorage->GetFileName().c_str());
  }
 
  // Close the final output file
  if (!NLogger::GetConsoleOutput()) {
    Printf("NGnTree::Process: [phase] final close start (%s)",
           outputData->GetHnSparseBase()->GetStorageTree()->GetFileName().c_str());
  }
  const auto closeStart = std::chrono::high_resolution_clock::now();
  outputData->GetHnSparseBase()->Close(true);
  const auto closeEnd = std::chrono::high_resolution_clock::now();
  if (!NLogger::GetConsoleOutput()) {
    const auto closeSec = std::chrono::duration_cast<std::chrono::duration<double>>(closeEnd - closeStart).count();
    Printf("NGnTree::Process: [phase] final close done (%.2f s)", closeSec);
  }
 
  if (!NLogger::GetConsoleOutput()) {
    Printf("NGnTree::Process: [phase] cleanup start (%s)", jobDir.c_str());
  }
  const auto cleanupStart = std::chrono::high_resolution_clock::now();
  gSystem->Exec(TString::Format("rm -fr %s", jobDir.c_str()));
  const auto cleanupEnd = std::chrono::high_resolution_clock::now();
  if (!NLogger::GetConsoleOutput()) {
    const auto cleanupSec = std::chrono::duration_cast<std::chrono::duration<double>>(cleanupEnd - cleanupStart).count();
    Printf("NGnTree::Process: [phase] cleanup done (%.2f s)", cleanupSec);
  }
  gROOT->SetBatch(batch); // Restore ROOT batch mode
  return true;
}
 
TList * NGnTree::GetOutput(std::string name)
{
 
  if (name.empty()) {
    name = fBinning->GetCurrentDefinitionName();
  }
  if (fOutputs.find(name) == fOutputs.end()) {
    fOutputs[name] = new TList();
    fOutputs[name]->SetName(name.c_str());
  }
  return fOutputs[name];
}
 
NGnTree * NGnTree::Open(const std::string & filename, const std::string & branches, const std::string & treename)
{
 
  NLogDebug("Opening '%s' with branches='%s' and treename='%s' ...", filename.c_str(), branches.c_str(),
            treename.c_str());
 
  TFile * file = TFile::Open(filename.c_str());
  if (!file) {
    NLogError("NGnTree::Open: Cannot open file '%s'", filename.c_str());
    return nullptr;
  }
 
  TTree * tree = (TTree *)file->Get(treename.c_str());
  if (!tree) {
    NLogError("NGnTree::Open: Cannot get tree '%s' from file '%s'", treename.c_str(), filename.c_str());
    return nullptr;
  }
 
  return Open(tree, branches, file);
}
 
NGnTree * NGnTree::Open(TTree * tree, const std::string & branches, TFile * file)
{
 
  NBinning * hnstBinning = (NBinning *)tree->GetUserInfo()->At(0);
  if (!hnstBinning) {
    NLogError("NGnTree::Open: Cannot get binning from tree '%s'", tree->GetName());
    return nullptr;
  }
  NStorageTree * hnstStorageTree = (NStorageTree *)tree->GetUserInfo()->At(1);
  if (!hnstStorageTree) {
    NLogError("NGnTree::Open: Cannot get tree storage info from tree '%s'", tree->GetName());
    return nullptr;
  }
 
  std::map<std::string, TList *> outputs;
  TDirectory *                   dir = nullptr;
  if (file) {
    dir    = (TDirectory *)file->Get("outputs");
    auto l = dir->GetListOfKeys();
    for (auto kv : *l) {
      TObject * obj = dir->Get(kv->GetName());
      if (!obj) continue;
      TList * l = dynamic_cast<TList *>(obj);
      if (!l) continue;
      outputs[l->GetName()] = l;
      NLogDebug("Imported output list for binning '%s' with %d object(s) from file '%s'", l->GetName(), l->GetEntries(),
                file->GetName());
    }
  }
  // TDirectory * dir = (TDirectory *)tree->GetUserInfo()->FindObject("outputs");
  // if (dir) {
  //   dir->Print();
  // }
 
  NGnTree * ngnt = new NGnTree(hnstBinning, hnstStorageTree);
 
  if (!hnstStorageTree->SetFileTree(file, tree)) return nullptr;
  // if (!ngnt->InitBinnings({})) return nullptr;
  // ngnt->Print();
  // Get list of branches
  std::vector<std::string> enabledBranches;
  if (!branches.empty()) {
    enabledBranches = Ndmspc::NUtils::Tokenize(branches, ',');
    NLogTrace("NGnTree::Open: Enabled branches: %s", NUtils::GetCoordsString(enabledBranches, -1).c_str());
    hnstStorageTree->SetEnabledBranches(enabledBranches);
  }
  else {
    // loop over all branches and set address
    for (auto & kv : hnstStorageTree->GetBranchesMap()) {
      NLogTrace("NGnTree::Open: Enabled branches: %s", kv.first.c_str());
    }
  }
  // Set all branches to be read
  hnstStorageTree->SetBranchAddresses();
  ngnt->SetOutputs(outputs);
 
  NGnNavigator * nav = new NGnNavigator();
  nav->SetGnTree(ngnt);
  ngnt->SetNavigator(nav);
 
  return ngnt;
}
 
void NGnTree::SetNavigator(NGnNavigator * navigator)
{
 
  if (fNavigator) {
    NLogTrace("NGnTree::SetNavigator: Replacing existing navigator ...");
    SafeDelete(fNavigator);
  }
 
  fNavigator = navigator;
}
 
bool NGnTree::Close(bool write)
{
 
  if (!fTreeStorage) {
    NLogError("NGnTree::Close: Storage tree is not initialized in NGnTree !!!");
    return false;
  }
 
  return fTreeStorage->Close(write, fOutputs);
}
 
Int_t NGnTree::GetEntry(Long64_t entry, bool checkBinningDef)
{
  if (!fTreeStorage) {
    NLogError("NGnTree::GetEntry: Storage tree is not initialized in NGnTree !!!");
    return -1;
  }
 
  int bytes =
      fTreeStorage->GetEntry(entry, fBinning->GetPoint(0, fBinning->GetCurrentDefinitionName()), checkBinningDef);
  if (fTreeStorage->GetBranch("_params")) fParameters = (NParameters *)fTreeStorage->GetBranch("_params")->GetObject();
  return bytes;
}
Int_t NGnTree::GetEntry(std::vector<std::vector<int>> /*range*/, bool checkBinningDef)
{
 
  return GetEntry(0, checkBinningDef);
}
 
void NGnTree::Play(int timeout, std::string binning, std::vector<int> outputPointIds,
                   std::vector<std::vector<int>> ranges, Option_t * option)
{
  TString opt = option;
  opt.ToUpper();
 
  std::string annimationTempDir =
      TString::Format("%s/.ndmspc/animation/%d", gSystem->Getenv("HOME"), gSystem->GetPid()).Data();
  gSystem->Exec(TString::Format("mkdir -p %s", annimationTempDir.c_str()));
 
  if (binning.empty()) {
    binning = fBinning->GetCurrentDefinitionName();
  }
 
  NBinningDef * binningDef = fBinning->GetDefinition(binning);
  if (!binningDef) {
    NLogError("NGnTree::Play: Binning definition '%s' not found in NGnTree !!!", binning.c_str());
    NLogError("Available binning definitions:");
    for (auto & name : fBinning->GetDefinitionNames()) {
      if (name == fBinning->GetCurrentDefinitionName())
        NLogError(" [*] %s", name.c_str());
      else
        NLogError(" [ ] %s", name.c_str());
    }
    return;
  }
 
  THnSparse * bdContent = (THnSparse *)binningDef->GetContent()->Clone();
 
  std::string bdContentName = TString::Format("bdContent_%s", binning.c_str()).Data();
  // Set axis ranges if provided
  if (!ranges.empty()) NUtils::SetAxisRanges(bdContent, ranges, false, true);
 
  Long64_t                                        linBin = 0;
  std::unique_ptr<ROOT::Internal::THnBaseBinIter> iter{bdContent->CreateIter(true /*use axis range*/)};
  std::vector<Long64_t>                           ids;
  // std::vector<Long64_t> ids = binningDef->GetIds();
  while ((linBin = iter->Next()) >= 0) {
    // NLogDebug("Original content bin %lld: %f", linBin, bdContentOrig->GetBinContent(linBin));
    ids.push_back(linBin);
  }
  if (ids.empty()) {
    NLogWarning("NGnTree::Play: No entries found in binning definition '%s' !!!", binning.c_str());
    return;
  }
  // return;
 
  TCanvas * c1 = nullptr;
 
  c1 = (TCanvas *)gROOT->GetListOfCanvases()->FindObject("c1");
  if (c1 == nullptr) c1 = new TCanvas("c1", "NGnTree::Play", 800, 600);
  c1->Clear();
  c1->cd();
  c1->DivideSquare(outputPointIds.size() > 0 ? outputPointIds.size() + 1 : 1);
  gSystem->ProcessEvents();
 
  binningDef->Print();
  bdContent->Reset();
 
  // loop over all ids and print them
  for (auto id : ids) {
    // for (int id = 0; id < GetEntries(); id++) {
    GetEntry(id);
    fBinning->GetPoint()->Print();
    TList * l = (TList *)fTreeStorage->GetBranch("_outputPoint")->GetObject();
    if (!l || l->IsEmpty()) {
      NLogWarning("NGnTree::Play: No 'outputPoint' for entry %lld !!!", id);
      continue;
    }
    else {
      // NLogInfo("Output for entry %lld:", id);
      // l->Print(opt.Data());
 
      if (outputPointIds.empty()) {
        outputPointIds.resize(l->GetEntries());
        for (int i = 0; i < l->GetEntries(); i++) {
          outputPointIds[i] = i;
        }
      }
      int n = outputPointIds.size();
 
      Double_t v = 1.0;
      for (int i = 0; i < n; i++) {
        // NLogDebug("Drawing output object id %d (list index %d) on pad %d", outputPointIds[i], i, i + 1);
 
        c1->cd(i + 2);
        TObject * obj = l->At(outputPointIds[i]);
        if (obj) {
          if (obj->InheritsFrom(TH1::Class())) {
            TH1 * h = (TH1 *)obj;
            h->SetDirectory(nullptr);
            // Draw a clone to avoid transferring ownership or modifying
            // the original object stored in the TList (can cause
            // TPad/TList removal during drawing and lead to crashes).
            TH1 * hclone = (TH1 *)h->Clone();
            if (hclone) {
              hclone->SetDirectory(nullptr);
              hclone->Draw();
            }
          }
          // obj->Print();
        }
        if (obj->InheritsFrom(TH1::Class()) && i == 0) {
          TH1 * h = (TH1 *)obj;
          v       = h->GetMean();
          NLogDebug("Mean value from histogram [%s]: %f", h->GetName(), v);
        }
      }
      bdContent->SetBinContent(fBinning->GetPoint()->GetStorageCoords(), 1);
      c1->cd(1);
      TH1 * bdProj = (TH1 *)gROOT->FindObjectAny("bdProj");
      if (bdProj) {
        delete bdProj;
        bdProj = nullptr;
      }
      if (bdContent->GetNdimensions() == 1) {
        bdProj = bdContent->Projection(0, "O");
      }
      else if (bdContent->GetNdimensions() == 2) {
        bdProj = bdContent->Projection(0, 1, "O");
      }
      else if (bdContent->GetNdimensions() == 3) {
        bdProj = bdContent->Projection(0, 1, 2, "O");
      }
      else {
        NLogError("NGnTree::Play: Cannot project THnSparse with %d dimensions", bdContent->GetNdimensions());
      }
      if (bdProj) {
        bdProj->SetName("bdProj");
        bdProj->SetTitle(TString::Format("Binning '%s' content projection", binning.c_str()).Data());
        bdProj->SetMinimum(0);
        // bdProj->SetDirectory(nullptr);
        bdProj->Draw("colz");
        // c1->ModifiedUpdate();
      }
    }
    if (c1) {
      c1->ModifiedUpdate();
      c1->SaveAs(TString::Format("%s/ndmspc_play_%06lld.png", annimationTempDir.c_str(), bdContent->GetNbins()).Data());
    }
    gSystem->ProcessEvents();
    if (timeout > 0) gSystem->Sleep(timeout);
    NLogInfo("%d", id);
  }
 
  NLogInfo("Creating animation gif from %s/ndmspc_play_*.png ...", annimationTempDir.c_str());
  gSystem->Exec(
      TString::Format("magick -delay 20 -loop 0 %s/ndmspc_play_*.png ndmspc_play.gif", annimationTempDir.c_str()));
  gSystem->Exec(TString::Format("rm -fr %s", annimationTempDir.c_str()));
  NLogInfo("Animation saved to ndmspc_play.gif");
 
  delete bdContent;
}
 
TList * NGnTree::Projection(const json & cfg, std::string binningName)
{
 
  // SetInput(); // Set input to selfp
  SetInput(Ndmspc::NGnTree::Open(fTreeStorage->GetFileName()));
  Ndmspc::NGnProcessFuncPtr processFunc = [](Ndmspc::NBinningPoint * point, TList * output, TList * /*outputPoint*/,
                                             int /*threadId*/) {
    // NLogInfo("Thread ID: %d", threadId);
    TH1::AddDirectory(kFALSE); // Prevent histograms from being associated with the current directory
    point->Print();
    json cfg = point->GetCfg();
 
    Printf("Processing THnSparse projection with configuration: %s", cfg.dump().c_str());
 
    Ndmspc::NGnTree * ngntIn = point->GetInput();
    // ngntIn->Print();
    // ngntIn->GetEntry(0);
    ngntIn->GetEntry(point->GetEntryNumber());
 
    // loop over all cfg["objects"]
    for (auto & [objName, objCfg] : cfg["objects"].items()) {
      NLogInfo("Processing object '%s' ...", objName.c_str());
 
      THnSparse * hns = (THnSparse *)(ngntIn->GetStorageTree()->GetBranchObject(objName));
      if (hns == nullptr) {
        NLogError("NGnTree::Projection: THnSparse 'hns' not found in storage tree !!!");
        return;
      }
      // hns->Print("all");
      // loop over cfg["objects"][objName] array of projection dimension names
      for (size_t i = 0; i < objCfg.size(); i++) {
 
        NLogInfo("Processing projection %zu for object '%s' ...", i, objName.c_str());
        std::vector<int>         dims;
        std::vector<std::string> dimNames = cfg["objects"][objName][i].get<std::vector<std::string>>();
        for (const auto & dimName : dimNames) {
          NLogDebug("Looking for dimension name '%s' in THnSparse ...", dimName.c_str());
          int dim = -1;
          for (int i = 0; i < hns->GetNdimensions(); i++) {
            if (dimName == hns->GetAxis(i)->GetName()) {
              dim = i;
              break;
            }
          }
          if (dim >= 0)
            dims.push_back(dim);
          else {
            NLogError("NGnTree::Projection: Dimension name '%s' not found in THnSparse !!!", dimName.c_str());
          }
        }
        // Print dims
        NLogInfo("Projecting THnSparse on dimensions: %s", NUtils::GetCoordsString(dims, -1).c_str());
        TH1 * hPrev = (TH1 *)output->At(i);
        TH1 * hProj = NUtils::ProjectTHnSparse(hns, dims, "O");
        hProj->SetName(TString::Format("%s_proj_%s", objName.c_str(), NUtils::Join(dims, '_').c_str()).Data());
        if (hPrev) {
          hPrev->Add(hProj);
        }
        else {
          output->Add(hProj);
        }
      }
    }
    output->Print();
  };
 
  // NBinningDef *                 binningDef = fInput->GetBinning()->GetDefinition(binningName);
  NBinningDef * binningDef = GetBinning()->GetDefinition(binningName);
  THnSparse *   hnsIn      = binningDef->GetContent();
  // std::vector<std::vector<int>> ranges{{0, 2, 2}, {2, 1, 1}};
  std::vector<std::vector<int>> ranges = cfg["ranges"].get<std::vector<std::vector<int>>>();
  NUtils::SetAxisRanges(hnsIn, ranges); // Set the ranges for the axes
  Long64_t                                        linBin = 0;
  std::unique_ptr<ROOT::Internal::THnBaseBinIter> iter{hnsIn->CreateIter(true /*use axis range*/)};
  std::vector<Long64_t>                           ids;
  // std::vector<Long64_t> ids = binningDef->GetIds();
  while ((linBin = iter->Next()) >= 0) {
    ids.push_back(linBin);
  }
  if (ids.empty()) {
    NLogWarning("NGnTree::Projection: No entries found in binning definition '%s' !!!", binningDef->GetName());
    binningDef->RefreshIdsFromContent();
    return nullptr;
  }
 
  binningDef->GetIds() = ids;
 
  // NUtils::SetAxisRanges(, std::vector<std::vector<int>> ranges)
  Process(processFunc, cfg);
 
  // Refresh binning definition ids from content after processing
  binningDef->RefreshIdsFromContent();
 
  // GetInput()->Close(false);
  return GetOutput(fBinning->GetCurrentDefinitionName());
 
  // Close(false);
}
 
NGnNavigator * NGnTree::Reshape(std::string binningName, std::vector<std::vector<int>> levels, int level,
                                std::map<int, std::vector<int>> ranges, std::map<int, std::vector<int>> rangesBase)
{
 
  NGnNavigator navigator;
  navigator.SetGnTree(this);
 
  return navigator.Reshape(binningName, levels, level, ranges, rangesBase);
}
 
NGnNavigator * NGnTree::GetResourceStatisticsNavigator(std::string binningName, std::vector<std::vector<int>> levels,
                                                       int level, std::map<int, std::vector<int>> ranges,
                                                       std::map<int, std::vector<int>> rangesBase)
{
 
  if (binningName.empty()) {
    binningName = fBinning->GetCurrentDefinitionName();
  }
 
  THnSparse * hns = (THnSparse *)fOutputs[binningName]->FindObject("resource_monitor");
  if (!hns) {
    NLogError("NGnTree::Draw: Resource monitor THnSparse not found in outputs !!!");
    return nullptr;
  }
  hns->Print("all");
  // return nullptr;
 
  auto ngnt = new NGnTree(hns, "stat", "/tmp/hnst_imported_for_drawing.root");
  if (ngnt->IsZombie()) {
    NLogError("NGnTree::GetResourceStatisticsNavigator: Failed to import resource monitor THnSparse !!!");
    return nullptr;
  }
  ngnt->Print();
  ngnt->Close();
 
  // return nullptr;
  auto ngnt2 = NGnTree::Open("/tmp/hnst_imported_for_drawing.root");
  auto nav   = ngnt2->Reshape("default", levels, level, ranges, rangesBase);
  // // nav->Export("/tmp/hnst_imported_for_drawing.json", {}, "ws://localhost:8080/ws/root.websocket");
  // // nav->Draw();
  return nav;
}
 
bool NGnTree::InitParameters(const std::vector<std::string> & paramNames)
{
 
  if (fParameters) {
    NLogTrace("NGnTree::InitParameters: Replacing existing parameters ...");
    delete fParameters;
  }
 
  if (paramNames.empty()) {
    NLogTrace("NGnTree::InitParameters: No parameter names provided, skipping ...");
    return false;
  }
 
  fParameters = new NParameters(paramNames, "results", "Results");
 
  return true;
}
 
NGnTree * NGnTree::Import(const std::string & findPath, const std::string & fileName,
                          const std::vector<std::string> & headers, const std::string & outputFile, bool close)
{
 
  // remove trailing slash from findPath if exists
  std::string findPathClean = findPath;
  if (!findPathClean.empty() && findPathClean.back() == '/') {
    findPathClean.pop_back();
  }
 
  std::vector<std::string> paths = NUtils::Find(findPathClean, fileName);
  NLogInfo("NGnTree::Import: Found %zu files to import ...", paths.size());
 
  TObjArray * ngntArray = NUtils::AxesFromDirectory(paths, findPathClean, fileName, headers);
  int         nDirAxes  = ngntArray->GetEntries();
 
  NGnTree * ngntFirst = NGnTree::Open(paths[0]);
  // Add all axes from ngntFirst to ngntArray
  for (const auto & axis : ngntFirst->GetBinning()->GetAxes()) {
    ngntArray->Add(axis->Clone());
  }
  ngntFirst->Close(false);
 
  std::map<std::string, std::vector<std::vector<int>>> b;
 
  for (int i = 0; i < ngntArray->GetEntries(); i++) {
    TAxis * axis = (TAxis *)ngntArray->At(i);
    b[axis->GetName()].push_back({1});
  }
 
  NGnTree * ngnt = new NGnTree(ngntArray, outputFile);
  ngnt->SetIsPureCopy(true);
 
  // return nullptr;
  ngnt->GetBinning()->AddBinningDefinition("default", b);
 
  json cfg;
  cfg["basedir"]  = findPathClean;
  cfg["filename"] = fileName;
  cfg["nDirAxes"] = nDirAxes;
  cfg["headers"]  = headers;
  // cfg["ndmspc"]["shared"]["currentFileName"]  = "";
  Ndmspc::NGnProcessFuncPtr processFunc = [](Ndmspc::NBinningPoint * point, TList * /*output*/, TList * outputPoint,
                                             int /*threadId*/) {
    // point->Print();
 
    json        cfg      = point->GetCfg();
    std::string filename = cfg["basedir"].get<std::string>();
    filename += "/";
    for (auto & header : cfg["headers"]) {
      filename += point->GetBinLabel(header.get<std::string>());
      filename += "/";
    }
    // filename += "/";
    // filename += point->GetBinLabel("c");
    // filename += point->GetBinLabel("year");
    // filename += "/";
    filename += cfg["filename"].get<std::string>();
    NGnTree * ngnt = (NGnTree *)point->GetTempObject("file");
    if (!ngnt || filename.compare(ngnt->GetStorageTree()->GetFileName()) != 0) {
      NLogInfo("NGnTree::Import: Opening file '%s' ...", filename.c_str());
      if (ngnt) {
        NLogDebug("NGnTree::Import: Closing previously opened file '%s' ...",
                  ngnt->GetStorageTree()->GetFileName().c_str());
        ngnt->Close(false);
        // delete ngnt;
        point->SetTempObject("file", nullptr);
      }
      ngnt = NGnTree::Open(filename.c_str());
      if (!ngnt || ngnt->IsZombie()) {
        NLogError("NGnTree::Import: Cannot open file '%s'", filename.c_str());
        return;
      }
      point->SetTempObject("file", ngnt);
    }
 
    int         nDirAxes  = cfg["nDirAxes"].get<int>();
    Int_t *     coords    = point->GetCoords();
    std::string coordsStr = NUtils::GetCoordsString(NUtils::ArrayToVector(coords, point->GetNDimensionsContent()));
    NLogInfo("NGnTree::Import: Processing point with coords %s ...", coordsStr.c_str());
 
    Long64_t entryNumber =
        ngnt->GetBinning()->GetContent()->GetBin(&coords[3 * nDirAxes], kFALSE); // skip first 3 dir axes
    NLogInfo("NGnTree::Import: Corresponding entry number in file '%s' is %lld", filename.c_str(), entryNumber);
 
    ngnt->GetEntry(entryNumber);
 
    // // add outputPoint content to outputPoint list
    // TList * inputOutputPoint = (TList *)ngnt->GetStorageTree()->GetBranch("_outputPoint")->GetObject();
    // for (int i = 0; i < inputOutputPoint->GetEntries(); i++) {
    //   outputPoint->Add(inputOutputPoint->At(i));
    // }
 
    // set all branches from ngnt to branch addresses in current object
    for (const auto & kv : ngnt->GetStorageTree()->GetBranchesMap()) {
      // check if branch exists in current storage tree
      if (point->GetStorageTree()->GetBranch(kv.first) == nullptr) {
        NLogTrace("NGnTree::Import: Adding branch '%s' to storage tree ...", kv.first.c_str());
        point->GetStorageTree()->AddBranch(kv.first, nullptr, kv.second.GetObjectClassName());
      }
      NLogTrace("NGnTree::Import: Setting branch address for branch '%s' ...", kv.first.c_str());
      point->GetTreeStorage()->GetBranch(kv.first)->SetAddress(kv.second.GetObject());
    }
    outputPoint->Add(new TNamed("source_file", filename));
 
    // ngnt->Print();
 
    // NLogInfo("NGnTree::Import: nDirAxes=%d ...", cfg["nDirAxes"].get<int>());
 
    // json & tempCfg = point->GetTempCfg();
    // if (tempCfg["test"].is_null()) {
    //   NLogInfo("Setting temp cfg test value to 42");
    //   tempCfg["test"] = 42;
    // }
    // NLogInfo("Temp cfg test value: %d", tempCfg["test"].get<int>());
 
    // f->ls();
  };
  Ndmspc::NGnBeginFuncPtr beginFunc = [](Ndmspc::NBinningPoint * /*point*/, int /*threadId*/) {
    TH1::AddDirectory(kFALSE); // Prevent histograms from being associated with the current directory
  };
 
  Ndmspc::NGnEndFuncPtr endFunc = [](Ndmspc::NBinningPoint * point, int /*threadId*/) {
    NGnTree * ngnt = (NGnTree *)point->GetTempObject("file");
    if (ngnt) {
      NLogDebug("NGnTree::Import: Closing last file '%s' ...", ngnt->GetStorageTree()->GetFileName().c_str());
      // ngnt->Close(false);
      // delete ngnt;
      point->SetTempObject("file", nullptr);
    }
  };
 
  ngnt->Process(processFunc, cfg, "", beginFunc, endFunc);
  if (close) {
    ngnt->Close(true);
    delete ngnt;
    ngnt = NGnTree::Open(outputFile.c_str());
  }
  return ngnt;
}
 
} // namespace Ndmspc