TRestDataSet.cxx

/*************************************************************************
 * This file is part of the REST software framework.                     *
 *                                                                       *
 * Copyright (C) 2016 GIFNA/TREX (University of Zaragoza)                *
 * For more information see https://gifna.unizar.es/trex                 *
 *                                                                       *
 * REST is free software: you can redistribute it and/or modify          *
 * it under the terms of the GNU General Public License as published by  *
 * the Free Software Foundation, either version 3 of the License, or     *
 * (at your option) any later version.                                   *
 *                                                                       *
 * REST is distributed in the hope that it will be useful,               *
 * but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *
 * GNU General Public License for more details.                          *
 *                                                                       *
 * You should have a copy of the GNU General Public License along with   *
 * REST in $REST_PATH/LICENSE.                                           *
 * If not, see https://www.gnu.org/licenses/.                            *
 * For the list of contributors see $REST_PATH/CREDITS.                  *
 *************************************************************************/
 
#include "TRestDataSet.h"
 
#include "TRestRun.h"
#include "TRestTools.h"
 
ClassImp(TRestDataSet);
 
TRestDataSet::TRestDataSet() { Initialize(); }
 
TRestDataSet::TRestDataSet(const char* cfgFileName, const std::string& name) : TRestMetadata(cfgFileName) {
    LoadConfigFromFile(fConfigFileName, name);
}
 
TRestDataSet::~TRestDataSet() {}
 
void TRestDataSet::Initialize() { SetSectionName(this->ClassName()); }
 
void TRestDataSet::GenerateDataSet() {
    if (fTree != nullptr) {
        RESTWarning << "Tree has already been loaded. Skipping TRestDataSet::GenerateDataSet ... "
                    << RESTendl;
        return;
    }
 
    if (fFileSelection.empty()) FileSelection();
 
    // We are not ready yet
    if (fFileSelection.empty()) {
        RESTError << "File selection is empty " << RESTendl;
        return;
    }
 
    TRestRun run(fFileSelection.front());
    std::vector<std::string> finalList;
    finalList.push_back("runOrigin");
    finalList.push_back("eventID");
    finalList.push_back("timeStamp");
 
    auto obsNames = run.GetAnalysisTree()->GetObservableNames();
    for (const auto& obs : fObservablesList) {
        if (std::find(obsNames.begin(), obsNames.end(), obs) != obsNames.end()) {
            finalList.push_back(obs);
        } else {
            RESTWarning << " Observable " << obs << " not found in observable list, skipping..." << RESTendl;
        }
    }
 
    for (const auto& name : obsNames) {
        for (const auto& pcs : fProcessObservablesList) {
            if (name.find(pcs) == 0) finalList.push_back(name);
        }
    }
 
    // Remove duplicated observables if any
    std::sort(finalList.begin(), finalList.end());
    finalList.erase(std::unique(finalList.begin(), finalList.end()), finalList.end());
 
    if (fMT)
        ROOT::EnableImplicitMT();
    else
        ROOT::DisableImplicitMT();
 
    RESTInfo << "Initializing dataset" << RESTendl;
    fDataFrame = ROOT::RDataFrame("AnalysisTree", fFileSelection);
 
    RESTInfo << "Making cuts" << RESTendl;
    fDataFrame = MakeCut(fCut);
 
    // Adding new user columns added to the dataset
    for (const auto& [cName, cExpression] : fColumnNameExpressions) {
        RESTInfo << "Adding column to dataset: " << cName << RESTendl;
        finalList.emplace_back(cName);
        fDataFrame = DefineColumn(cName, cExpression);
    }
 
    RegenerateTree(finalList);
 
    RESTInfo << " - Dataset generated!" << RESTendl;
}
 
