#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include "fastjet/ClusterSequenceArea.hh"
#include "UWEvent.hh"
#include "print_jets.hh"
#include "SimpleHist.hh"
#include "AverageAndError.hh"
#include "fastjet/Selector.hh"
#include "fastjet/tools/JetMedianBackgroundEstimator.hh"
//#include "fastjet/tools/GridMedianBackgroundEstimator.hh"
#include "fastjet/tools/Subtractor.hh" 

using namespace std;
using namespace fastjet;


//
// Run this program with
//    gunzip -c ../event-files/pythia-Zprime.UW.gz | ./04-subtraction
//

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 = "04-subtraction.hist";

  // max number of events read in (even if input file contains more)
  int maxiev = 1000;
  
  // filename containing the events
//  string datafile =  "./data/";
//  datafile += "pythia-dijets.UW";
  
  // antikt, R=0.6
  JetDefinition jet_def(antikt_algorithm,0.6);
  // kt, R=0.5
  JetDefinition jet_def_for_rho(kt_algorithm,0.5);

  // set up area definition
  AreaDefinition area_def(active_area,GhostedAreaSpec(rapmax + 1.));

  // initialize histograms
  SimpleHist inv_mass_hist(500,1500,100);  
  SimpleHist inv_mass_sub_hist(inv_mass_hist);  

  // initialize AverageAndError
  AverageAndError avg_rho;
  
  // initialize background estimator
  JetMedianBackgroundEstimator bge(sel_rapmax, jet_def_for_rho, area_def);
//  GridMedianBackgroundEstimator bge(rapmax, 0.55);
  
  // define subtractor     
  Subtractor sub(&bge);
  
  // 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
       ClusterSequenceArea cs(input_particles, jet_def, area_def);
//       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);

       // if there aren't two jets, skip event
       if (jets.size() < 2) {cout << "SKIPPING EVENT" << endl; continue;}
       
       // calculate invariant mass of two jets, add it to histogram
       PseudoJet sum = jets[0] + jets[1];
       inv_mass_hist.add_entry(sum.m());

       // pass the input particles to the backgrund estimator, collect rho for average
       bge.set_particles(input_particles);
       avg_rho.add(bge.rho());
       
       // calculate inv. mass again, but AFTER subtracting the jets
       sum = sub(jets[0]) + sub(jets[1]);
       inv_mass_sub_hist.add_entry(sum.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 << "# jet_def for rho: " << jet_def_for_rho.description() << endl;
  output_file << "# " << area_def.description() << endl;
  output_file << "# jets selected using: " << full_selector.description() << endl;
  output_file << "# background estimator: " << bge.description() << endl;
  output_file << "# average rho = " << avg_rho.average() << endl;
  output_file << "# jets subtracted with: " << sub.description() << endl;
  output(inv_mass_hist,&output_file); // in SimpleHist.hh
  output_file << endl << endl;
  output(inv_mass_sub_hist,&output_file); // in SimpleHist.hh
  output_file << endl << endl;

  output_file.close();
  
  
}