// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/WFinder.hh"
#include "Rivet/Projections/VetoedFinalState.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/HeavyHadrons.hh"

namespace Rivet {


  /// @brief ATLAS W+b measurement
  class ATLAS_Wbb: public Analysis {
  public:

    ATLAS_Wbb(string name = "ATLAS_Wbb")
      : Analysis(name)
    {
      // the electron mode is used by default
      _mode = 1;
    }


    void init() {
      FinalState fs;
      addProjection(fs, "FinalState");

      Cut cuts = Cuts::etaIn(-2.5,2.5) & (Cuts::pT >= 25*GeV);

      // W finder for electrons and muons
      WFinder wf(fs, cuts, _mode==3? PID::MUON : PID::ELECTRON, 0.0*GeV, MAXDOUBLE, 0.0, 0.1,
                 WFinder::CLUSTERNODECAY, WFinder::NOTRACK, WFinder::TRANSMASS);
      addProjection(wf, "WF");

      // jets
		VetoedFinalState jet_fs(fs);
      jet_fs.addVetoOnThisFinalState(getProjection<WFinder>("WF"));
      FastJets jetpro1(jet_fs, FastJets::ANTIKT, 0.4);
      jetpro1.useInvisibles();
      addProjection(jetpro1, "AntiKtJets04");
      
      // book histograms
      _njet     = bookHisto1D(1, 1, _mode); // dSigma / dNjet
      _jet1_bPt = bookHisto1D(2, 1, _mode); // dSigma / dBjetPt for Njet = 1
      _jet2_bPt = bookHisto1D(2, 2, _mode); // dSigma / dBjetPt for Njet = 2


		_h_hadron_pT = bookHisto1D("hadlev_B_hadron_pT", 50,0,100);
		_h_matching_hadron_pT = bookHisto1D("hadlev_B_hadron_pT_matching_jet", 50,0,100);
		_h_bjet_pT_single = bookHisto1D("hadlev_bjet_pT_single_hadron_0GeV", 20,0,200);
		_h_bjet_pT_double = bookHisto1D("hadlev_bjet_pT_double_hadron_0GeV", 20,0,200);
		_h_bjet_pT_single_5GeV = bookHisto1D("hadlev_bjet_pT_single_hadron_5GeV", 20,0,200);
		_h_bjet_pT_double_5GeV = bookHisto1D("hadlev_bjet_pT_double_hadron_5GeV", 20,0,200);
		_h_bjet_num_matching = bookHisto1D("hadlev_bjet_number_matching_B_hadrons", 6,-0.5,5.5);
		_h_bjet_num_matching_5GeV = bookHisto1D("hadlev_bjet_number_matching_B_hadrons_5GeV", 6,-0.5,5.5);

		_h_bjet_bb_over_b = bookScatter2D("hadlev_bjet_bb_over_b_hadron");
		_h_bjet_bb_over_b_5GeV = bookScatter2D("hadlev_bjet_bb_over_b_hadron_5GeV");


		_h_quark_pT = bookHisto1D("quarklev_b_quark_pT", 50,0,100);
		_h_matching_quark_pT = bookHisto1D("quarklev_b_quark_pT_matching_jet", 50,0,100);
		_h_bjet_pT_single_quark = bookHisto1D("quarklev_bjet_pT_single_quark_0GeV", 20,0,200);
		_h_bjet_pT_double_quark = bookHisto1D("quarklev_bjet_pT_double_quark_0GeV", 20,0,200);
		_h_bjet_pT_single_quark_5GeV = bookHisto1D("quarklev_bjet_pT_single_quark_5GeV", 20,0,200);
		_h_bjet_pT_double_quark_5GeV = bookHisto1D("quarklev_bjet_pT_double_quark_5GeV", 20,0,200);
		_h_bjet_num_matching_quark = bookHisto1D("quarklev_bjet_number_matching_b_quarks", 6,-0.5,5.5);
		_h_bjet_num_matching_quark_5GeV = bookHisto1D("quarklev_bjet_number_matching_b_quarks_5GeV", 6,-0.5,5.5);

		_h_bjet_bb_over_b_quark = bookScatter2D("quarklev_bjet_bb_over_b_quarks");
		_h_bjet_bb_over_b_quark_5GeV = bookScatter2D("quarklev_bjet_bb_over_b_quarks_5GeV");



    }



