// Example showing usage of JetsWithoutJets classes.
//
// Compile it with "make example" and run it with
//
//   ./example < ../data/single-event.dat
//
// Copyright (c) 2013
// Daniele Bertolini and Jesse Thaler
//
// $Id: $
//----------------------------------------------------------------------
// This file is part of FastJet contrib.
//
// It is free software; you can redistribute it and/or modify it under
// the terms of the GNU General Public License as published by the
// Free Software Foundation; either version 2 of the License, or (at
// your option) any later version.
//
// It is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
// or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
// License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this code. If not, see <http://www.gnu.org/licenses/>.
//----------------------------------------------------------------------

#include <iomanip>
#include <iostream>
#include <sstream>
#include <ctime>


#include "fastjet/PseudoJet.hh"
#include "fastjet/ClusterSequence.hh"
#include "fastjet/JetDefinition.hh"
#include "fastjet/tools/Filter.hh"

#include "JetsWithoutJets.hh" // In external code, this should be fastjet/contrib/JetsWithoutJets.hh

using namespace std;
using namespace fastjet;
using namespace fastjet::contrib;

// forward declaration to make things clearer
void read_event(vector<PseudoJet> &event);
void analyze(const vector<PseudoJet> & input_particles);

//----------------------------------------------------------------------
int main(){

  //----------------------------------------------------------
  // read in input particles
  vector<PseudoJet> event;
  read_event(event);
  cout << "#########"  << endl;
  cout << "## Read an event with " << event.size() << " particles" << endl;

  //----------------------------------------------------------
  // illustrate how this JetsWithoutJets contrib works

  analyze(event);
  return 0;
}

// read in input particles
void read_event(vector<PseudoJet> &event){  
  string line;
  while (getline(cin, line)) {
    istringstream linestream(line);
    // take substrings to avoid problems when there are extra "pollution"
    // characters (e.g. line-feed).
    if (line.substr(0,4) == "#END") {return;}
    if (line.substr(0,1) == "#") {continue;}
    double px,py,pz,E;
    linestream >> px >> py >> pz >> E;
    PseudoJet particle(px,py,pz,E);

    // push event onto back of full_event vector
    event.push_back(particle);
  }
}

