TTree* filter_tree(TTree* t, const char* selection, const char* vars, TFile* fout);

void filter_tree(const char* ntuple, const char* files, const char* selection, const char* vars, const char* fout);


TTree* filter_tree(TTree* t, const char* selection, const char* vars, TFile* fout)
{

  // ttf implements the event selection
  TTreeFormula ttf("filter_tree",selection,t);

  // turn off all branches in input
  t->SetBranchStatus("*",0);

  // tokenize array of variable names
  TString s(vars);
  TObjArray* vlist = s.Tokenize(", ");
  // loop over variable names, turning on associated branch in input
  std::cout<<"filter_tree: writing out the following variables:"<<std::endl;
  for(int i=0; i<vlist->GetEntries(); i++){
    TObjString* obj = dynamic_cast<TObjString*>(vlist->At(i));
    TString v=obj->GetString();
    std::cout<<v<<std::endl;
    t->SetBranchStatus(v.Data(),1);
  }
  std::cout<<"end of variables"<<std::endl;

  fout->cd();
  TTree* tc = t->CloneTree(0);

  // go through and turn on branches needed by selection  
  std::cout<<"Turning on branches"<<std::endl;
  for(Int_t j=0; j< ttf.GetNcodes(); j++){
    std::cout<<"Enabling branch : "<<ttf.GetLeaf(j)->GetBranch()->GetName()<<std::endl;
    t->SetBranchStatus(ttf.GetLeaf(j)->GetBranch()->GetName(),1);
  }
  ttf.UpdateFormulaLeaves();

  Long64_t nbr=0;
  std::cout.precision(3);
  Long64_t i=0;
  Long64_t npass=0;
  bool keep_going=true;
  Int_t current_tree=-1;

  while(true) {
    if(t->GetTreeNumber()!=current_tree) {
      current_tree=t->GetTreeNumber(); 
      ttf.UpdateFormulaLeaves();
    }
    // load branches needed by selection
    for(Int_t j=0; j< ttf.GetNcodes(); j++){
      Long64_t nb = ttf.GetLeaf(j)->GetBranch()->GetEntry(i);
      if(nb<=0){
	// done reading
	std::cout<<"Should be done (1)? : [i,j] = "<<i<<" , "<<j<<std::endl;
	keep_going=false;
	break;
      }
      nbr+=nb;
    }
    if(!keep_going) break;
    // evaluate the selection, if true, fill output
    Double_t val = ttf.EvalInstance();
    if(val > 0.0){
      npass++;
      Long64_t nb=t->GetEntry(i,0);
      if(nb<=0){
	// done reading
	std::cout<<"Should be done (2)?"<<std::endl;
	keep_going=false;
	continue;
      }
      nbr+=nb;      
      tc->Fill();
    }
    // how's it going
    if(i%100000==0){
      std::cout<<i<<" entries  read : "<<npass<<" filled : "
      <<nbr/(1024.0*1024.0)<<" Mb read"<<std::endl;
    }
    
    i++; // don't forget
  }
  std::cout<<"Done                                              "<<std::endl;
  std::cout<<"filter_tree: output has "
	   <<npass<<" entries."<<std::endl;  
  return tc;
}

void filter_tree(const char* ntuple, const char* files, const char* selection, const char* vars, const char* fout){
  std::cout<<"filter_tree: building chain of '"
	   <<ntuple<<"' ntuples from "<<files<<std::endl;
  TChain c(ntuple,ntuple);
  c.Add(files);
  /*
  std::cout<<"filter_tree: using "<<c.GetNtrees()
	   <<" trees with a total of "
	   <<c.GetEntries()<<" entries."<<std::endl;
  */
  std::cout<<"filter_tree: using "<<c.GetNtrees()<<" trees."<<std::endl;
  
  TFile f(fout,"recreate");
  filter_tree(&c,selection,vars,&f);
  
  f.Write();
  
}