#include <ROOT/RDataFrame.hxx>
#include "ROOT/RVec.hxx"
#include <ROOT/RLogger.hxx>
#include <ROOT/RSnapshotOptions.hxx>
#include <ROOT/RDFHelpers.hxx>

#include "TROOT.h"
#include "TFile.h"
#include <string>

using namespace std;
using namespace ROOT;
using namespace ROOT::VecOps;
double nPIM_nominal=3.8089049E7;

std::vector<ROOT::RDF::RResultPtr<TH1>> doHisto(ROOT::RDF::RNode df, int icut,string cutName) {
  std::vector<ROOT::RDF::RResultPtr<TH1>> ret;
  ret.push_back(df.Histo1D({Form("hWCAL_%i",icut),"hWCAL",800,0.,2.},"ewcal"));
  ret.push_back(df.Histo1D({Form("hVETO_%i",icut),"hVETO",800,0,0.1},"eveto"));
  ret.push_back(df.Histo1D({Form("hHCAL_%i",icut),"hHCAL",800,0,80},"ehcal"));
  ret.push_back(df.Histo2D({Form("hWCAL_HCAL_%i",icut),"hWCAL_HCAL;WCAL;HCAL",500,0,10,800,0,80},"ewcal","ehcal"));
  ret.push_back(df.Histo1D({Form("hRHCAL0_%i",icut),"hRHCAL0",500,0,1},"RHCAL0"));

  return ret;
}

std::vector<ROOT::RDF::RResultPtr<TH1>> allH;
void postCutProcess(int idx, string cutName, ROOT::RDF::RNode df) {
  auto ret = doHisto(df, idx, cutName);
  for (auto h : ret) {
    allH.push_back(h);
  }
}


int main(){

   
  ROOT::EnableImplicitMT(); 

  string fNameOut="fout.root";
  ROOT::RDataFrame d("tout","tmp2/h_MC_*.root"); 
  ROOT::RDF::Experimental::AddProgressBar(d);
 

 
  //0. Defines
  auto d1=d.Define("eveto","VETO_ene[0]+VETO_ene[1]+VETO_ene[2]");
  d1=d1.Define("ehcal","ehcal0+ehcal1+ehcal2");

  int hcalM=0;
  auto Rhcal =[hcalM](RVec<double> &ehcal){
    double rval=0;
    double sum=0;
    for (int iX=0;iX<3;iX++){
      for (int iY=0;iY<3;iY++){
	sum+=ehcal[9*hcalM+3*iX+iY];
      }
    }
    rval=(sum-ehcal[9*hcalM+3*1+1])/sum;
    return rval;
  };
  d1=d1.Define("RHCAL0",Rhcal,{"HCAL_ene"});
 
  //1. All events
  postCutProcess(0,"all",d1);
   
  //2. Select MIPS in WCAL.
  auto d2=d1.Filter("ewcal>0.08&&ewcal<0.3");
  postCutProcess(1,"WCAL_MIP",d2);

  //3. Small-energy in HCAL
  auto d3=d2.Filter("ehcal<6.2");
  postCutProcess(2,"HCAL",d3);

  
  ROOT::RDF::RSnapshotOptions opts;
  opts.fMode = "recreate";
  d3.Snapshot("tout",fNameOut.c_str(),"",opts);
  
  TFile *fout=new TFile(fNameOut.c_str(),"update");
  for (auto h:allH){
    h->Write();
  }
  fout->Close();
  




}