// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/FastJets.hh"
#include "Rivet/Projections/SmearedJets.hh"

namespace Rivet {


  /// @brief CDF properties of 6-jet events with large 6-jet mass
  class CDF_1997_S3541940 : public Analysis {
  public:

    RIVET_DEFAULT_ANALYSIS_CTOR(CDF_1997_S3541940);


    void init() {

      // Find true jets
      const FinalState fs(Cuts::abseta < 4.2);
      FastJets fj(fs, FastJets::CDFJETCLU, 0.7);

      // Smear jet energy and mass with the 10% uncertainty quoted in the paper
      SmearedJets sj_E(fj, [](const Jet& jet){ return P4_SMEAR_MASS_GAUSS(P4_SMEAR_E_GAUSS(jet, 0.1*jet.E()), 0.1*jet.mass()); });
      declare(sj_E, "Jets");

      book(_h_m6J ,1, 1, 1);
      book(_h_X3ppp ,2, 1, 1);
      book(_h_X4ppp ,3, 1, 1);
      book(_h_costheta3ppp ,4, 1, 1);
      book(_h_psi3ppp ,5, 1, 1);
      book(_h_f3ppp ,6, 1, 1);
      book(_h_f4ppp ,6, 1, 2);
      book(_h_f5ppp ,6, 1, 3);
      book(_h_XApp ,7, 1, 1);
      book(_h_XCp ,8, 1, 1);
      book(_h_XE ,9, 1, 1);
      book(_h_psiAppBpp ,10, 1, 1);
      book(_h_psiCpDp ,11, 1, 1);
      book(_h_psiEF ,12, 1, 1);
      book(_h_fApp ,13, 1, 1);
      book(_h_fBpp ,14, 1, 1);
      book(_h_fCp ,15, 1, 1);
      book(_h_fDp ,16, 1, 1);
      book(_h_fE ,17, 1, 1);
      book(_h_fF ,18, 1, 1);
    }


    void analyze(const Event& event) {

      const Jets alljets = apply<JetAlg>(event, "Jets").jets(Cuts::Et > 20*GeV && Cuts::abseta < 3, cmpMomByEt);

      Jets jets;
      double sumEt = 0.0;
      FourMomentum jetsystem(0.0, 0.0, 0.0, 0.0);
      for (const Jet& jet : alljets) {
        double Et = jet.Et();
        bool separated = true;
        for (const Jet& ref : jets) {
          if (deltaR(jet, ref) < 0.9) {
            separated = false;
            break;
          }
        }
        if (!separated) continue;
        jets.push_back(jet);
        sumEt += Et;
        jetsystem += jet.momentum();
        if (jets.size() >= 6) break;
      }

      if (jets.size() < 6) vetoEvent;
      if (sumEt < 320.0*GeV) vetoEvent;

      double m6J = _safeMass(jetsystem);
      if (m6J < 520.0*GeV) vetoEvent;

      if (getLog().isActive(Log::DEBUG)) {
        stringstream ss;
        ss << "Jets:\n";
        for (const Jet& j : jets) ss << j << "\n";
        MSG_DEBUG(ss.str());
      }

      const LorentzTransform cms_boost = LorentzTransform::mkFrameTransformFromBeta(jetsystem.betaVec());
      vector<FourMomentum> jets6;
      for (Jet jet : jets) {
        jets6.push_back(cms_boost.transform(jet.momentum()));
      }
      std::sort(jets6.begin(), jets6.end(), FourMomentum::byEDescending());

      FourMomentum pE, pF;
      vector<FourMomentum> jets5(_reduce(jets6, pE, pF));
      std::sort(jets5.begin(), jets5.end(), FourMomentum::byEDescending());

      FourMomentum pCp, pDp;
      vector<FourMomentum> jets4(_reduce(jets5, pCp, pDp));
      std::sort(jets4.begin(), jets4.end(), FourMomentum::byEDescending());

      FourMomentum pApp, pBpp;
      vector<FourMomentum> jets3(_reduce(jets4, pApp, pBpp));
      std::sort(jets3.begin(), jets3.end(), FourMomentum::byEDescending());
      FourMomentum p3ppp(jets3[0]);
      FourMomentum p4ppp(jets3[1]);
      FourMomentum p5ppp(jets3[2]);

      double X3ppp = 2.0*p3ppp.E()/m6J;
      if (X3ppp > 0.9) vetoEvent;

      FourMomentum pAV = cms_boost.transform(_avg_beam_in_lab(m6J, jetsystem.rapidity()));
      double costheta3ppp = pAV.p3().unit().dot(p3ppp.p3().unit());
      if (fabs(costheta3ppp) > 0.9) vetoEvent;

      // 3-jet-system variables
      _h_m6J->fill(m6J);
      _h_X3ppp->fill(X3ppp);
      _h_X4ppp->fill(2.0*p4ppp.E()/m6J);
      _h_costheta3ppp->fill(costheta3ppp);
      double psi3ppp = _psi(p3ppp, pAV, p4ppp, p5ppp);
      _h_psi3ppp->fill(psi3ppp);
      _h_f3ppp->fill(_safeMass(p3ppp)/m6J);
      _h_f4ppp->fill(_safeMass(p4ppp)/m6J);
      _h_f5ppp->fill(_safeMass(p5ppp)/m6J);

      // 4 -> 3 jet variables
      _h_fApp->fill(_safeMass(pApp)/m6J);
      _h_fBpp->fill(_safeMass(pApp)/m6J);
      _h_XApp->fill(pApp.E()/(pApp.E()+pBpp.E()));
      double psiAppBpp = _psi(pApp, pBpp, pApp+pBpp, pAV);
      _h_psiAppBpp->fill(psiAppBpp);

      // 5 -> 4 jet variables
      _h_fCp->fill(_safeMass(pCp)/m6J);
      _h_fDp->fill(_safeMass(pDp)/m6J);
      _h_XCp->fill(pCp.E()/(pCp.E()+pDp.E()));
      double psiCpDp = _psi(pCp, pDp, pCp+pDp, pAV);
      _h_psiCpDp->fill(psiCpDp);

      // 6 -> 5 jet variables
      _h_fE->fill(_safeMass(pE)/m6J);
      _h_fF->fill(_safeMass(pF)/m6J);
      _h_XE->fill(pE.E()/(pE.E()+pF.E()));
      double psiEF = _psi(pE, pF, pE+pF, pAV);
      _h_psiEF->fill(psiEF);
    }


