// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/ZFinder.hh"
#include "Rivet/Projections/WFinder.hh"
#include "Rivet/Projections/UnstableParticles.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Math/LorentzTrans.hh"

namespace Rivet {


  class MC_VH2BB : public Analysis {
  public:

    /// @name Constructors etc.
    //@{

    /// Constructor
    MC_VH2BB()
      : Analysis("MC_VH2BB")
    {    }

    //@}


    /// @name Analysis methods
    //@{

    /// Book histograms and initialise projections before the run
    void init() {

      FinalState fs;
      Cut cut = Cuts::abseta < 3.5 && Cuts::pT > 25*GeV;
      ZFinder zeefinder(fs, cut, PID::ELECTRON, 65*GeV, 115*GeV, 0.2);
      declare(zeefinder, "ZeeFinder");
      ZFinder zmmfinder(fs, cut, PID::MUON, 65*GeV, 115*GeV, 0.2);
      declare(zmmfinder, "ZmmFinder");
      WFinder wefinder(fs, cut, PID::ELECTRON, 60*GeV, 100*GeV, 25*GeV, 0.2);
      declare(wefinder, "WeFinder");
      WFinder wmfinder(fs, cut, PID::MUON, 60*GeV, 100*GeV, 25*GeV, 0.2);
      declare(wmfinder, "WmFinder");

      declare(fs, "FinalState");
      declare(FastJets(fs, FastJets::ANTIKT, 0.4), "AntiKT04");
      declare(FastJets(fs, FastJets::ANTIKT, 0.5), "AntiKT05");
      declare(FastJets(fs, FastJets::ANTIKT, 0.6), "AntiKT06");

      /// Book histograms
      book(_h_jet_bb_Delta_eta ,"jet_bb_Delta_eta", 50, 0, 4);
      book(_h_jet_bb_Delta_phi ,"jet_bb_Delta_phi", 50, 0, 1);
      book(_h_jet_bb_Delta_pT ,"jet_bb_Delta_pT", 50,0, 500);
      book(_h_jet_bb_Delta_R ,"jet_bb_Delta_R", 50, 0, 5);
      book(_h_jet_b_jet_eta ,"jet_b_jet_eta", 50, -4, 4);
      book(_h_jet_b_jet_multiplicity ,"jet_b_jet_multiplicity", 11, -0.5, 10.5);
      book(_h_jet_b_jet_phi ,"jet_b_jet_phi", 50, 0, 1);
      book(_h_jet_b_jet_pT ,"jet_b_jet_pT", 50, 0, 500);
      book(_h_jet_H_eta_using_bb ,"jet_H_eta_using_bb", 50, -4, 4);
      book(_h_jet_H_mass_using_bb ,"jet_H_mass_using_bb", 50, 50, 200);
      book(_h_jet_H_phi_using_bb ,"jet_H_phi_using_bb", 50, 0, 1);
      book(_h_jet_H_pT_using_bb ,"jet_H_pT_using_bb", 50, 0, 500);
      book(_h_jet_eta ,"jet_eta", 50, -4, 4);
      book(_h_jet_multiplicity ,"jet_multiplicity", 11, -0.5, 10.5);
      book(_h_jet_phi ,"jet_phi", 50, 0, 1);
      book(_h_jet_pT ,"jet_pT", 50, 0, 500);
      book(_h_jet_VBbb_Delta_eta ,"jet_VBbb_Delta_eta", 50, 0, 4);
      book(_h_jet_VBbb_Delta_phi ,"jet_VBbb_Delta_phi", 50, 0, 1);
      book(_h_jet_VBbb_Delta_pT ,"jet_VBbb_Delta_pT", 50, 0, 500);
      book(_h_jet_VBbb_Delta_R ,"jet_VBbb_Delta_R", 50, 0, 8);

      book(_h_VB_eta ,"VB_eta", 50, -4, 4);
      book(_h_VB_mass ,"VB_mass", 50, 60, 110);
      book(_h_Z_multiplicity ,"Z_multiplicity", 11, -0.5, 10.5);
      book(_h_W_multiplicity ,"W_multiplicity", 11, -0.5, 10.5);
      book(_h_VB_phi ,"VB_phi", 50, 0, 1);
      book(_h_VB_pT ,"VB_pT", 50, 0, 500);

      book(_h_jet_bVB_angle_Hframe ,"jet_bVB_angle_Hframe", 50, 0, 1);
      book(_h_jet_bb_angle_Hframe ,"jet_bb_angle_Hframe", 50, 0, 1);
      book(_h_jet_bVB_cosangle_Hframe ,"jet_bVB_cosangle_Hframe", 50, -1, 1);
      book(_h_jet_bb_cosangle_Hframe ,"jet_bb_cosangle_Hframe", 50, -1, 1);
    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {
      const double JETPTCUT = 30*GeV;

      const ZFinder& zeefinder = apply<ZFinder>(event, "ZeeFinder");
      const ZFinder& zmmfinder = apply<ZFinder>(event, "ZmmFinder");
      const WFinder& wefinder = apply<WFinder>(event, "WeFinder");
      const WFinder& wmfinder = apply<WFinder>(event, "WmFinder");
      const Particles vectorBosons = zeefinder.bosons() + zmmfinder.bosons() + wefinder.bosons() + wmfinder.bosons();
      _h_Z_multiplicity->fill(zeefinder.bosons().size() + zmmfinder.bosons().size());
      _h_W_multiplicity->fill(wefinder.bosons().size() + wmfinder.bosons().size());

      const Jets jets = apply<FastJets>(event, "AntiKT04").jetsByPt(JETPTCUT);
      _h_jet_multiplicity->fill(jets.size());

      // Identify the b-jets
      Jets bjets;
      for (const Jet& jet : jets) {
        const double jetEta = jet.eta();
        const double jetPhi = jet.phi();
        const double jetPt = jet.pT();
        _h_jet_eta->fill(jetEta);
        _h_jet_phi->fill(jetPhi/2/M_PI);
        _h_jet_pT->fill(jetPt/GeV);

        if (jet.bTagged() && jet.pT() > JETPTCUT) {
          bjets.push_back(jet);
          _h_jet_b_jet_eta->fill(jetEta);
          _h_jet_b_jet_phi->fill(jetPhi/2/M_PI);
          _h_jet_b_jet_pT->fill(jetPt);
        }
      }
      _h_jet_b_jet_multiplicity->fill(bjets.size());

      // Plot vector boson properties
      for (const Particle& v : vectorBosons) {
        _h_VB_phi->fill(v.phi()/2/M_PI);
        _h_VB_pT->fill(v.pT());
        _h_VB_eta->fill(v.eta());
        _h_VB_mass->fill(v.mass());
      }

      // rest of analysis requires at least 1 b jets
      if(bjets.empty()) vetoEvent;

      // Construct Higgs candidates from pairs of b-jets
      for (size_t i = 0; i < bjets.size()-1; ++i) {
        for (size_t j = i+1; j < bjets.size(); ++j) {
          const Jet& jet1 = bjets[i];
          const Jet& jet2 = bjets[j];

          const double deltaEtaJJ = fabs(jet1.eta() - jet2.eta());
          const double deltaPhiJJ = deltaPhi(jet1.momentum(), jet2.momentum());
          const double deltaRJJ = deltaR(jet1.momentum(), jet2.momentum());
          const double deltaPtJJ = fabs(jet1.pT() - jet2.pT());
          _h_jet_bb_Delta_eta->fill(deltaEtaJJ);
          _h_jet_bb_Delta_phi->fill(deltaPhiJJ/M_PI);
          _h_jet_bb_Delta_pT->fill(deltaPtJJ);
          _h_jet_bb_Delta_R->fill(deltaRJJ);

          const FourMomentum phiggs = jet1.momentum() + jet2.momentum();
          _h_jet_H_eta_using_bb->fill(phiggs.eta());
          _h_jet_H_mass_using_bb->fill(phiggs.mass());
          _h_jet_H_phi_using_bb->fill(phiggs.phi()/2/M_PI);
          _h_jet_H_pT_using_bb->fill(phiggs.pT());

          for (const Particle& v : vectorBosons) {
            const double deltaEtaVH = fabs(phiggs.eta() - v.eta());
            const double deltaPhiVH = deltaPhi(phiggs, v.momentum());
            const double deltaRVH = deltaR(phiggs, v.momentum());
            const double deltaPtVH = fabs(phiggs.pT() - v.pT());
            _h_jet_VBbb_Delta_eta->fill(deltaEtaVH);
            _h_jet_VBbb_Delta_phi->fill(deltaPhiVH/M_PI);
            _h_jet_VBbb_Delta_pT->fill(deltaPtVH);
            _h_jet_VBbb_Delta_R->fill(deltaRVH);

            // Calculate boost angles
            const vector<double> angles = boostAngles(jet1.momentum(), jet2.momentum(), v.momentum());
            _h_jet_bVB_angle_Hframe->fill(angles[0]/M_PI);
            _h_jet_bb_angle_Hframe->fill(angles[1]/M_PI);
            _h_jet_bVB_cosangle_Hframe->fill(cos(angles[0]));
            _h_jet_bb_cosangle_Hframe->fill(cos(angles[1]));
          }

        }
      }
    }


    /// Normalise histograms etc., after the run
    void finalize() {
      scale(_h_jet_bb_Delta_eta, crossSection()/sumOfWeights());
      scale(_h_jet_bb_Delta_phi, crossSection()/sumOfWeights());
      scale(_h_jet_bb_Delta_pT, crossSection()/sumOfWeights());
      scale(_h_jet_bb_Delta_R, crossSection()/sumOfWeights());
      scale(_h_jet_b_jet_eta, crossSection()/sumOfWeights());
      scale(_h_jet_b_jet_multiplicity, crossSection()/sumOfWeights());
      scale(_h_jet_b_jet_phi, crossSection()/sumOfWeights());
      scale(_h_jet_b_jet_pT, crossSection()/sumOfWeights());
      scale(_h_jet_H_eta_using_bb, crossSection()/sumOfWeights());
      scale(_h_jet_H_mass_using_bb, crossSection()/sumOfWeights());
      scale(_h_jet_H_phi_using_bb, crossSection()/sumOfWeights());
      scale(_h_jet_H_pT_using_bb, crossSection()/sumOfWeights());
      scale(_h_jet_eta, crossSection()/sumOfWeights());
      scale(_h_jet_multiplicity, crossSection()/sumOfWeights());
      scale(_h_jet_phi, crossSection()/sumOfWeights());
      scale(_h_jet_pT, crossSection()/sumOfWeights());
      scale(_h_jet_VBbb_Delta_eta, crossSection()/sumOfWeights());
      scale(_h_jet_VBbb_Delta_phi, crossSection()/sumOfWeights());
      scale(_h_jet_VBbb_Delta_pT, crossSection()/sumOfWeights());
      scale(_h_jet_VBbb_Delta_R, crossSection()/sumOfWeights());

      scale(_h_VB_eta, crossSection()/sumOfWeights());
      scale(_h_VB_mass, crossSection()/sumOfWeights());
      scale(_h_Z_multiplicity, crossSection()/sumOfWeights());
      scale(_h_W_multiplicity, crossSection()/sumOfWeights());
      scale(_h_VB_phi, crossSection()/sumOfWeights());
      scale(_h_VB_pT, crossSection()/sumOfWeights());

      scale(_h_jet_bVB_angle_Hframe, crossSection()/sumOfWeights());
      scale(_h_jet_bb_angle_Hframe, crossSection()/sumOfWeights());
      scale(_h_jet_bVB_cosangle_Hframe, crossSection()/sumOfWeights());
      scale(_h_jet_bb_cosangle_Hframe, crossSection()/sumOfWeights());
    }


    /// This should take in the four-momenta of two b's (jets/hadrons) and a vector boson, for the process VB*->VBH with H->bb
    /// It should return the smallest angle between the virtual vector boson and one of the b's, in the rest frame of the Higgs boson.
    /// It should also return (as the second element of the vector) the angle between the b's, in the rest frame of the Higgs boson.
    vector<double> boostAngles(const FourMomentum& b1, const FourMomentum& b2, const FourMomentum& vb) {
      const FourMomentum higgsMomentum = b1 + b2;
      const FourMomentum virtualVBMomentum = higgsMomentum + vb;
      const LorentzTransform lt = LorentzTransform::mkFrameTransformFromBeta(higgsMomentum.betaVec());

      const FourMomentum virtualVBMomentumBOOSTED = lt.transform(virtualVBMomentum);
      const FourMomentum b1BOOSTED = lt.transform(b1);
      const FourMomentum b2BOOSTED = lt.transform(b2);

      const double angle1 = b1BOOSTED.angle(virtualVBMomentumBOOSTED);
      const double angle2 = b2BOOSTED.angle(virtualVBMomentumBOOSTED);

      const double anglebb = mapAngle0ToPi(b1BOOSTED.angle(b2BOOSTED));

      vector<double> rtn;
      rtn.push_back(angle1 < angle2 ? angle1 : angle2);
      rtn.push_back(anglebb);
      return rtn;
    }

    //@}


  private:

    /// @name Histograms
    //@{

    Histo1DPtr _h_Z_multiplicity, _h_W_multiplicity;
    Histo1DPtr _h_jet_bb_Delta_eta, _h_jet_bb_Delta_phi, _h_jet_bb_Delta_pT, _h_jet_bb_Delta_R;
    Histo1DPtr _h_jet_b_jet_eta, _h_jet_b_jet_multiplicity, _h_jet_b_jet_phi, _h_jet_b_jet_pT;
    Histo1DPtr _h_jet_H_eta_using_bb, _h_jet_H_mass_using_bb, _h_jet_H_phi_using_bb, _h_jet_H_pT_using_bb;
    Histo1DPtr _h_jet_eta, _h_jet_multiplicity, _h_jet_phi, _h_jet_pT;
    Histo1DPtr _h_jet_VBbb_Delta_eta, _h_jet_VBbb_Delta_phi, _h_jet_VBbb_Delta_pT, _h_jet_VBbb_Delta_R;
    Histo1DPtr _h_VB_eta, _h_VB_mass, _h_VB_phi, _h_VB_pT;
    Histo1DPtr _h_jet_bVB_angle_Hframe, _h_jet_bb_angle_Hframe, _h_jet_bVB_cosangle_Hframe, _h_jet_bb_cosangle_Hframe;
    //Histo1DPtr _h_jet_cuts_bb_deltaR_v_HpT;

    //@}

  };


  // This global object acts as a hook for the plugin system
  RIVET_DECLARE_PLUGIN(MC_VH2BB);

}