void TRestDataSet::RegenerateTree(std::vector<std::string> finalList) {
    RESTInfo << "Generating snapshot." << RESTendl;
    std::string user = getenv("USER");
    std::string fOutName = "/tmp/rest_output_" + user + ".root";
    if (!finalList.empty())
        fDataFrame.Snapshot("AnalysisTree", fOutName, finalList);
    else
        fDataFrame.Snapshot("AnalysisTree", fOutName);
 
    RESTInfo << "Re-importing analysis tree." << RESTendl;
    fDataFrame = ROOT::RDataFrame("AnalysisTree", fOutName);
 
    TFile* f = TFile::Open(fOutName.c_str());
    fTree = (TChain*)f->Get("AnalysisTree");
}
 
std::vector<std::string> TRestDataSet::FileSelection() {
    fFileSelection.clear();
 
    std::time_t time_stamp_start = REST_StringHelper::StringToTimeStamp(fFilterStartTime);
    std::time_t time_stamp_end = REST_StringHelper::StringToTimeStamp(fFilterEndTime);
 
    if (!time_stamp_end || !time_stamp_start) {
        RESTError << "TRestDataSet::FileSelect. Start or end dates not properly formed. Please, check "
                     "REST_StringHelper::StringToTimeStamp documentation for valid formats"
                  << RESTendl;
        return fFileSelection;
    }
 
    std::vector<std::string> fileNames = TRestTools::GetFilesMatchingPattern(fFilePattern);
 
    RESTInfo << "TRestDataSet::FileSelection. Starting file selection." << RESTendl;
    RESTInfo << "Total files : " << fileNames.size() << RESTendl;
    RESTInfo << "This process may take long computation time in case there are many files." << RESTendl;
 
    fTotalDuration = 0;
    std::cout << "Processing file selection.";
    int cnt = 1;
    for (const auto& file : fileNames) {
        if (cnt % 100 == 0) {
            std::cout << std::endl;
            std::cout << "Files processed: " << cnt << " ." << std::flush;
        }
        cnt++;
        TRestRun run(file);
        std::cout << "." << std::flush;
        double runStart = run.GetStartTimestamp();
        double runEnd = run.GetEndTimestamp();
 
        if (runStart < time_stamp_start || runEnd > time_stamp_end) {
            RESTInfo << "Rejecting file out of date range: " << file << RESTendl;
            continue;
        }
 
        int n = 0;
        bool accept = true;
        for (const auto& md : fFilterMetadata) {
            std::string mdValue = run.GetMetadataMember(md);
 
            if (!fFilterContains[n].empty())
                if (mdValue.find(fFilterContains[n]) == std::string::npos) accept = false;
 
            if (fFilterGreaterThan[n] != -1) {
                if (StringToDouble(mdValue) <= fFilterGreaterThan[n]) accept = false;
            }
 
            if (fFilterLowerThan[n] != -1)
                if (StringToDouble(mdValue) >= fFilterLowerThan[n]) accept = false;
 
            if (fFilterEqualsTo[n] != -1)
                if (StringToDouble(mdValue) != fFilterEqualsTo[n]) accept = false;
 
            n++;
        }
 
        if (!accept) continue;
 
        Double_t acc = 0;
        for (auto& [name, properties] : fQuantity) {
            std::string value = run.ReplaceMetadataMembers(properties.metadata);
            const Double_t val = REST_StringHelper::StringToDouble(value);
 
            if (properties.strategy == "accumulate") {
                acc += val;
                properties.value = StringWithPrecision(val, 2);
            }
 
            if (properties.strategy == "max")
                if (properties.value.empty() || REST_StringHelper::StringToDouble(properties.value) < val)
                    properties.value = value;
 
            if (properties.strategy == "min")
                if (properties.value.empty() || REST_StringHelper::StringToDouble(properties.value) > val)
                    properties.value = value;
 
            if (properties.strategy == "unique") {
                if (properties.value.empty())
                    properties.value = value;
                else if (properties.value != value) {
                    RESTWarning << "TRestDataSet::FileSelection. Relevant quantity retrieval." << RESTendl;
                    RESTWarning << "A unique metadata member used for the `" << name
                                << "` quantity is not unique!!" << RESTendl;
                    RESTWarning << "Pre-registered value : " << properties.value << " New value : " << value
                                << RESTendl;
                }
            }
 
            if (properties.strategy == "last") properties.value = value;
        }
 
        if (run.GetStartTimestamp() < fStartTime) fStartTime = run.GetStartTimestamp();
 
        if (run.GetEndTimestamp() > fEndTime) fEndTime = run.GetEndTimestamp();
 
        fTotalDuration += run.GetEndTimestamp() - run.GetStartTimestamp();
        fFileSelection.push_back(file);
    }
    std::cout << std::endl;
 
    return fFileSelection;
}
 