	  //=========================================================



    void analyze(const Event& event) {

      const double weight = event.weight();

      //  retrieve W boson candidate
      const WFinder& wf = applyProjection<WFinder>(event, "WF");
      if( wf.bosons().size() != 1 )  vetoEvent; // only one W boson candidate
      if( !(wf.mT() > 60.0*GeV) )    vetoEvent;
      //const Particle& Wbboson  = wf.boson();


      // retrieve constituent neutrino
      const Particle& neutrino = wf.constituentNeutrino();
      if( !(neutrino.pT() > 25.0*GeV) )  vetoEvent;

      // retrieve constituent lepton
      const Particle& lepton = wf.constituentLepton();
      if( !(lepton.pT() > 25.0*GeV) )  vetoEvent;



      //---------------------------
	   Particles weakBs, bquarks;	
		// Loop over the unstable particles
		foreach (const Particle& p, particles(event.genEvent()) ) {
			const PdgId pid = p.pid();
			//look for b quarks:
			if(abs(pid)==5) {
				bool good_B = true;
				const HepMC::GenParticle* pgen = p.genParticle();
				HepMC::GenVertex* vgen = pgen -> end_vertex();
				if(vgen) {
					//make sure this is the last b in the chain:
					for (HepMC::GenVertex::particles_out_const_iterator it = vgen->particles_out_const_begin(); it !=  vgen->particles_out_const_end(); ++it) {
						if (  abs( (*it)->pdg_id() ) ==5 ) {
							good_B = false;
						}
					}
				}
				if(good_B) {
					bquarks += p;
				}
			}
			// Look for B hadrons:
			if (PID::hasBottom(pid)) {
				bool good_B = true;
				const HepMC::GenParticle* pgen = p.genParticle();
				HepMC::GenVertex* vgen = pgen -> end_vertex();
				if(vgen) {
					// make sure there are no more B hadrons in the chain:
					for (HepMC::GenVertex::particles_out_const_iterator it = vgen->particles_out_const_begin(); it !=  vgen->particles_out_const_end(); ++it) {
						if ( PID::hasBottom( (*it)->pdg_id() ) ) {
							good_B = false;
						}
					}
				}
				if ( good_B ) {
					weakBs += p;
				}
			}
		
		}//loop over particles


      if (weakBs.size() <1 && bquarks.size() <1 ) vetoEvent;




      //---------------------------
      // -- get the b-jets:
      const Jets& jets = applyProjection<JetAlg>(event, "AntiKtJets04").jetsByPt(Cuts::pT >25.0*GeV && Cuts::absrap <2.1);
      
      int goodjets = 0;
      Jets good_jets, b_jets, b_quark_jets;
		vector<bool> hadron_used;
		for(uint i=0; i<weakBs.size(); i++) hadron_used.push_back(false);
      foreach(const Jet& jet, jets) {
			//reapply cuts, not sure why this is not working...
			if(jet.pT()<25*GeV) continue;
			if(fabs(jet.rap())>2.1) continue;
			
			int nmatch=0, nmatch_5GeV=0, nmatch_quark=0, nmatch_quark_5GeV=0;

			//veto overlaps with W lepton:
			if(deltaR(jet, lepton) < 0.5)  continue;

			good_jets.push_back(jet);
			goodjets++;

			int iB=-1;
			foreach(const Particle& bhadron, weakBs) {
				iB++;
				if(hadron_used[iB]) continue;
				if( deltaR(jet, bhadron) <= 0.3 ) {
					hadron_used[iB]=true;
					nmatch++;
					if(bhadron.pT()>5*GeV) {
						if(nmatch_5GeV==0) b_jets.push_back(jet);
						nmatch_5GeV++;
					}
				}
	      } // end loop on b-hadrons  
			foreach(const Particle& bquark, bquarks) {
				if( deltaR(jet, bquark) <= 0.3 ) {
					nmatch_quark++;
					if(bquark.pT()>5*GeV) nmatch_quark_5GeV++;
				}
	      } // end loop on b-quarks  
			

			_h_bjet_num_matching->fill(nmatch, weight);
			_h_bjet_num_matching_5GeV->fill(nmatch_5GeV, weight);
			if(nmatch == 1) _h_bjet_pT_single->fill(jet.pT()/GeV, weight);
			if(nmatch > 1) _h_bjet_pT_double->fill(jet.pT()/GeV, weight);
			if(nmatch_5GeV == 1) _h_bjet_pT_single_5GeV->fill(jet.pT()/GeV, weight);
			if(nmatch_5GeV > 1) _h_bjet_pT_double_5GeV->fill(jet.pT()/GeV, weight);

			_h_bjet_num_matching_quark->fill(nmatch_quark, weight);
			_h_bjet_num_matching_quark_5GeV->fill(nmatch_quark_5GeV, weight);
			if(nmatch_quark == 1) _h_bjet_pT_single_quark->fill(jet.pT()/GeV, weight);
			if(nmatch_quark > 1) _h_bjet_pT_double_quark->fill(jet.pT()/GeV, weight);
			if(nmatch_quark_5GeV == 1) _h_bjet_pT_single_quark_5GeV->fill(jet.pT()/GeV, weight);
			if(nmatch_quark_5GeV > 1) _h_bjet_pT_double_quark_5GeV->fill(jet.pT()/GeV, weight);

      }
      


		//now loop the other way, to avoid double counting matched hadrons:
		foreach(const Particle& bhadron, weakBs) {
			if( fabs(bhadron.rapidity()) > 2.4) continue;
			 _h_hadron_pT->fill(bhadron.pT(), weight);
			bool match = false;
			foreach(const Jet& jet, good_jets) if( deltaR(jet, bhadron) <= 0.3 ) match=true;
			if( match ) _h_matching_hadron_pT->fill(bhadron.pT(), weight);
		}
		foreach(const Particle& bquark, bquarks) {
			if( fabs(bquark.rapidity()) > 2.4) continue;
			_h_quark_pT->fill(bquark.pT(), weight);
			bool match = false;
			foreach(const Jet& jet, good_jets) if( deltaR(jet, bquark) <= 0.3 ) match=true;
			if( match ) _h_matching_quark_pT->fill(bquark.pT(), weight);
		}







		//nw fill the data plots:
      if( goodjets > 2 )  vetoEvent; // at most two jets
      if( b_jets.size()==0 )  vetoEvent; // at least one of them b-tagged
		
      _njet->fill(goodjets, weight);
      _njet->fill(3. , weight);

      if(     goodjets == 1)  _jet1_bPt->fill(b_jets[0].pT()*GeV, weight);
      else if(goodjets == 2)  _jet2_bPt->fill(b_jets[0].pT()*GeV, weight);
    }