// Main code
void analyze(const vector<PseudoJet> & input_particles) {
  
   // Jet parameters. 
   double Rjet = 1.0;
   double pTcut = 200.0;

   // Subjet parameters. Need for trimming.
   double Rsub = 0.2;
   double fcut = 0.05;
   double ptsubcut = 50.0;  // additional ptsubcut for subjet multiplicity

   //////////
   // Basic event shape analysis.
   //////////
   
   // Selectors of the type SelectorEventTrimmer are trimmers of the event.
   Selector trimmer=SelectorShapeTrimming(Rjet,pTcut,Rsub,fcut);
   
   // Classes derived from JetLikeEventShape take as input a vector<PseudoJet>
   // representing the event, and return a number.
   // They have two constructors: you can specify Rjet and pTcut solely or
   // in addition to those also Rsub and fcut (if you want to implement
   // built-in trimming at the same time).
   
   // N.B.: trimming an event and then calculating the event shape is NOT
   // the same as calculating the trimmed event shape.  See the README for
   // more information.

   // ShapeJetMultiplicity = jet counting
   ShapeJetMultiplicity Nj(Rjet,pTcut);
   ShapeJetMultiplicity Nj_trim(Rjet,pTcut,Rsub,fcut);  // with trimming
   
   // ShapeScalarPt = HT (summed jet pT)
   ShapeScalarPt HT(Rjet,pTcut);
   ShapeScalarPt HT_trim(Rjet,pTcut,Rsub,fcut); // with trimming
   
   // ShapeMissingPt = missing pT
   ShapeMissingPt HTmiss(Rjet,pTcut);
   ShapeMissingPt HTmiss_trim(Rjet,pTcut,Rsub,fcut); // with trimming
   
   cout << setprecision(6);
   cout << "#########" << endl;
   cout << "## Example of basic event shape analysis" << endl;
   cout << "#########" << endl;
   cout << "Jet parameters: R_jet=" << Rjet << ", pTcut=" << pTcut <<endl;
   cout << "Trimming parameters: R_sub=" << Rsub << ", fcut=" << fcut <<endl;
   cout << "#####" << endl;
   cout << "N_jet=" << Nj(input_particles) << endl;
   cout << "N_jet_trim (using built-in)=" << Nj_trim(input_particles) << endl;
   cout << "N_jet_trim (using trimmer)=" << Nj(trimmer(input_particles)) << endl; // N.B.: Not exactly equivalent
   cout << "#####" << endl;
   cout << "H_T=" << HT(input_particles) << endl;
   cout << "H_T_trim (using built-in)=" << HT_trim(input_particles) << endl;
   cout << "H_T_trim (using trimmer)=" << HT(trimmer(input_particles)) << endl; // N.B.: Not exactly equivalent
   cout << "#####" << endl;
   cout << "Missing H_T=" << HTmiss(input_particles) << endl;
   cout << "Missing H_T_trim (using built-in)=" << HTmiss_trim(input_particles) << endl;
   cout << "Missing H_T_trim (using trimmer)=" << HTmiss(trimmer(input_particles)) << endl; // N.B.: Not exactly equivalent
   cout << "#####" << endl;

   //////////
   // Testing different trimming methods
   //////////

   // Different trimming mechanisms
   Selector eventShapeTrimmer=SelectorShapeTrimming(Rjet,pTcut,Rsub,fcut);
   JetShapeTrimmer jetShapeTrimmer(Rsub,fcut);
   Filter treeTrimmer(Rsub, SelectorPtFractionMin(fcut));

   // Clustering with anti-kt algorithm
   JetAlgorithm algorithm = antikt_algorithm; 
   JetDefinition jetDef = JetDefinition(algorithm,Rjet);
   ClusterSequence clust_seq(input_particles,jetDef);
   PseudoJet hardestJet  = sorted_by_pt(clust_seq.inclusive_jets(pTcut))[0];

   // Clustering after applying event shape trimming
   ClusterSequence clust_seq_estrim(eventShapeTrimmer(input_particles),jetDef);
   PseudoJet hardestJet_estrim  = sorted_by_pt(clust_seq_estrim.inclusive_jets(pTcut))[0];

   double untrimmedMass = hardestJet.m();
   double treeTrimmedMass = treeTrimmer(hardestJet).m();
   double jetShapeTrimmedMass = jetShapeTrimmer(hardestJet).m();
   double eventShapeTrimmedMass = hardestJet_estrim.m();
   
   cout << setprecision(6);
   cout << "#########" << endl;
   cout << "## Example of different trimming methods" << endl;
   cout << "#########" << endl;
   cout << "Anti-kT jet parameters: R_jet=" << Rjet << ", pTcut=" << pTcut <<endl;
   cout << "Trimming parameters: R_sub=" << Rsub << ", fcut=" << fcut <<endl;
   cout << "Untrimmed Mass = " <<  untrimmedMass << endl;
   cout << "Tree Trimmed Mass = " <<  treeTrimmedMass << endl;
   cout << "Jet Shape Trimmed Mass = " <<  jetShapeTrimmedMass << endl;
   cout << "Event Shape Trimmed Mass = " <<  eventShapeTrimmedMass << endl;
   cout << "#####" << endl;

   //////////
   // Variable pT analysis.
   //////////

   // These use classes derived from JetLikeEventShape_VariablePtCut, where 
   // the full pT information is kept in calculating the event shape.
   // One can calculate the value of the event shape at a given pT cut, or
   // one can calculate the inverse function to find the value of the pT cut
   // needed to return a given value of the event shape.
  
   ShapeJetMultiplicity_VariablePtCut Nj_allPt(Rjet);
   ShapeJetMultiplicity_VariablePtCut Nj_allPt_trim(Rjet,Rsub,fcut); // with trimming
   Nj_allPt.set_input(input_particles);
   Nj_allPt_trim.set_input(input_particles);
   
   cout << "#########" << endl;
   cout << "## Example of variable pT analysis" << endl;
   cout << "#########" << endl;
   cout << "Jet parameters: R_jet=" << Rjet << ", pTcut=" << pTcut << " (unless otherwise stated)" <<endl;
   cout << "Trimming parameters: R_sub=" << Rsub << ", fcut=" << fcut <<endl;
   cout << "#####" << endl;
   cout << "N_jet from full pT_cut function=" <<  Nj_allPt.eventShapeFor(pTcut) << endl;
   cout << "N_jet_trim from full pT_cut function=" << Nj_allPt_trim.eventShapeFor(pTcut) << endl;
   cout << "#####" << endl;
   cout << "Variable pT_cut" << endl;
   double pTcut_values[] = {10,20,50,100,150}; 
   for (unsigned int i=0; i<5; i++){
   	cout << "pT_cut=" << pTcut_values[i] << ", N_jet=" << Nj_allPt.eventShapeFor(pTcut_values[i]) << endl;
   }
   cout << "#####" << endl;
   cout << "pT of the hardest jet=" <<  Nj_allPt.ptCutFor(1) << endl;
   cout << "pT of the hardest trimmed jet=" << Nj_allPt_trim.ptCutFor(1) << endl;
   cout << "#####" << endl;

   // Can get the full arrays too, though these are not output on the example.ref file
   std::vector<double> ptval =  Nj_allPt.ptCutArray();
   std::vector<double> NjPtval =  Nj_allPt.eventShapeArray();

   //////////
   // Variable R analysis.
   //////////

   // These use the class ShapeJetMultiplicity_VariableR (at present there is no
   // general class for variable R observables because of the high computational time). 
   // Here, the full R information is kept.

   ShapeJetMultiplicity_VariableR Nj_allR(pTcut);
   ShapeJetMultiplicity_VariableR Nj_allR_trim(pTcut,Rsub,fcut); // with trimming
   Nj_allR.set_input(input_particles);
   Nj_allR_trim.set_input(input_particles);
   
   cout << "#########" << endl;
   cout << "## Example of variable R analysis" << endl;
   cout << "#########" << endl;
   cout << "Jet parameters: R_jet=" << Rjet << " (unless otherwise stated), pTcut=" << pTcut <<endl;
   cout << "Trimming parameters: R_sub=" << Rsub << ", fcut=" << fcut <<endl;
   cout << "#####" << endl;
   cout << "N_jet from full R function=" <<  Nj_allR.eventShapeFor(Rjet) << endl;
   cout << "N_jet_trim from full R function=" <<  Nj_allR_trim.eventShapeFor(Rjet) << endl;
   cout << "#####" << endl;
   cout << "Variable Rjet" << endl;
   double Rjet_values[] = {0.2,0.4,0.6,0.8,1}; 
   for(unsigned int i=0; i<5; i++){
   	cout << "R=" << Rjet_values[i] << ", N_jet=" << Nj_allR.eventShapeFor(Rjet_values[i]) << endl;
   }
   
   // Can get the full arrays too, though these are not output on the example.ref file
   std::vector<double> Rval =  Nj_allR.RArray();
   std::vector<double> NjRval =  Nj_allR.eventShapeArray();


   //////////
   // Finding Jet Axes
   //////////

   EventShapeDensity_JetAxes axes_eventShape(Rjet,pTcut);
   axes_eventShape.set_input(input_particles);  //Loads particles and finds axes and weights

   //Get jet axes (by default sorted by pt). Pt of each axis corresponds to its pt weight. 	
   vector<PseudoJet> axes = axes_eventShape.axes();
   //Get N_jet weights
   vector<double> Njw = axes_eventShape.Njet_weights();
   
   //If you want turn on global consistency and recalculate axes and weights
   axes_eventShape.setGlobalConsistencyCheck(true);
   axes_eventShape.find_axes_and_weights(); // uses stored information to save time; don't need to set_input again
   vector<PseudoJet> axes_gc = axes_eventShape.axes();
   vector<double> Njw_gc = axes_eventShape.Njet_weights();
   
   //Output hardest jet eta/phi/pT/N_jet weight
   cout << "#########" << endl;
   cout << "## Hardest jet axis and weights" << endl;
   cout << "#########" << endl;
   cout << "Without global consistency:" << endl;
   cout << "Rap_hardest_jet=" <<  axes[0].rap() <<endl;
   cout << "Phi_hardest_jet=" <<  axes[0].phi() <<endl;
   cout << "pt_weight_hardest_jet=" <<  axes[0].pt() <<endl;
   cout << "N_jet_weight_hardest_jet=" <<  Njw[0] <<endl;
   cout << "#" << endl;
   cout << "With global consistency:" << endl;
   cout << "Rap_hardest_jet=" <<  axes_gc[0].rap() <<endl;
   cout << "Phi_hardest_jet=" <<  axes_gc[0].phi() <<endl;
   cout << "pt_weight_hardest_jet=" <<  axes_gc[0].pt() <<endl;
   cout << "N_jet_weight_hardest_jet=" <<  Njw_gc[0] <<endl;

   //////////
   // Showing a wide range of event shapes
   //////////
   
   // This emphasizes that the jet-like event shapes are all derived from 
   // MyFunctionOfVectorOfPseudoJets
   vector<JetLikeEventShape*> myShapes;
   vector<string > myNames;
   
   myShapes.push_back( new ShapeJetMultiplicity(Rjet,pTcut) );
   myNames.push_back( "Njet");
   myShapes.push_back( new ShapeJetMultiplicity(Rjet,pTcut,Rsub,fcut) ); //with trimming
   myNames.push_back( "Njet_trim");

   myShapes.push_back( new ShapeScalarPt(Rjet,pTcut) );
   myNames.push_back( "HT");
   myShapes.push_back( new ShapeScalarPt(Rjet,pTcut,Rsub,fcut) ); //with trimming
   myNames.push_back( "HT_trim");

   myShapes.push_back( new ShapeMissingPt(Rjet,pTcut) );
   myNames.push_back( "MissHT");
   myShapes.push_back( new ShapeMissingPt(Rjet,pTcut,Rsub,fcut) ); //with trimming
   myNames.push_back( "MissHT_trim");
   
   myShapes.push_back( new ShapeSummedMass(Rjet,pTcut) );
   myNames.push_back( "SigmaM");
   myShapes.push_back( new ShapeSummedMass(Rjet,pTcut,Rsub,fcut) ); //with trimming
   myNames.push_back( "SigmaM_trim");

   myShapes.push_back( new ShapeSummedMassSquared(Rjet,pTcut) );
   myNames.push_back( "SigmaM2");
   myShapes.push_back( new ShapeSummedMassSquared(Rjet,pTcut,Rsub,fcut) ); //with trimming
   myNames.push_back( "SigmaM2_trim");

   myShapes.push_back( new ShapeTrimmedSubjetMultiplicity(Rjet,pTcut,Rsub,fcut,ptsubcut) ); //with trimming
   myNames.push_back( "SigmaNsub_trim");
   
   for (int n = -2; n <= 2; n++) {
      myShapes.push_back( new ShapeScalarPtToN(n,Rjet,pTcut) );
      myShapes.push_back( new ShapeScalarPtToN(n,Rjet,pTcut,Rsub,fcut) ); //with trimming

      stringstream myStream;
      myStream << "GenHT_" << n;
      myNames.push_back( myStream.str());
      myNames.push_back( myStream.str() + "_trim");
   }
     
   // outputing the variables
   cout << "#########" << endl;
   cout << "## Examples showing more general shapes" << endl;
   cout << "#########" << endl;
   cout << "Jet parameters: R_jet=" << Rjet << ", pTcut=" << pTcut << endl;
   cout << "Trimming parameters: R_sub=" << Rsub << ", fcut=" << fcut << endl;
   cout << "#####" << endl;
   for (unsigned int i = 0; i < myShapes.size(); i++) {
      cout << myShapes.at(i)->description() << endl; // ideally each even shape should have a description
      
      myShapes.at(i)->setUseStorageArray(false);
      double valueNoStorage = myShapes.at(i)->result(input_particles);
      myShapes.at(i)->setUseStorageArray(true);
      double valueYesStorage = myShapes.at(i)->result(input_particles);
            
      cout << myNames.at(i) << "="
           << valueNoStorage << endl;
      cout << "(Storage array works? " << (valueNoStorage == valueYesStorage ? "Yes" : "No!!") << ")" << endl;
      cout << "#" << endl;
           
      delete myShapes.at(i);
   }
   cout << "#####" << endl;

   
   // timing tests for the developers
   double do_timing_test = false;
   if (do_timing_test) {
      clock_t clock_begin, clock_end;
      double num_iter;

      cout << setprecision(6);

      num_iter = 1000;
      clock_begin = clock();
      ShapeScalarPt NjetA(Rjet,pTcut);
      cout << "timing " << NjetA.description() << endl;
      for (int t = 0; t < num_iter; t++) {
         NjetA(input_particles);
      }
      clock_end = clock();
      cout << "Method A: " << (clock_end-clock_begin)/double(CLOCKS_PER_SEC*num_iter)*1000 << " ms per Njet"<< endl;
      
      num_iter = 1000;
      clock_begin = clock();
      JetLikeEventShape NjetB(new FunctorJetScalarPt(),Rjet,pTcut);
      NjetB.setUseStorageArray(false);
      cout << "timing " << NjetB.description() << endl;
      for (int t = 0; t < num_iter; t++) {
         NjetB(input_particles);
      }
      clock_end = clock();
      cout << "Method B: " << (clock_end-clock_begin)/double(CLOCKS_PER_SEC*num_iter)*1000 << " ms per Njet"<< endl;
      
      
   }

}