ROOT::RDF::RNode TRestDataSet::Range(size_t from, size_t to) { return fDataFrame.Range(from, to); }
 
ROOT::RDF::RNode TRestDataSet::ApplyRange(size_t from, size_t to) {
    fDataFrame = fDataFrame.Range(from, to);
    RegenerateTree();
    return fDataFrame;
}
 
ROOT::RDF::RNode TRestDataSet::MakeCut(const TRestCut* cut) {
    auto df = fDataFrame;
 
    if (cut == nullptr) return df;
 
    auto paramCut = cut->GetParamCut();
    auto obsList = df.GetColumnNames();
    for (const auto& [param, condition] : paramCut) {
        if (std::find(obsList.begin(), obsList.end(), param) != obsList.end()) {
            std::string pCut = param + condition;
            RESTDebug << "Applying cut " << pCut << RESTendl;
            df = df.Filter(pCut);
        } else {
            RESTWarning << " Cut observable " << param << " not found in observable list, skipping..."
                        << RESTendl;
        }
    }
 
    auto cutString = cut->GetCutStrings();
    for (const auto& pCut : cutString) {
        bool added = false;
        for (const auto& obs : obsList) {
            if (pCut.find(obs) != std::string::npos) {
                RESTDebug << "Applying cut " << pCut << RESTendl;
                df = df.Filter(pCut);
                added = true;
                break;
            }
        }
 
        if (!added) {
            RESTWarning << " Cut string " << pCut << " not found in observable list, skipping..." << RESTendl;
        }
    }
 
    return df;
}
 
size_t TRestDataSet::GetEntries() {
    auto nEntries = fDataFrame.Count();
    if (*nEntries == (long long unsigned int)GetTree()->GetEntries()) return *nEntries;
    RESTWarning << "TRestDataSet::GetEntries. Number of tree entries is not the same as RDataFrame entries."
                << RESTendl;
    RESTWarning << "Returning RDataFrame entries" << RESTendl;
    return *nEntries;
}
 
ROOT::RDF::RNode TRestDataSet::DefineColumn(const std::string& columnName, const std::string& formula) {
    auto df = fDataFrame;
 
    std::string evalFormula = formula;
    for (auto const& [name, properties] : fQuantity)
        evalFormula = REST_StringHelper::Replace(evalFormula, name, properties.value);
 
    df = df.Define(columnName, evalFormula);
 
    return df;
}
 
