#include <Rtypes.h>
#include <TFile.h>
#include <TTree.h>
#include <TBranch.h>
#include <TClonesArray.h>
#include <TLorentzVector.h>

#include <string>
#include <cstdio>
#include <cstdlib>
#include <iostream>

Float_t gen_rand() {
	return rand() / static_cast<Double_t>(RAND_MAX);
}

void gen_tree(const std::string pathFile = "file_tcatlv.root") {
	TFile *tfOut = TFile::Open(pathFile.c_str(), "RECREATE");
	tfOut->cd();
	TTree *ttOut = new TTree("Events", "Tree of TClonesArrays of TLorentzVectors");
	auto *tcaTest = new TClonesArray("TLorentzVector", 8);
	size_t nVar;
	ttOut->Branch("TestVar1", tcaTest);
	ttOut->Branch("nTestVar1", &nVar);
	constexpr size_t nEntries = 10;
	for (size_t jEntry = 0; jEntry < nEntries; ++jEntry) {
		tcaTest->Clear();
		nVar = rand() % 8;
		tcaTest->Expand(nVar);
		std::cerr << "jEntry: " << jEntry << "\tnVar: " << nVar << "\tsize: " << tcaTest->GetSize() << std::endl;
		for (size_t iVar = 0; iVar < nVar; ++iVar) {
			const auto &pVar = new((*tcaTest)[iVar]) TLorentzVector(gen_rand(), gen_rand(), gen_rand(), gen_rand());
			std::fprintf(stderr, "(%F, %F, %F, %F) ", pVar->Px(), pVar->Py(), pVar->Pz(), pVar->M());
		}
		std::cerr << std::endl;
		ttOut->Fill();
	}
	ttOut->Write();
	tfOut->Close();
}

int main(int argc, char* argv[]) {
	if (argc > 1) {
		gen_tree(argv[1]);
	} else {
		gen_tree();
	}
	return 0;
}

#include <Rtypes.h>
#include <TFile.h>
#include <TTree.h>
#include <TBranch.h>
#include <TH1.h>
#include <TLorentzVector.h>
#include <ROOT/RDataFrame.hxx>
#include <Math/Vector4D.h>

#include <vector>
#include <string>

void reproduce(const std::string pathFileIn = "file_tcatlv.root", const std::string pathFileOut = "file_out.root") {
	TFile *tfIn = TFile::Open(pathFileIn.c_str());
	TTree *ttIn = tfIn->Get<TTree>("Events");
	ROOT::RDataFrame dfIn(*ttIn);
	std::cerr << "Defining TestVar2 from TestVar1: TLorentzVector -> LorentzVector";
	auto df1 = dfIn.Define("TestVar2", [](const TClonesArray &tcaTest, const size_t n) {
		const size_t nGot = tcaTest.GetSize();
		std::cerr << "TestVar1.GetSize(): " << nGot << "\tn: " << n <<  std::endl;
		ROOT::RVec<ROOT::Math::PtEtaPhiMVector> vTest(n, ROOT::Math::PtEtaPhiMVector());
		for (size_t i = 0; i < n; ++i) {
			const auto &tlv = *static_cast<TLorentzVector *>(tcaTest[i]);
			vTest[i].SetPxPyPzE(tlv.Px(), tlv.Py(), tlv.Pz(), tlv.E());
		}
		return vTest;
	}, { "TestVar1", "nTestVar1" });
	std::cerr << "Redefining TestVar1: TLorentzVector -> LorentzVector" << std::endl;
	auto df2 = df1.Redefine("TestVar1", [](const TClonesArray &tcaTest, const size_t n) {
		const size_t nGot = tcaTest.GetSize();
		std::cerr << "TestVar1.GetSize(): " << nGot << "\tn: " << n <<  std::endl;
		ROOT::RVec<ROOT::Math::PtEtaPhiMVector> vTest(n, ROOT::Math::PtEtaPhiMVector());
		for (size_t i = 0; i < n; ++i) {
			const auto &tlv = *static_cast<TLorentzVector *>(tcaTest[i]);
			vTest[i].SetPxPyPzE(tlv.Px(), tlv.Py(), tlv.Pz(), tlv.E());
		}
		return vTest;
	}, { "TestVar1", "nTestVar1" });
	std::cerr << "Creating a custom column TestVar2" << std::endl;
	auto df3 = df2.Define("TestVar3", [](ROOT::RVec<ROOT::Math::PtEtaPhiMVector> TestVar2) {
		return ROOT::VecOps::Map(TestVar2, [](ROOT::Math::PtEtaPhiMVector lv){return lv.Pt();});
	}, {"TestVar2"});
	std::cerr << "Columns: {" << std::flush;
	for (const auto nameCol: df3.GetColumnNames()) {
		std::cerr << nameCol << ", " << std::flush;
	}
	std::cerr << "}" << std::endl;
	std::cerr << "Generating histogram hTest3 with TestVar3 in df3" << std::endl;
	TFile *tfOut = TFile::Open(pathFileOut.c_str(), "RECREATE");
	tfOut->cd();
	df3.Histo1D({ "hTest3", "", 30, 0., 3. }, "TestVar3")->Write();
	std::cerr << "Closing the output file" << std::endl;
	tfOut->Close();
	std::cerr << "Snapshotting df3 {TestVar3} as tree3 into the output file" << std::endl;
	df3.Snapshot("tree3", pathFileOut, {"TestVar3"});
	std::cerr << "Snapshotting df3 {TestVar2} as tree2 into the output file" << std::endl;
	df3.Snapshot("tree2", pathFileOut, {"TestVar2"});
	std::cerr << "Snapshotting df3 {TestVar1} as tree1 into the output file" << std::endl;
	df3.Snapshot("tree1", pathFileOut, {"TestVar1"});
	std::cerr << "Closing the input file" << std::endl;
	tfIn->Close();
	std::cerr << "Complete!" << std::endl;
}

int main(int argc, char* argv[]) {
	if (argc > 1) {
		reproduce(argv[1], argv[2]);
	} else {
		reproduce();
	}
}