    void finalize() {

      // Print summary info
      const double xs_pb(crossSection() / picobarn);
      const double sumw(sumOfWeights());
      MSG_INFO("Cross-Section/pb: " << xs_pb      );
      MSG_INFO("Sum of weights  : " << sumw       );
      MSG_INFO("nEvents         : " << numEvents());

      const double normfac(xs_pb / sumw);

      scale(_njet,     normfac);
      scale(_jet1_bPt, normfac);
      scale(_jet2_bPt, normfac);


	  scale( _h_hadron_pT, normfac);
	  scale( _h_matching_hadron_pT, normfac);

	  scale( _h_bjet_pT_single, normfac);
	  scale( _h_bjet_pT_double, normfac);
	  scale( _h_bjet_pT_single_5GeV, normfac);
	  scale( _h_bjet_pT_double_5GeV, normfac);
	  scale( _h_bjet_num_matching, normfac);
	  scale( _h_bjet_num_matching_5GeV, normfac);

	  divide( _h_bjet_pT_double, _h_bjet_pT_single, _h_bjet_bb_over_b);
	  divide( _h_bjet_pT_double_5GeV, _h_bjet_pT_single_5GeV, _h_bjet_bb_over_b_5GeV);

	  scale( _h_quark_pT, normfac);
	  scale( _h_matching_quark_pT, normfac);

	  scale( _h_bjet_pT_single_quark, normfac);
	  scale( _h_bjet_pT_double_quark, normfac);
	  scale( _h_bjet_pT_single_quark_5GeV, normfac);
	  scale( _h_bjet_pT_double_quark_5GeV, normfac);
	  scale( _h_bjet_num_matching_quark, normfac);
	  scale( _h_bjet_num_matching_quark_5GeV, normfac);

	  divide( _h_bjet_pT_double_quark, _h_bjet_pT_single_quark, _h_bjet_bb_over_b_quark);
	  divide( _h_bjet_pT_double_quark_5GeV, _h_bjet_pT_single_quark_5GeV, _h_bjet_bb_over_b_quark_5GeV);

    }