void TRestDataSet::PrintMetadata() {
    TRestMetadata::PrintMetadata();
 
    RESTMetadata << " - StartTime : " << REST_StringHelper::ToDateTimeString(fStartTime) << RESTendl;
    RESTMetadata << " - EndTime : " << REST_StringHelper::ToDateTimeString(fEndTime) << RESTendl;
    RESTMetadata << " - Path : " << TRestTools::SeparatePathAndName(fFilePattern).first << RESTendl;
    RESTMetadata << " - File pattern : " << TRestTools::SeparatePathAndName(fFilePattern).second << RESTendl;
    RESTMetadata << "  " << RESTendl;
    RESTMetadata << " - Accumulated run time (seconds) : " << fTotalDuration << RESTendl;
    RESTMetadata << " - Accumulated run time (hours) : " << fTotalDuration / 3600. << RESTendl;
    RESTMetadata << " - Accumulated run time (days) : " << fTotalDuration / 3600. / 24. << RESTendl;
 
    RESTMetadata << "  " << RESTendl;
 
    if (!fObservablesList.empty()) {
        RESTMetadata << " Single observables added:" << RESTendl;
        RESTMetadata << " -------------------------" << RESTendl;
        for (const auto& l : fObservablesList) RESTMetadata << " - " << l << RESTendl;
 
        RESTMetadata << "  " << RESTendl;
    }
 
    if (!fProcessObservablesList.empty()) {
        RESTMetadata << " Process observables added: " << RESTendl;
        RESTMetadata << " -------------------------- " << RESTendl;
        for (const auto& l : fProcessObservablesList) RESTMetadata << " - " << l << RESTendl;
 
        RESTMetadata << "  " << RESTendl;
    }
 
    if (!fFilterMetadata.empty()) {
        RESTMetadata << " Metadata filters: " << RESTendl;
        RESTMetadata << " ----------------- " << RESTendl;
        RESTMetadata << " - StartTime : " << fFilterStartTime << RESTendl;
        RESTMetadata << " - EndTime : " << fFilterEndTime << RESTendl;
        int n = 0;
        for (const auto& mdFilter : fFilterMetadata) {
            RESTMetadata << " - " << mdFilter << ".";
 
            if (!fFilterContains[n].empty()) RESTMetadata << " Contains: " << fFilterContains[n];
            if (fFilterGreaterThan[n] != -1) RESTMetadata << " Greater than: " << fFilterGreaterThan[n];
            if (fFilterLowerThan[n] != -1) RESTMetadata << " Lower than: " << fFilterLowerThan[n];
            if (fFilterEqualsTo[n] != -1) RESTMetadata << " Equals to: " << fFilterEqualsTo[n];
 
            RESTMetadata << RESTendl;
            n++;
        }
 
        RESTMetadata << "  " << RESTendl;
    }
 
    if (!fQuantity.empty()) {
        RESTMetadata << " Relevant quantities: " << RESTendl;
        RESTMetadata << " -------------------- " << RESTendl;
 
        for (auto const& [name, properties] : fQuantity) {
            RESTMetadata << " - Name : " << name << ". Value : " << properties.value
                         << ". Strategy: " << properties.strategy << RESTendl;
            RESTMetadata << " - Metadata: " << properties.metadata << RESTendl;
            RESTMetadata << " - Description: " << properties.description << RESTendl;
            RESTMetadata << " " << RESTendl;
        }
    }
 
    if (!fColumnNameExpressions.empty()) {
        RESTMetadata << " New columns added to generated dataframe: " << RESTendl;
        RESTMetadata << " ---------------------------------------- " << RESTendl;
        for (const auto& [cName, cExpression] : fColumnNameExpressions) {
            RESTMetadata << " - Name : " << cName << RESTendl;
            RESTMetadata << " - Expression: " << cExpression << RESTendl;
            RESTMetadata << " " << RESTendl;
        }
    }
 
    if (fMergedDataset) {
        RESTMetadata << " " << RESTendl;
        RESTMetadata << "This is a combined dataset." << RESTendl;
        RESTMetadata << " -------------------- " << RESTendl;
        RESTMetadata << " - Relevant quantities have been removed!" << RESTendl;
        RESTMetadata << " - Dataset metadata properties correspond to the first file in the list."
                     << RESTendl;
        RESTMetadata << " " << RESTendl;
        RESTMetadata << "List of imported files: " << RESTendl;
        RESTMetadata << " -------------------- " << RESTendl;
        for (const auto& fn : fImportedFiles) RESTMetadata << " - " << fn << RESTendl;
    }
 
    RESTMetadata << " " << RESTendl;
    if (fMT)
        RESTMetadata << " - Multithreading was enabled" << RESTendl;
    else
        RESTMetadata << " - Multithreading was NOT enabled" << RESTendl;
 
    RESTMetadata << "----" << RESTendl;
}
 
