#include <iostream>
#include <fstream>
#include "TFile.h"
#include "TTree.h"
#include "TBranch.h"
#include "TEventList.h"
#include "TChain.h"

/*Run on the grid:
  prun --exec "echo %IN > input.txt; root -l -q -b 'thinGrid_chain.C+(\"input.txt\")'" --outputs _thin.root --match "*D3PD*"
  --athenaTag=15.6.9.12,32,setup,AtlasProduction --inDS --outDS --noBuild
*/
void thinGrid_chain(string strInFile, string strBranchNameTxt = "bNames.txt", int dpdType=1,string strOutFile = "_thin.root"){
  

  std::string argStr;
  std::ifstream ifs(strInFile.c_str());
  std::getline(ifs,argStr);
  // split by ','
  std::vector<std::string> fileList;
  for (size_t i=0,n; i <= argStr.length(); i=n+1)
    {
      n = argStr.find_first_of(',',i);
      if (n == string::npos)
        n = argStr.length();
      string tmp = argStr.substr(i,n-i);
      fileList.push_back(tmp);
    }
  TChain *chain = new TChain("physics");
    
  for (int nFile=0; nFile < fileList.size();nFile++){
    chain->Add(fileList[nFile].c_str());
  }
  chain->LoadTree(0);
  chain->SetBranchStatus("*",0);    
  cout << "Opening branch names: " << strBranchNameTxt << endl;
  ifstream bNames (strBranchNameTxt.c_str());
      
  while (!bNames.eof()){
    string line;
    getline (bNames,line);
    if(line!=""){
      chain->SetBranchStatus(line.c_str(),1);
    }
  }
  TFile *newfile = new TFile(strOutFile.c_str(),"recreate");
  cout <<"file"<< endl;
  TTree *newtree = (TTree*)chain->GetTree()->CloneTree(0);
  cout<<" tree"<< endl;
  TEventList *evlist=new TEventList("evlist");
  string strThinCut = "";
  if (dpdType == 1) {
    strThinCut = "((mu_staco_n > 0  && mu_staco_pt[0]>10000)||(mu_muid_n > 0  && mu_muid_pt[0]>10000)) && el_n > 0 && el_pt>10000";
    //strThinCut = "((mu_staco_n > 0  )||(mu_muid_n > 0  )) && el_n > 0";
  } else {
    strThinCut = "mu_n > 0 && el_n > 0";
  }
  cout<<"cut "<< endl;
  chain->Draw(">>evlist",strThinCut.c_str());
  int evPassed=evlist->GetN();  
  cout <<"evlist "<< endl;
  cout << "Number of events after Skimming/Thinning: " << evPassed << endl;
  for(int i=0;i<evPassed;i++){
    chain->GetEntry(evlist->GetEntry(i));
    newtree->Fill();
  }
  cout <<"filled"<< endl;
  delete chain;
  delete evlist;
  newfile->Write();
  delete newfile;
}