#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include "fastjet/ClusterSequence.hh"
#include "UWEvent.hh"
#include "print_jets.hh"
#include "SimpleHist.hh"
#include "AverageAndError.hh"
#include "fastjet/Selector.hh"
#include "fastjet/tools/MassDropTagger.hh"
#include "fastjet/tools/Filter.hh" 

using namespace std;
using namespace fastjet;


//
// Run this program with
//    gunzip -c ../event-files/pythia-ZZ-500.UW.gz | ./05-boosted-Z
//

int main() {

  // select the kind of jets to look at
  //
  double rapmax = 2.;
  Selector sel_rapmax = SelectorAbsRapMax(rapmax);
  // in each event select only 2 hardest jets, followed by selection of jets in |y| < 2
  Selector full_selector = sel_rapmax * SelectorNHardest(2);
  string outfile = "05-boosted-Z.hist";

  // max number of events read in (even if input file contains more)
  int maxiev = 2;
  
  // filename containing the events
//  string datafile =  "./data/";
//  datafile += "pythia-dijets.UW";
  
  // CA, R=1.2
  JetDefinition jet_def(cambridge_algorithm,1.2);

  // initialize histograms
  SimpleHist inv_mass_hist(50,150,100);  
  SimpleHist inv_mass_tag_hist(inv_mass_hist);  
  SimpleHist inv_mass_tag_filt_hist(inv_mass_hist);  

  // initialize average and error
  AverageAndError avg_pt;
  
  // define mass drop tagger
  MassDropTagger md_tagger(0.667, 0.09);
  
  // define a filter. Do it here, so that one can output its
  // description at the end.
  Filter filter;
      
  // read file containing events
//  ifstream input_file;
//  input_file.open(datafile.c_str());
  int iev=0;
  vector<PseudoJet> input_particles;
  while (readUWEvent(cin, input_particles) && ++iev <= maxiev) {
//  while (readUWEvent(input_file, input_particles) && ++iev <= maxiev) {
    cerr << "Event " << iev << ", size = " << input_particles.size() << endl; 

       // cluster the input particles
       ClusterSequence cs(input_particles, jet_def);
  
       // extract the jets
       vector<PseudoJet> jets = sorted_by_pt(cs.inclusive_jets());

       // select the jets to be used for in. mass evaluation
       jets = full_selector(jets);
       
       for(unsigned int i=0; i < jets.size(); i++) {

          //calculate the average pt of the two hardest jets
          avg_pt.add(jets[i].perp());

          // calculate the untagged invariant mass
          inv_mass_hist.add_entry(jets[i].m());

          // try to tag the jet
          PseudoJet tagged = md_tagger(jets[i]);
 
          // if tagging successful, go on to filter
          if ( tagged != 0 ) {
	  
             // calculate the invariant mass of the tagged jet
             inv_mass_tag_hist.add_entry(tagged.m());
	     
	     // now on to filter
             PseudoJet parent1 = tagged.pieces()[0];
             PseudoJet parent2 = tagged.pieces()[1];
             double   Rbb = parent1.delta_R(parent2);
             double   Rfilt = min(Rbb/2, 0.3); // somewhat arbitrary choice
             unsigned nfilt = 3;               // number of pieces we'll take
             
	     // print out in first event
	     if (iev < 1 ) {cout << "Subjet separation (Rbb) = " << Rbb << ", Rfilt = " << Rfilt << endl;}

             // define filter and use it
	     filter = Filter(Rfilt, SelectorNHardest(nfilt));
             PseudoJet filtered = filter(tagged);

             // calculate the invariant mass of the tagged and filtered jet
             inv_mass_tag_filt_hist.add_entry(filtered.m());
	  }    
       }
       
       // print out jets in first event
       //if (iev < 1 ) { print_jets(jets); } // in print_jets.hh
  }
    
  // open output file for histograms
  ofstream output_file;
  output_file.open(outfile.c_str());
      
  // output histograms
  output_file << "# jet_def: " << jet_def.description() << endl;
  output_file << "# jets selected using: " << full_selector.description() << endl;
  output_file << "# average pt = " << avg_pt.average() << endl;
  output_file << "# tagger: " << md_tagger.description() << endl;
  output_file << "# filter: " << filter.description() << endl;
  output_file << "#" << endl;
  output_file << "# no tagging, no filtering" << endl;  
  output(inv_mass_hist,&output_file); // in SimpleHist.hh
  output_file << endl << endl;
  output_file << "# tagging, no filtering" << endl;  
  output(inv_mass_tag_hist,&output_file); // in SimpleHist.hh
  output_file << endl << endl;
  output_file << "# both tagging and filtering" << endl;  
  output(inv_mass_tag_filt_hist,&output_file); // in SimpleHist.hh
  output_file << endl << endl;

  output_file.close();
  
}