void TRestDataSet::InitFromConfigFile() {
    TRestMetadata::InitFromConfigFile();
 
    TiXmlElement* filterDefinition = GetElement("filter");
    while (filterDefinition != nullptr) {
        std::string metadata = GetFieldValue("metadata", filterDefinition);
        if (metadata.empty() || metadata == "Not defined") {
            RESTError << "Filter key defined without metadata member!" << RESTendl;
            exit(1);
        }
 
        fFilterMetadata.push_back(metadata);
 
        std::string contains = GetFieldValue("contains", filterDefinition);
        if (contains == "Not defined") contains = "";
        Double_t greaterThan = StringToDouble(GetFieldValue("greaterThan", filterDefinition));
        Double_t lowerThan = StringToDouble(GetFieldValue("lowerThan", filterDefinition));
        Double_t equalsTo = StringToDouble(GetFieldValue("equalsTo", filterDefinition));
 
        fFilterContains.push_back(contains);
        fFilterGreaterThan.push_back(greaterThan);
        fFilterLowerThan.push_back(lowerThan);
        fFilterEqualsTo.push_back(equalsTo);
 
        filterDefinition = GetNextElement(filterDefinition);
    }
 
    TiXmlElement* observablesDefinition = GetElement("observables");
    while (observablesDefinition != nullptr) {
        std::string observables = GetFieldValue("list", observablesDefinition);
        if (observables.empty() || observables == "Not defined") {
            RESTError << "<observables key does not contain a list!" << RESTendl;
            exit(1);
        }
 
        std::vector<std::string> obsList = REST_StringHelper::Split(observables, ",");
 
        fObservablesList.insert(fObservablesList.end(), obsList.begin(), obsList.end());
 
        observablesDefinition = GetNextElement(observablesDefinition);
    }
 
    TiXmlElement* obsProcessDefinition = GetElement("processObservables");
    while (obsProcessDefinition != nullptr) {
        std::string observables = GetFieldValue("list", obsProcessDefinition);
        if (observables.empty() || observables == "Not defined") {
            RESTError << "<processObservables key does not contain a list!" << RESTendl;
            exit(1);
        }
 
        std::vector<std::string> obsList = REST_StringHelper::Split(observables, ",");
 
        for (const auto& l : obsList) fProcessObservablesList.push_back(l);
 
        obsProcessDefinition = GetNextElement(obsProcessDefinition);
    }
 
    TiXmlElement* quantityDefinition = GetElement("quantity");
    while (quantityDefinition != nullptr) {
        std::string name = GetFieldValue("name", quantityDefinition);
        if (name.empty() || name == "Not defined") {
            RESTError << "<quantity key does not contain a name!" << RESTendl;
            exit(1);
        }
 
        std::string metadata = GetFieldValue("metadata", quantityDefinition);
        if (metadata.empty() || metadata == "Not defined") {
            RESTError << "<quantity key does not contain a metadata value!" << RESTendl;
            exit(1);
        }
 
        std::string strategy = GetFieldValue("strategy", quantityDefinition);
        if (strategy.empty() || strategy == "Not defined") {
            strategy = "unique";
        }
 
        std::string description = GetFieldValue("description", quantityDefinition);
 
        RelevantQuantity quantity;
        quantity.metadata = metadata;
        quantity.strategy = strategy;
        quantity.description = description;
        quantity.value = "";
 
        fQuantity[name] = quantity;
 
        quantityDefinition = GetNextElement(quantityDefinition);
    }
 
    TiXmlElement* columnDefinition = GetElement("addColumn");
    while (columnDefinition != nullptr) {
        std::string name = GetFieldValue("name", columnDefinition);
        if (name.empty() || name == "Not defined") {
            RESTError << "<define key does not contain a name name!" << RESTendl;
            exit(1);
        }
 
        std::string expression = GetFieldValue("expression", columnDefinition);
        if (expression.empty() || expression == "Not defined") {
            RESTError << "<addColumn key does not contain a expression value!" << RESTendl;
            exit(1);
        }
 
        fColumnNameExpressions.push_back({name, expression});
 
        columnDefinition = GetNextElement(columnDefinition);
    }
 
    fCut = (TRestCut*)InstantiateChildMetadata("TRestCut");
}
 
