<commit_before>// Copyright (C) 2011 <NAME> <<EMAIL>>

// This program 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 2 of the License, or
// (at your option) any later version.
//
// This program 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 received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.


#include "Selector.h"

#include <iostream>
#include <iomanip>
#include <sstream>
#include <limits>
#include <cassert>

#include <TFile.h>
#include <TObjArray.h>
#include <TChainElement.h>
#include <TPRegexp.h>
#include <TTreeFormula.h>
#include <TTreeFormulaManager.h>
#include <TEntryList.h>

#include "util.h"
#include "File.h"
#include "Settings.h"


using namespace std;


namespace froast {


void Selector::copyObject(TDirectory *tdir, const TString &objName) {
	TObject *keepObj = 0;
	tdir->GetObject(objName.Data(), keepObj);
	if (keepObj != 0) {
		TTree *keepTree = dynamic_cast<TTree*>(keepObj);
		if (keepTree != 0) {
			TTree *cloned = keepTree->CloneTree();
			cloned->Write();
		}
	} else {
		throw runtime_error(string("Object ") + objName.Data() + " not found in TDirectory");
	}
}


void Selector::mapMulti(TChain *chain, const TString &selector, const TString &tag, const TString &keep) {
	TObjArray *chainElems = chain->GetListOfFiles();
	for (int chainEntry = 0; chainEntry < chainElems->GetEntriesFast(); ++chainEntry) {
		TChainElement *e = dynamic_cast<TChainElement*>(chainElems->At(chainEntry));
		string treeName = e->GetName();
		string inFileName = e->GetTitle();

		string outFileName = (File(inFileName).base() % tag.Data()).path();
		cerr << "Mapping " << inFileName << ":" << treeName << " to " << outFileName << endl;

		TFile inFile(inFileName.c_str(), "read");
		TTree *inTree; inFile.GetObject(treeName.c_str(), inTree);
		
		TFile outFile(outFileName.c_str(), "recreate");
		outFile.SetCompressionLevel(GSettings::get("froast.tfile.compression.level", 1));

		TSelector *sel = TSelector::GetSelector(selector.Data());
		if (sel == 0) throw runtime_error(string("Cannot load selector ") + selector.Data());
		inTree->Process(sel);
		delete sel;

		TObjArray* keeps = keep.Tokenize(",");
		for (int i = 0; i < keeps->GetEntriesFast(); ++i) {
			TString keepObjName = dynamic_cast<TObjString*>(keeps->At(i))->GetString().Strip(TString::kBoth);
			cerr << "Copying object " << keepObjName << " to output" << endl;
			copyObject(&inFile, keepObjName);
		}
		delete keeps;
		
		Settings::global().writeToGDirectory();
		outFile.Write();
		outFile.Close();
		inFile.Close();
	}
}


void Selector::mapSingle(const TString &inFileName, const TString &mappers, const TString &outFileName, bool noRecompile) {
	TFile inFile(inFileName.Data(), "read");
	TFile outFile(outFileName.Data(), "recreate");
	outFile.SetCompressionLevel(GSettings().get("froast.tfile.compression.level", 1));

	TPRegexp mapperSpecExpr("^(.*)\\((.*)\\)$");
	TPRegexp xxExp("\\+\\+$");

	TObjArray* mapperNames = mappers.Tokenize(",");
	for (int mapperEntry = 0; mapperEntry < mapperNames->GetEntriesFast(); ++mapperEntry) {
		TString mapper = dynamic_cast<TObjString*>(mapperNames->At(mapperEntry))->GetString().Strip(TString::kBoth);
		TObjArray* matches = mapperSpecExpr.MatchS(mapper);
		if (matches->GetEntriesFast() == 3) {
			TString fctName = dynamic_cast<TObjString*>(matches->At(1))->GetString();
			if (noRecompile) xxExp.Substitute(fctName, "+"); 
			TString objName = dynamic_cast<TObjString*>(matches->At(2))->GetString();
			assert ((fctName != 0) && (objName != 0));
			delete matches;
			
			cerr << "Applying " << fctName << "(" << objName << ")" << endl;
			
			TObject *inObj; inFile.GetObject(objName.Data(), inObj);
			if (inObj != 0) {
				TTree *inTree = dynamic_cast<TTree*>(inObj);
				if (inTree != 0) {
					if (fctName == "copy") {
						TTree *copied = inTree->CloneTree();
						copied->Write();
					} else {
						TSelector *sel = TSelector::GetSelector(fctName.Data());
						if (sel == 0) throw runtime_error(string("Cannot load selector ") + fctName.Data());
						inTree->Process(sel);
						delete sel;
					}
				} else throw invalid_argument(string("Objects of type ") + inObj->Class()->GetName() + " not supported yet");
			} else {
				delete matches;
				throw runtime_error(string("Object ") + objName.Data() + " not found in TDirectory");
			}
		} else {
			delete matches;
			throw invalid_argument(string("Invalid mapper specification: \"") + mapper.Data() + "\"");
		}
	}
	delete mapperNames;
	
	Settings::global().writeToGDirectory();
	outFile.Write();
	outFile.Close();
	inFile.Close();
}


void Selector::mapMulti(const TString &fileName, const TString &mappers, const TString &tag, bool noRecompile) {
	cerr << TString::Format("Selector::map(%s, %s, %s)", fileName.Data(), mappers.Data(), tag.Data()) << endl;

	TChain chain("");
	chain.Add(fileName.Data());
	TObjArray *chainElems = chain.GetListOfFiles();
	for (int chainEntry = 0; chainEntry < chainElems->GetEntriesFast(); ++chainEntry) {
		TChainElement *e = dynamic_cast<TChainElement*>(chainElems->At(chainEntry));
		string inFileName = e->GetTitle();

		string outFileName = (File(inFileName).base() % tag.Data()).path();
		cerr << "Mapping " << inFileName << " to " << outFileName << endl;
		// Don't recompile even if fct ends with "++" after first run:
		mapSingle(inFileName, mappers, outFileName, (chainEntry > 0) || noRecompile);
	}
	cerr << "Selector::map(...) finished" << endl;
}



// Based in part on TTreePlayer::scan (Copyright (C) 1995-2000, <NAME>
// and <NAME>)
void Selector::tabulate(TTree *chain, ostream &out, const TString &varexp, const TString &selection, ssize_t nEntries, ssize_t startEntry) {
	cerr << TString::Format("Selector::tabulate(TChain*, ostream, \"%s\", \"%s\")", varexp.Data(), selection.Data()) << endl;

	ssize_t logEvery = GSettings::get("selector.log.every", 10000);

	const TString FS_TSV = "tsv";
	const TString FS_JSON = "json";
	
	out.unsetf(ios::fixed | ios::scientific);
	out << std::setprecision(std::numeric_limits<double>::digits10);

	TPRegexp varexpExpr("^(([^>]|>[^>])*)\\s*(>>\\s*(\\w+)\\s*(\\((.*)\\))?)?$");
	vector<TString> varexpParts;
	Util::match(varexp, varexpExpr, varexpParts, TString::kBoth);

	if (varexpParts.size() <= 1) throw invalid_argument("Invalid varexp");
	vector<TString> functions;
	Util::split(varexpParts[1], ":", functions, TString::kBoth);
	TString format = FS_TSV;
	if (varexpParts.size() > 4) format = varexpParts[4];
	vector<TString> labels;
	if (varexpParts.size() > 5) {
		if (varexpParts.size() > 6) Util::split(varexpParts[6], ":", labels, TString::kBoth);
		size_t nSpecLabels = labels.size();
		labels.resize(functions.size());
		for (size_t i = nSpecLabels; i < labels.size(); ++i)
			labels[i] = functions[i];
	}
	
	cerr << "Tabulation expression: ";
	for (size_t i = 0; i < functions.size(); ++i) cerr << (i>0 ? ":" : "") << functions[i];
	cerr << " >> " << format;
	if (labels.size() > 0) {
		cerr << "(";
		for (size_t i = 0; i < labels.size(); ++i) cerr << (i>0 ? ":" : "") << labels[i];
		cerr << ")";
	}
	cerr << endl;

	const size_t ncols = functions.size();

	TList tformulas;
	
	TTreeFormula *select  = 0;
	if (selection.Length() > 0) {
		select = new TTreeFormula("Selection", selection.Data(), chain);
		if (!select) throw invalid_argument("Invalid selection expression");
		if (!select->GetNdim()) { throw invalid_argument("Invalid selection expression"); }
		tformulas.Add(select);
	}
	
	vector<TTreeFormula*> colFormulas; colFormulas.reserve(ncols);
	for (size_t col = 0; col < ncols; ++col) {
		colFormulas.push_back(new TTreeFormula(TString::Format("col%lli", (long long)(col)).Data(), functions[col].Data(), chain));
		tformulas.Add(colFormulas[col]);<commit_msg>Fixed column-name genaration for JSON output in Selector::tabulate<commit_after>// Copyright (C) 2011 <NAME> <<EMAIL>>

// This program 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 2 of the License, or
// (at your option) any later version.
//
// This program 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 received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.


#include "Selector.h"

#include <iostream>
#include <iomanip>
#include <sstream>
#include <limits>
#include <cassert>

#include <TFile.h>
#include <TObjArray.h>
#include <TChainElement.h>
#include <TPRegexp.h>
#include <TTreeFormula.h>
#include <TTreeFormulaManager.h>
#include <TEntryList.h>

#include "util.h"
#include "File.h"
#include "Settings.h"


using namespace std;


namespace froast {


void Selector::copyObject(TDirectory *tdir, const TString &objName) {
	TObject *keepObj = 0;
	tdir->GetObject(objName.Data(), keepObj);
	if (keepObj != 0) {
		TTree *keepTree = dynamic_cast<TTree*>(keepObj);
		if (keepTree != 0) {
			TTree *cloned = keepTree->CloneTree();
			cloned->Write();
		}
	} else {
		throw runtime_error(string("Object ") + objName.Data() + " not found in TDirectory");
	}
}


void Selector::mapMulti(TChain *chain, const TString &selector, const TString &tag, const TString &keep) {
	TObjArray *chainElems = chain->GetListOfFiles();
	for (int chainEntry = 0; chainEntry < chainElems->GetEntriesFast(); ++chainEntry) {
		TChainElement *e = dynamic_cast<TChainElement*>(chainElems->At(chainEntry));
		string treeName = e->GetName();
		string inFileName = e->GetTitle();

		string outFileName = (File(inFileName).base() % tag.Data()).path();
		cerr << "Mapping " << inFileName << ":" << treeName << " to " << outFileName << endl;

		TFile inFile(inFileName.c_str(), "read");
		TTree *inTree; inFile.GetObject(treeName.c_str(), inTree);
		
		TFile outFile(outFileName.c_str(), "recreate");
		outFile.SetCompressionLevel(GSettings::get("froast.tfile.compression.level", 1));

		TSelector *sel = TSelector::GetSelector(selector.Data());
		if (sel == 0) throw runtime_error(string("Cannot load selector ") + selector.Data());
		inTree->Process(sel);
		delete sel;

		TObjArray* keeps = keep.Tokenize(",");
		for (int i = 0; i < keeps->GetEntriesFast(); ++i) {
			TString keepObjName = dynamic_cast<TObjString*>(keeps->At(i))->GetString().Strip(TString::kBoth);
			cerr << "Copying object " << keepObjName << " to output" << endl;
			copyObject(&inFile, keepObjName);
		}
		delete keeps;
		
		Settings::global().writeToGDirectory();
		outFile.Write();
		outFile.Close();
		inFile.Close();
	}
}


void Selector::mapSingle(const TString &inFileName, const TString &mappers, const TString &outFileName, bool noRecompile) {
	TFile inFile(inFileName.Data(), "read");
	TFile outFile(outFileName.Data(), "recreate");
	outFile.SetCompressionLevel(GSettings().get("froast.tfile.compression.level", 1));

	TPRegexp mapperSpecExpr("^(.*)\\((.*)\\)$");
	TPRegexp xxExp("\\+\\+$");

	TObjArray* mapperNames = mappers.Tokenize(",");
	for (int mapperEntry = 0; mapperEntry < mapperNames->GetEntriesFast(); ++mapperEntry) {
		TString mapper = dynamic_cast<TObjString*>(mapperNames->At(mapperEntry))->GetString().Strip(TString::kBoth);
		TObjArray* matches = mapperSpecExpr.MatchS(mapper);
		if (matches->GetEntriesFast() == 3) {
			TString fctName = dynamic_cast<TObjString*>(matches->At(1))->GetString();
			if (noRecompile) xxExp.Substitute(fctName, "+"); 
			TString objName = dynamic_cast<TObjString*>(matches->At(2))->GetString();
			assert ((fctName != 0) && (objName != 0));
			delete matches;
			
			cerr << "Applying " << fctName << "(" << objName << ")" << endl;
			
			TObject *inObj; inFile.GetObject(objName.Data(), inObj);
			if (inObj != 0) {
				TTree *inTree = dynamic_cast<TTree*>(inObj);
				if (inTree != 0) {
					if (fctName == "copy") {
						TTree *copied = inTree->CloneTree();
						copied->Write();
					} else {
						TSelector *sel = TSelector::GetSelector(fctName.Data());
						if (sel == 0) throw runtime_error(string("Cannot load selector ") + fctName.Data());
						inTree->Process(sel);
						delete sel;
					}
				} else throw invalid_argument(string("Objects of type ") + inObj->Class()->GetName() + " not supported yet");
			} else {
				delete matches;
				throw runtime_error(string("Object ") + objName.Data() + " not found in TDirectory");
			}
		} else {
			delete matches;
			throw invalid_argument(string("Invalid mapper specification: \"") + mapper.Data() + "\"");
		}
	}
	delete mapperNames;
	
	Settings::global().writeToGDirectory();
	outFile.Write();
	outFile.Close();
	inFile.Close();
}


void Selector::mapMulti(const TString &fileName, const TString &mappers, const TString &tag, bool noRecompile) {
	cerr << TString::Format("Selector::map(%s, %s, %s)", fileName.Data(), mappers.Data(), tag.Data()) << endl;

	TChain chain("");
	chain.Add(fileName.Data());
	TObjArray *chainElems = chain.GetListOfFiles();
	for (int chainEntry = 0; chainEntry < chainElems->GetEntriesFast(); ++chainEntry) {
		TChainElement *e = dynamic_cast<TChainElement*>(chainElems->At(chainEntry));
		string inFileName = e->GetTitle();

		string outFileName = (File(inFileName).base() % tag.Data()).path();
		cerr << "Mapping " << inFileName << " to " << outFileName << endl;
		// Don't recompile even if fct ends with "++" after first run:
		mapSingle(inFileName, mappers, outFileName, (chainEntry > 0) || noRecompile);
	}
	cerr << "Selector::map(...) finished" << endl;
}



// Based in part on TTreePlayer::scan (Copyright (C) 1995-2000, <NAME>
// and <NAME>)
void Selector::tabulate(TTree *chain, ostream &out, const TString &varexp, const TString &selection, ssize_t nEntries, ssize_t startEntry) {
	cerr << TString::Format("Selector::tabulate(TChain*, ostream, \"%s\", \"%s\")", varexp.Data(), selection.Data()) << endl;

	ssize_t logEvery = GSettings::get("selector.log.every", 10000);

	const TString FS_TSV = "tsv";
	const TString FS_JSON = "json";
	
	out.unsetf(ios::fixed | ios::scientific);
	out << std::setprecision(std::numeric_limits<double>::digits10);

	TPRegexp varexpExpr("^(([^>]|>[^>])*)\\s*(>>\\s*(\\w+)\\s*(\\((.*)\\))?)?$");
	vector<TString> varexpParts;
	Util::match(varexp, varexpExpr, varexpParts, TString::kBoth);

	if (varexpParts.size() <= 1) throw invalid_argument("Invalid varexp");
	vector<TString> functions;
	Util::split(varexpParts[1], ":", functions, TString::kBoth);
	TString format = FS_TSV;
	if (varexpParts.size() > 4) format = varexpParts[4];
	vector<TString> labels;
	if (varexpParts.size() > 5) {
		if (varexpParts.size() > 6) Util::split(varexpParts[6], ":", labels, TString::kBoth);
		size_t nSpecLabels = labels.size();
		labels.resize(functions.size());
		for (size_t col = nSpecLabels; col < labels.size(); ++col) {
			labels[col] = functions[col];
			if (format == FS_JSON) {
				for (ssize_t i = 0; i < labels[col].Length(); ++i)
					if (labels[col](i) == '.') labels[col](i) = '$';
			}
		}
	}
	
	cerr << "Tabulation expression: ";
	for (size_t i = 0; i < functions.size(); ++i) cerr << (i>0 ? ":" : "") << functions[i];
	cerr << " >> " << format;
	if (labels.size() > 0) {
		cerr << "(";
		for (size_t i = 0; i < labels.size(); ++i) cerr << (i>0 ? ":" : "") << labels[i];
		cerr << ")";
	}
	cerr << endl;

	const size_t ncols = functions.size();

	TList tformulas;
	
	TTreeFormula *select  = 0;
	if (selection.Length() > 0) {
		select = new TTreeFormula("Selection", selection.Data(), chain);
		if (!select) throw invalid_argument("Invalid selection expression");
		if (!select->GetNdim()) { throw invalid_argument("Invalid selection expression"); }
		tformulas.Add(select);
	}
	
	vector<TTreeFormula*> colFormulas; colFormulas.reserve(ncols);
	for (size_t col = 0; col < ncols; ++col) {
		colFormulas.push_back(new TTreeFormula(TString::Format("col%lli", (long long)(col)).Data(), functions[col].Data(), chain));
		tformulas.Add(colFormulas[col]);
{'fraction_non_alphanumeric': 0.17840489311451582, 'fraction_numerical': 0.008296493671652637, 'mean_word_length': 4.702348066298343, 'pattern_counts': {'":': 10, '<': 238, '<?xml version=': 0, '>': 151, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2799544', 'n_tokens_mistral': 6096, 'n_tokens_neox': 5562, 'n_words': 1349}