// -*- C++ -*-
#include "Rivet/Analysis.hh"
#include "Rivet/RivetAIDA.hh"
#include "Rivet/Tools/Logging.hh"
#include "Rivet/Math/MathUtils.hh"
#include "Rivet/Tools/ParticleIdUtils.hh"
#include "HepMC/GenParticle.h"
#include "HepMC/GenVertex.h"
#include "HepMC/GenRanges.h"
#include "Rivet/Tools/BinnedHistogram.hh"

//#include "TransverseThrust.h" // The guy from the eventshapelib

#include <set>
#include <assert.h>
#include <sstream>
using namespace HepMC;
namespace Rivet {


class MC_EVENTSHAPES : public Analysis {
public:

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

    /// Constructor
    MC_EVENTSHAPES()
        : Analysis("MC_EVENTSHAPES")
    {


    }

    //@}


public:


    // based on PDGID, return true if particle ID is partonic
    bool isParton(int pdgid) {
      bool dec = false;
      int partonIDs[9] = {1, 2, 3, 4, 5, 6, 7, 8, 21};
      for (int i=0; i<9; i++) {
        if (fabs(pdgid) == partonIDs[i]) dec=true;
      }
      return dec;
    }

    // Return the number of MPI vertices
    int nMPIVertices(const Event& event, int gen) {
      // Get the HepMC record 
      const GenEvent & hepmc = event.genEvent();
      // Iterate over GenParticles
      int nvertices = 0;
      int MPIvertices = 0;
      // The int gen toggles the MC generator
      // 0 = Sherpa
      // 1 = Pythia8

      if (gen == 0) {
        for (HepMC::GenEvent::vertex_const_iterator vtx = hepmc.vertices_begin() ;
               vtx != hepmc.vertices_end() ; ++vtx) {
          // Loop over and count incoming GenParticles
          int npin = 0;
          int npartin = 0;
          for (HepMC::GenVertex::particles_in_const_iterator pin = (*vtx)->particles_in_const_begin();
               pin != (*vtx)->particles_in_const_end() ; ++pin) {
            npin++;
            int pdgid = (*pin)->pdg_id();
            if ( isParton(pdgid) ) npartin++;
          }

          // Loop over and count outgoing GenParticles
          int npout = 0;
          int npartout = 0;
          for (HepMC::GenVertex::particles_out_const_iterator pout = (*vtx)->particles_out_const_begin();
               pout != (*vtx)->particles_out_const_end() ; ++pout) {
            npout++;
            int pdgid = (*pout)->pdg_id();
            if ( isParton(pdgid) )  npartout++;
          }
          nvertices++;
          
          // Make sure we found an MPI vertex 
          if (npin ==  npartin && npout == npartout && npin == 2 && npartin == npartout) MPIvertices++;
        }
      }
      else if (gen == 1) {
        for (HepMC::GenEvent::vertex_const_iterator vtx = hepmc.vertices_begin() ;
               vtx != hepmc.vertices_end() ; ++vtx) {
          // Loop over and count incoming GenParticles
          int npin = 0;
          int npartin = 0;
          for (HepMC::GenVertex::particles_in_const_iterator pin = (*vtx)->particles_in_const_begin();
               pin != (*vtx)->particles_in_const_end() ; ++pin) {
            npin++;
            int status = (*pin)->status();
            if ( status==31 ) npartin++;
            // check for status 21 or 31
          }
          
          // Make sure we found an MPI vertex 
          if (npin ==  npartin && npin == 2 ) MPIvertices++;
        }
      }
      return MPIvertices;
    }
    
    // Return the number of partons
    int nPartons(const Event& event) {
      // Get the HepMC record 
      const GenEvent & hepmc = event.genEvent();
      // Iterate over GenParticles
      int nvertices = 0;
      int HARDvertices = 0;
      int njets = 0;
      for (HepMC::GenEvent::vertex_const_iterator vtx = hepmc.vertices_begin() ;
             vtx != hepmc.vertices_end() ; ++vtx) {
        // Loop over and count incoming GenParticles
        int npin = 0;
        int npartin = 0;
        for (HepMC::GenVertex::particles_in_const_iterator pin = (*vtx)->particles_in_const_begin();
             pin != (*vtx)->particles_in_const_end() ; ++pin) {
          npin++;
          int pdgid = (*pin)->pdg_id();
          if ( isParton(pdgid) ) npartin++;
        }

        // Loop over and count outgoing GenParticles and
        int npout = 0;
        int npartout = 0;
        int nElectrons = 0;
        int nPositrons = 0;
        for (HepMC::GenVertex::particles_out_const_iterator pout = (*vtx)->particles_out_const_begin();
             pout != (*vtx)->particles_out_const_end() ; ++pout) {
          npout++;
          int pdgid = (*pout)->pdg_id();
          if ( isParton(pdgid) )  npartout++;
          if (pdgid == 11 || pdgid ==13) nElectrons ++;
          if (pdgid == -11 || pdgid ==-13 ) nPositrons ++;
        }
        nvertices++;
        
        // Make sure we found the Z-production vertex
        if (npin ==  npartin && nElectrons == 1 && nPositrons == 1) {
          //(*vtx)->print();
          HARDvertices++;
          njets = npartout;
        }
      }
      // Sanity check
      //if (HARDvertices !=1) cout << "Oh nooooooooooo" << endl;


      return njets;
    }
    /// @name Analysis methods
    //@{

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

    }       


    /// Perform the per-event analysis
    void analyze(const Event& event) {
     
        // Get generator weight
        const double weight = event.weight();


        // 0=Sherpa, 1=Pythia8
        const double NMPI = nMPIVertices(event, 0);

        const double nJets = nPartons(event);
    }


    /// Normalise histograms etc., after the run
    void finalize() {



    }


private:


};

// This global object acts as a hook for the plugin system
DECLARE_RIVET_PLUGIN(MC_EVENTSHAPES);


}