void TRestDataSet::Export(const std::string& filename, std::vector<std::string> excludeColumns) {
    RESTInfo << "Exporting dataset" << RESTendl;
 
    std::vector<std::string> columns = fDataFrame.GetColumnNames();
    if (!excludeColumns.empty()) {
        columns.erase(std::remove_if(columns.begin(), columns.end(),
                                     [&excludeColumns](std::string elem) {
                                         return std::find(excludeColumns.begin(), excludeColumns.end(),
                                                          elem) != excludeColumns.end();
                                     }),
                      columns.end());
 
        RESTInfo << "Re-Generating snapshot." << RESTendl;
        std::string user = getenv("USER");
        std::string fOutName = "/tmp/rest_output_" + user + ".root";
        fDataFrame.Snapshot("AnalysisTree", fOutName, columns);
 
        RESTInfo << "Re-importing analysis tree." << RESTendl;
        fDataFrame = ROOT::RDataFrame("AnalysisTree", fOutName);
 
        TFile* f = TFile::Open(fOutName.c_str());
        fTree = (TChain*)f->Get("AnalysisTree");
    }
 
    if (TRestTools::GetFileNameExtension(filename) == "txt" ||
        TRestTools::GetFileNameExtension(filename) == "csv") {
        if (excludeColumns.empty()) {
            RESTInfo << "Re-Generating snapshot." << RESTendl;
            std::string user = getenv("USER");
            std::string fOutName = "/tmp/rest_output_" + user + ".root";
            fDataFrame.Snapshot("AnalysisTree", fOutName);
 
            TFile* f = TFile::Open(fOutName.c_str());
            fTree = (TChain*)f->Get("AnalysisTree");
        }
 
        std::vector<std::string> dataTypes;
        for (int n = 0; n < fTree->GetListOfBranches()->GetEntries(); n++) {
            std::string bName = fTree->GetListOfBranches()->At(n)->GetName();
            std::string type = fTree->GetLeaf((TString)bName)->GetTypeName();
            dataTypes.push_back(type);
            if (type != "Double_t" && type != "Int_t") {
                RESTError << "Branch name : " << bName << " is type : " << type << RESTendl;
                RESTError << "Only Int_t and Double_t types are allowed for "
                             "exporting to ASCII table"
                          << RESTendl;
                RESTError << "File will not be generated" << RESTendl;
                return;
            }
        }
 
        FILE* f = fopen(filename.c_str(), "wt");
        fprintf(f, "### TRestDataSet generated file\n");
        fprintf(f, "### \n");
        fprintf(f, "### StartTime : %s\n", fFilterStartTime.c_str());
        fprintf(f, "### EndTime : %s\n", fFilterEndTime.c_str());
        fprintf(f, "###\n");
        fprintf(f, "### Accumulated run time (seconds) : %lf\n", fTotalDuration);
        fprintf(f, "### Accumulated run time (hours) : %lf\n", fTotalDuration / 3600.);
        fprintf(f, "### Accumulated run time (days) : %lf\n", fTotalDuration / 3600. / 24.);
        fprintf(f, "###\n");
        fprintf(f, "### Data path : %s\n", TRestTools::SeparatePathAndName(fFilePattern).first.c_str());
        fprintf(f, "### File pattern : %s\n", TRestTools::SeparatePathAndName(fFilePattern).second.c_str());
        fprintf(f, "###\n");
        if (!fFilterMetadata.empty()) {
            fprintf(f, "### Metadata filters : \n");
            int n = 0;
            for (const auto& md : fFilterMetadata) {
                fprintf(f, "### - %s.", md.c_str());
                if (!fFilterContains[n].empty()) fprintf(f, " Contains: %s.", fFilterContains[n].c_str());
                if (fFilterGreaterThan[n] != -1) fprintf(f, " Greater than: %6.3lf.", fFilterGreaterThan[n]);
                if (fFilterLowerThan[n] != -1) fprintf(f, " Lower than: %6.3lf.", fFilterLowerThan[n]);
                if (fFilterEqualsTo[n] != -1) fprintf(f, " Equals to: %6.3lf.", fFilterLowerThan[n]);
                fprintf(f, "\n");
                n++;
            }
        }
        fprintf(f, "###\n");
        fprintf(f, "### Relevant quantities: \n");
        for (auto& [name, properties] : fQuantity) {
            fprintf(f, "### - %s : %s - %s\n", name.c_str(), properties.value.c_str(),
                    properties.description.c_str());
        }
        fprintf(f, "###\n");
        fprintf(f, "### Observables list: ");
        for (int n = 0; n < fTree->GetListOfBranches()->GetEntries(); n++) {
            std::string bName = fTree->GetListOfBranches()->At(n)->GetName();
            fprintf(f, " %s", bName.c_str());
        }
        fprintf(f, "\n");
        fprintf(f, "###\n");
        fprintf(f, "### Data starts here\n");
 
        auto obsNames = fDataFrame.GetColumnNames();
        std::string obsListStr = "";
        for (const auto& l : obsNames) {
            if (!obsListStr.empty()) obsListStr += ":";
            obsListStr += l;
        }
 
        // We do this so that later we can recover the values using TTree::GetVal
        fTree->Draw((TString)obsListStr, "", "goff");
 
        for (unsigned int n = 0; n < fTree->GetEntries(); n++) {
            for (unsigned int m = 0; m < GetNumberOfBranches(); m++) {
                std::string bName = fTree->GetListOfBranches()->At(m)->GetName();
                if (m > 0) fprintf(f, "\t");
                if (dataTypes[m] == "Double_t")
                    if (bName == "timeStamp")
                        fprintf(f, "%010.0lf", fTree->GetVal(m)[n]);
                    else
                        fprintf(f, "%05.3e", fTree->GetVal(m)[n]);
                else
                    fprintf(f, "%06d", (Int_t)fTree->GetVal(m)[n]);
            }
            fprintf(f, "\n");
        }
        fclose(f);
 
        return;
    } else if (TRestTools::GetFileNameExtension(filename) == "root") {
        fDataFrame.Snapshot("AnalysisTree", filename);
 
        TFile* f = TFile::Open(filename.c_str(), "UPDATE");
        std::string name = this->GetName();
        if (name.empty()) name = "mock";
        this->Write(name.c_str());
        f->Close();
    } else {
        RESTWarning << "TRestDataSet::Export. Extension " << TRestTools::GetFileNameExtension(filename)
                    << " not recognized" << RESTendl;
    }
    RESTInfo << "Dataset generated: " << filename << RESTendl;
}
 
