// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/Projections/FinalState.hh"
#include "Rivet/Projections/LeptonFinder.hh"
#include "Rivet/Projections/MissingMomentum.hh"
#include "Rivet/Projections/Beam.hh"

namespace Rivet {


  /// @brief MC validation analysis for W polarisation
  class MC_WPOL : public Analysis {
  public:

    /// Constructor
    RIVET_DEFAULT_ANALYSIS_CTOR(MC_WPOL);


    /// @name Analysis methods
    /// @{

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

      declare(MissingMomentum(), "MET");

      LeptonFinder lf(0.1, Cuts::abspid == PID::ELECTRON);
      declare(lf, "Leptons");

      Beam beams;
      declare(beams, "Beams");

      vector<string> tags{"_wplus", "_wminus"};
      _h_dists.resize(tags.size());
      _h_histos.resize(tags.size());
      for (size_t i=0; i<tags.size(); ++i) {
        _h_dists[i].resize(11,Profile1DPtr());
        double sqrts = sqrtS()>0. ? sqrtS() : 14000.;
        book(_h_dists[i][0] ,"A0"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][1] ,"A1"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][2] ,"A2"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][3] ,"A3"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][4] ,"A4"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][5] ,"A5"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][6] ,"A6"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][7] ,"A7"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][8] ,"fL"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][9] ,"fR"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        book(_h_dists[i][10] ,"f0"+tags[i],logspace(100, 1.0, 0.5*sqrts));
        _h_histos[i].resize(4,Histo1DPtr());
        book(_h_histos[i][0] ,"thetastar"+tags[i],100,-1.0,1.0);
        book(_h_histos[i][1] ,"phistar"+tags[i],90,0.0,360.0);
        book(_h_histos[i][2] ,"thetastar_ptw20"+tags[i],100,-1.0,1.0);
        book(_h_histos[i][3] ,"phistar_ptw20"+tags[i],90,0.0,360.0);
      }
    }


    /// Perform the per-event analysis
    void analyze(const Event& event) {

      // MET cut
      const P4& pmiss = apply<MissingMom>(event, "MET").missingMom();
      if (pmiss.pT() < 25*GeV) vetoEvent;

      // Identify the closest-matching l+MET to m == mW
      const Particles& ls = apply<LeptonFinder>(event, "Leptons").particles();
      const int ifound = closestMassIndex(ls, pmiss, 80.4*GeV, 60*GeV, 100*GeV);
      if (ifound < 0) vetoEvent;

      const ParticlePair& beams = apply<Beam>(event, "Beams").beams();
      FourMomentum pb1(beams.second.momentum()), pb2(beams.first.momentum());
      const Particle& lepton = ls[ifound];
      FourMomentum pl(lepton.momentum());
      const size_t idx = lepton.charge3() > 0 ? 0 : 1;
      const FourMomentum plnu = lepton.mom() + pmiss;

      const LorentzTransform cms = LorentzTransform::mkFrameTransformFromBeta(plnu.betaVec());
      Matrix3 zrot(plnu.p3(), Vector3(0.0, 0.0, 1.0));
      pl = cms.transform(pl);
      pb1 = cms.transform(pb1);
      pb2 = cms.transform(pb2);
      Vector3 pl3 = pl.p3();
      Vector3 pb13 = pb1.p3();
      Vector3 pb23 = pb2.p3();
      pl3 = zrot*pl3;
      pb13 = zrot*pb13;
      pb23 = zrot*pb23;
      Vector3 xref(cos(pb13.theta())>cos(pb23.theta()) ? pb13 : pb23);
      Matrix3 xrot(Vector3(xref.x(), xref.y(), 0.0), Vector3(1.0, 0.0, 0.0));
      pl3 = xrot*pl3;

      double ptw(plnu.pT()/GeV);
      double thetas(pl3.theta()), phis(pl3.phi());
      double costhetas(cos(thetas)), sinthetas(sin(thetas));
      double cosphis(cos(phis)), sinphis(sin(phis));
      if (phis < 0.0) phis += 2.0*M_PI;

      _h_histos[idx][0]->fill(costhetas);
      _h_histos[idx][1]->fill(phis*180.0/M_PI);
      if (ptw > 20.0) {
        _h_histos[idx][2]->fill(costhetas);
        _h_histos[idx][3]->fill(phis*180.0/M_PI);
      }
      _h_dists[idx][0]->fill(ptw,10.0/3.0*(1.0-3.0*sqr(costhetas))+2.0/3.0);
      _h_dists[idx][1]->fill(ptw,10.0*sinthetas*costhetas*cosphis);
      _h_dists[idx][2]->fill(ptw,10.0*sqr(sinthetas)*(sqr(cosphis)-sqr(sinphis)));
      _h_dists[idx][3]->fill(ptw,4.0*sinthetas*cosphis);
      _h_dists[idx][4]->fill(ptw,4.0*costhetas);
      _h_dists[idx][5]->fill(ptw,4.0*sinthetas*sinphis);
      _h_dists[idx][6]->fill(ptw,10.0*costhetas*sinthetas*sinphis);
      _h_dists[idx][7]->fill(ptw,10.0*sqr(sinthetas)*cosphis*sinphis);
      _h_dists[idx][8]->fill(ptw,0.5*sqr(1.0-costhetas)-(1.0-2.0*sqr(costhetas)));
      _h_dists[idx][9]->fill(ptw,0.5*sqr(1.0+costhetas)-(1.0-2.0*sqr(costhetas)));
      _h_dists[idx][10]->fill(ptw,5.0*sqr(sinthetas)-3.0);
    }


    /// Normalise histograms etc., after the run
    void finalize() {
      for (size_t i = 0; i < _h_histos.size(); ++i) {
        for (Histo1DPtr histo : _h_histos[i]) {
          scale(histo, crossSection()/picobarn/sumOfWeights());
        }
      }
    }

    /// @}


  private:

    /// @name Histograms
    /// @{
    vector<vector<Profile1DPtr> > _h_dists;
    vector<vector<Histo1DPtr> > _h_histos;
    /// @}


  };


  RIVET_DECLARE_PLUGIN(MC_WPOL);

}