    void finalize() {
      normalize(_h_m6J);
      normalize(_h_X3ppp);
      normalize(_h_X4ppp);
      normalize(_h_costheta3ppp);
      normalize(_h_psi3ppp);
      normalize(_h_f3ppp);
      normalize(_h_f4ppp);
      normalize(_h_f5ppp);
      normalize(_h_XApp);
      normalize(_h_XCp);
      normalize(_h_XE);
      normalize(_h_psiAppBpp);
      normalize(_h_psiCpDp);
      normalize(_h_psiEF);
      normalize(_h_fApp);
      normalize(_h_fBpp);
      normalize(_h_fCp);
      normalize(_h_fDp);
      normalize(_h_fE);
      normalize(_h_fF);
    }


  private:

    vector<FourMomentum> _reduce(const vector<FourMomentum>& jets,
                                 FourMomentum& combined1,
                                 FourMomentum& combined2) {
      double minMass2 = 1e9;
      size_t idx1(jets.size()), idx2(jets.size());
      for (size_t i = 0; i < jets.size(); ++i) {
        for (size_t j = i+1; j < jets.size(); ++j) {
          double mass2 = FourMomentum(jets[i] + jets[j]).mass2();
          if (mass2 < minMass2) {
            idx1 = i;
            idx2 = j;
          }
        }
      }
      vector<FourMomentum> newjets;
      for (size_t i = 0; i < jets.size(); ++i) {
        if (i != idx1 && i != idx2) newjets.push_back(jets[i]);
      }
      newjets.push_back(jets[idx1] + jets[idx2]);
      combined1 = jets[idx1];
      combined2 = jets[idx2];
      return newjets;
    }


    FourMomentum _avg_beam_in_lab(const double& m, const double& y) {
      const double mt = m/2.0;
      FourMomentum beam1(mt, 0, 0,  mt);
      FourMomentum beam2(mt, 0, 0, -mt);
      if (fabs(y) > 1e-3) {
        FourMomentum boostvec(cosh(y), 0.0, 0.0, sinh(y));
        const LorentzTransform cms_boost = LorentzTransform::mkFrameTransformFromBeta(boostvec.betaVec()).inverse();
        beam1 = cms_boost.transform(beam1);
        beam2 = cms_boost.transform(beam2);
      }
      return (beam1.E() > beam2.E()) ? beam1 - beam2 : beam2 - beam1;
    }


    double _psi(const FourMomentum& p1, const FourMomentum& p2,
                const FourMomentum& p3, const FourMomentum& p4) {
      Vector3 p1xp2 = p1.p3().cross(p2.p3());
      Vector3 p3xp4 = p3.p3().cross(p4.p3());
      return mapAngle0ToPi(acos(p1xp2.unit().dot(p3xp4.unit())));
    }


    double _safeMass(const FourMomentum& p) {
      double mass2 = p.mass2();
      if (mass2 > 0.0) return sqrt(mass2);
      if (mass2 < -1e-5) MSG_WARNING("m2 = " << m2 << ". Assuming m2=0.");
      return 0.0;
    }


  private:

    Histo1DPtr _h_m6J;
    Histo1DPtr _h_X3ppp, _h_X4ppp;
    Histo1DPtr _h_costheta3ppp;
    Histo1DPtr _h_psi3ppp;
    Histo1DPtr _h_f3ppp, _h_f4ppp, _h_f5ppp;
    Histo1DPtr _h_XApp, _h_XCp, _h_XE;
    Histo1DPtr _h_psiAppBpp, _h_psiCpDp, _h_psiEF;
    Histo1DPtr _h_fApp, _h_fBpp, _h_fCp, _h_fDp, _h_fE, _h_fF;

  };



  RIVET_DECLARE_ALIASED_PLUGIN(CDF_1997_S3541940, CDF_1997_I442265);

}