TRestDataSet& TRestDataSet::operator=(TRestDataSet& dS) {
    SetName(dS.GetName());
    fFilterStartTime = dS.GetFilterStartTime();
    fFilterEndTime = dS.GetFilterEndTime();
    fStartTime = dS.GetStartTime();
    fEndTime = dS.GetEndTime();
    fFilePattern = dS.GetFilePattern();
    fObservablesList = dS.GetObservablesList();
    fFileSelection = dS.GetFileSelection();
    fProcessObservablesList = dS.GetProcessObservablesList();
    fFilterMetadata = dS.GetFilterMetadata();
    fFilterContains = dS.GetFilterContains();
    fFilterGreaterThan = dS.GetFilterGreaterThan();
    fFilterLowerThan = dS.GetFilterLowerThan();
    fFilterEqualsTo = dS.GetFilterEqualsTo();
    fQuantity = dS.GetQuantity();
    fColumnNameExpressions = dS.GetAddedColumns();
    fTotalDuration = dS.GetTotalTimeInSeconds();
    fCut = dS.GetCut();
 
    return *this;
}
 
Bool_t TRestDataSet::Merge(const TRestDataSet& dS) {
    auto obsNames = GetObservablesList();
    for (const auto& obs : fObservablesList) {
        if (std::find(obsNames.begin(), obsNames.end(), obs) != obsNames.end()) {
            RESTError << "Cannot merge dataSets with different observable list " << RESTendl;
            return false;
        }
    }
 
    if (fStartTime > dS.GetStartTime()) fStartTime = dS.GetStartTime();
    if (fEndTime < dS.GetEndTime()) fEndTime = dS.GetEndTime();
 
    auto fileSelection = dS.GetFileSelection();
    fFileSelection.insert(fFileSelection.end(), fileSelection.begin(), fileSelection.end());
 
    fTotalDuration += dS.GetTotalTimeInSeconds();
 
    return true;
}
 