  protected:

    size_t _mode;

  private:

    Histo1DPtr _njet;
    Histo1DPtr _jet1_bPt;
    Histo1DPtr _jet2_bPt;


	  Histo1DPtr _h_hadron_pT;
	  Histo1DPtr _h_matching_hadron_pT;
	  Histo1DPtr _h_bjet_pT_single;
	  Histo1DPtr _h_bjet_pT_double;
	  Histo1DPtr _h_bjet_pT_single_5GeV;
	  Histo1DPtr _h_bjet_pT_double_5GeV;
	  Histo1DPtr _h_bjet_num_matching;
	  Histo1DPtr _h_bjet_num_matching_5GeV;
	
	  Scatter2DPtr _h_bjet_bb_over_b;
	  Scatter2DPtr _h_bjet_bb_over_b_5GeV;

	  Histo1DPtr _h_quark_pT;
	  Histo1DPtr _h_matching_quark_pT;
	  Histo1DPtr _h_bjet_pT_single_quark;
	  Histo1DPtr _h_bjet_pT_double_quark;
	  Histo1DPtr _h_bjet_pT_single_quark_5GeV;
	  Histo1DPtr _h_bjet_pT_double_quark_5GeV;
	  Histo1DPtr _h_bjet_num_matching_quark;
	  Histo1DPtr _h_bjet_num_matching_quark_5GeV;
	
	  Scatter2DPtr _h_bjet_bb_over_b_quark;
	  Scatter2DPtr _h_bjet_bb_over_b_quark_5GeV;


    bool _isMuon;

  };

  class ATLAS_Wbb_EL : public ATLAS_Wbb {
  public:
    ATLAS_Wbb_EL()
      : ATLAS_Wbb("ATLAS_Wbb_EL")
    {
      _mode = 2;
    }
  };

  class ATLAS_Wbb_MU : public ATLAS_Wbb {
  public:
    ATLAS_Wbb_MU()
      : ATLAS_Wbb("ATLAS_Wbb_MU")
    {
      _mode = 3;
    }
  };

  // The hook for the plugin system
  DECLARE_RIVET_PLUGIN(ATLAS_Wbb);
  DECLARE_RIVET_PLUGIN(ATLAS_Wbb_EL);
  DECLARE_RIVET_PLUGIN(ATLAS_Wbb_MU);

}