void TRestDataSet::Import(const std::string& fileName) {
    if (TRestTools::GetFileNameExtension(fileName) != "root") {
        RESTError << "Datasets can only be imported from root files" << RESTendl;
        return;
    }
 
    TRestDataSet* dS = nullptr;
    TFile* file = TFile::Open(fileName.c_str(), "READ");
    if (file != nullptr) {
        TIter nextkey(file->GetListOfKeys());
        TKey* key;
        while ((key = (TKey*)nextkey())) {
            std::string kName = key->GetClassName();
            if (REST_Reflection::GetClassQuick(kName.c_str()) != nullptr &&
                REST_Reflection::GetClassQuick(kName.c_str())->InheritsFrom("TRestDataSet")) {
                dS = file->Get<TRestDataSet>(key->GetName());
                *this = *dS;
            }
        }
    }
 
    if (dS == nullptr) {
        RESTError << fileName << " is not a valid dataSet" << RESTendl;
        return;
    }
 
    if (fMT)
        ROOT::EnableImplicitMT();
    else
        ROOT::DisableImplicitMT();
 
    fDataFrame = ROOT::RDataFrame("AnalysisTree", fileName);
 
    fTree = (TChain*)file->Get("AnalysisTree");
}
 
void TRestDataSet::Import(std::vector<std::string> fileNames) {
    for (const auto& fN : fileNames)
        if (TRestTools::GetFileNameExtension(fN) != "root") {
            RESTError << "Datasets can only be imported from root files" << RESTendl;
            return;
        }
 
    int count = 0;
    auto it = fileNames.begin();
    while (it != fileNames.end()) {
        std::string fileName = *it;
        TFile* file = TFile::Open(fileName.c_str(), "READ");
        bool isValid = false;
        if (file != nullptr) {
            TIter nextkey(file->GetListOfKeys());
            TKey* key;
            while ((key = (TKey*)nextkey())) {
                std::string kName = key->GetClassName();
                if (REST_Reflection::GetClassQuick(kName.c_str()) != nullptr &&
                    REST_Reflection::GetClassQuick(kName.c_str())->InheritsFrom("TRestDataSet")) {
                    TRestDataSet* dS = file->Get<TRestDataSet>(key->GetName());
                    if (GetVerboseLevel() >= TRestStringOutput::REST_Verbose_Level::REST_Info)
                        dS->PrintMetadata();
 
                    if (count == 0) {
                        *this = *dS;
                        isValid = true;
                    } else {
                        isValid = Merge(*dS);
                    }
 
                    if (isValid) count++;
                }
            }
        } else {
            RESTError << "Cannot open " << fileName << RESTendl;
        }
 
        if (!isValid) {
            RESTError << fileName << " is not a valid dataSet skipping..." << RESTendl;
            it = fileNames.erase(it);
        } else {
            ++it;
        }
    }
 
    if (fileNames.empty()) {
        RESTError << "File selection is empty, dataSet will not be imported " << RESTendl;
        return;
    }
 
    RESTInfo << "Opening list of files. First file: " << fileNames[0] << RESTendl;
    fDataFrame = ROOT::RDataFrame("AnalysisTree", fileNames);
 
    if (fTree != nullptr) {
        delete fTree;
        fTree = nullptr;
    }
    fTree = new TChain("AnalysisTree");
 
    for (const auto& fN : fileNames) fTree->Add((TString)fN);
 
    fMergedDataset = true;
    fImportedFiles = fileNames;
 
    fQuantity.clear();
}