////////////////////////////////////////////////////////////////////////////////// 
/// Analysis skeleton
/// also for distributed analysis examples
/// Author: Ketevi A. Assamagan
/// BNL, July 22, 2004
///
/// DESCRIPTION:
///
/// This class is an analysis skeleton - The user can implement his analysis here
/// This class is also used for the demonstration of the distributed analysis
/// Some electron histograms are used for the distributed case. The user may
/// remove the histograms and the electron stuff if not needed.
/// Note t:he single algorithm structure as an analysis code does not scale
/// For detailed analysis examples, look in CVS: PhysicsAnalysis/AnalysisCommon/AnalysisExamples/
/// Ketevi A. Assamagan on June 22, 2004
///
///////////////////////////////////////////////////////////////////////////////////////////////////////


//////////////////////////////////////////////////////////////////////////////////

#include "GaudiKernel/MsgStream.h"
#include "GaudiKernel/AlgFactory.h"
#include "GaudiKernel/IToolSvc.h"

#include "StoreGate/StoreGateSvc.h"

#include "egammaEvent/ElectronContainer.h"

#include "McParticleEvent/TruthParticleContainer.h"

#include "GaudiKernel/ITHistSvc.h"
#include "TTree.h"

#include "UserAnalysis/AnalysisSkeleton.h"

#include <algorithm>
#include <math.h>
#include <functional>
#include <iostream>

static const double mZ = 91.19*GeV;
static const int  MAX_PARTICLES = 20;

using namespace Analysis;

//////////////////////////////////////////////////////////////////////////////////////
/// Constructor

AnalysisSkeleton::AnalysisSkeleton(const std::string& name,
  ISvcLocator* pSvcLocator) : CBNT_AthenaAwareBase(name, pSvcLocator),
  m_analysisTools( "AnalysisTools", this ) {

  /** switches to control the analysis through job options */

  /////////////// Remove this if not needed //////////////////////////////////////////
  /// This is here only for distributed analysis example

  declareProperty( "AnalysisTools", m_analysisTools );
  declareProperty("ElectronContainer", m_electronContainerName = "ElectronAODCollection");
  declareProperty("MCParticleContainer", m_truthParticleContainerName = "SpclMC");

  /** the cuts - default values - to be modified in job options */

  declareProperty("DeltaRMatchCut", m_deltaRMatchCut = 0.2);
  declareProperty("MaxDeltaR", m_maxDeltaR = 0.9999);

  /** for electrons */
  declareProperty("ElectronEtCut", m_etElecCut = 10.0*GeV);
  declareProperty("ElectronEtaCut", m_etaElecCut = 2.5);
  declareProperty("ElectronCone", m_elecCone = 0.9);

  ///////////////// Remove the above if not needed ///////////////////////////////////

  }



/////////////////////////////////////////////////////////////////////////////////////
/// Destructor - check up memory allocation
/// delete any memory allocation on the heap

AnalysisSkeleton::~AnalysisSkeleton() {}

////////////////////////////////////////////////////////////////////////////////////
/// Initialize
/// initialize StoreGate
/// get a handle on the analysis tools
/// book histograms

StatusCode AnalysisSkeleton::CBNT_initialize() {

  MsgStream mLog( messageService(), name() );

  mLog << MSG::DEBUG << "Initializing AnalysisSkeleton" << endreq;

  /** get a handle of StoreGate for access to the Event Store */
  StatusCode sc = service("StoreGateSvc", m_storeGate);
  if (sc.isFailure()) {
     mLog << MSG::ERROR
          << "Unable to retrieve pointer to StoreGateSvc"
          << endreq;
     return sc;
  }
  
  /// get a handle on the analysis tools
  sc = m_analysisTools.retrieve();
  if ( sc.isFailure() ) {
      mLog << MSG::ERROR << "Can't get handle on analysis tools" << endreq;
      return sc;
  }

  /** get a handle on the NTuple and histogramming service */
  sc = service("THistSvc", m_thistSvc);
  if (sc.isFailure()) {
     mLog << MSG::ERROR
          << "Unable to retrieve pointer to THistSvc"
          << endreq;
     return sc;
  }

  /////////////////////// To be removed if not needed /////////////////////////


  /** Athena-Aware NTuple (AAN) - this is the recommended option for users */

  /** now add branches and leaves to the AAN tree */

  addBranch("NElectrons",   m_aan_size, "NElectrons/i");
  addBranch("ElectronEta",   m_aan_eta);
  addBranch("ElectronPt",    m_aan_pt);
  addBranch("ElecPtRatio",   m_aan_elecetres);

  /// ROOT histograms ---------------------------------------

  m_h_elecpt     = new TH1F("elec_pt","pt el",50,0,250.*GeV);
  sc = m_thistSvc->regHist("/AANT/Electron/elec_pt",m_h_elecpt);

  m_h_eleceta    = new TH1F("elec_eta","eta el",70,-3.5,3.5);
  sc = m_thistSvc->regHist("/AANT/Electron/elec_eta",m_h_eleceta);

  m_h_tranverseMofW = new TH1F("trans_mass",       "mT(W)",    100, 0.0, 200.0*GeV);
  sc = m_thistSvc->regHist("/AANT/Truth/trans_mass",m_h_tranverseMofW);

  m_h_nuPt = new TH1F("nu_pt","pt nu",50,0,250.*GeV);
  sc = m_thistSvc->regHist("/AANT/Truth/nu_pt",m_h_nuPt);

  m_h_qqMass = new TH1F("qq_mass","mass_qq",50,0.,100.*GeV);
  sc = m_thistSvc->regHist("/AANT/Truth/qq_mass",m_h_qqMass);

  if (sc.isFailure()) { 
     mLog << MSG::ERROR << "ROOT Hist registration failed" << endreq; 
     return sc; 
  }
  /// end ROOT Histograms ------------------------------------------

  ////////////////////// The above to be removed if not needed ////////////////////
  
  return StatusCode::SUCCESS;
}		 

///////////////////////////////////////////////////////////////////////////////////
/// Finalize - delete any memory allocation from the heap

StatusCode AnalysisSkeleton::CBNT_finalize() {
  MsgStream mLog( messageService(), name() );
  
  return StatusCode::SUCCESS;

}

///////////////////////////////////////////////////////////////////////////////////
/// Clear - clear CBNT members
StatusCode AnalysisSkeleton::CBNT_clear() {
  /// For Athena-Aware NTuple

  m_aan_size = 0;
  m_aan_eta->clear();
  m_aan_pt->clear();
  m_aan_elecetres->clear();

  return StatusCode::SUCCESS;
}

//////////////////////////////////////////////////////////////////////////////////
/// Execute - on event by event

StatusCode AnalysisSkeleton::CBNT_execute() {
  MsgStream mLog( messageService(), name() );
  
  mLog << MSG::DEBUG << "in execute()" << endreq;

  /// remove the lines if not needed for your analysis
  /// they are here to exercice the Distributed Analysis

  ////////////////// To Be removed if NOT needed ////////////////////////////////

  StatusCode sc = electronSkeleton();
  if (sc.isFailure()) {
    mLog << MSG::ERROR << "The method electronSkeleton() failed" << endreq;
    return StatusCode::SUCCESS;
  }

  return StatusCode::SUCCESS;
}

//////////////////////////////////////////////////////////////////////////////////
/// Electron method - called by execute() on event by event
/// to be removed if not needed

StatusCode AnalysisSkeleton::electronSkeleton() {
  MsgStream mLog( messageService(), name() );
  
  mLog << MSG::DEBUG << "in electronSkeleton()" << endreq;

  /** get the MC truth particle AOD container from StoreGate */
  const TruthParticleContainer*  mcpartTES = 0;
  StatusCode sc=m_storeGate->retrieve( mcpartTES, m_truthParticleContainerName);
  if( sc.isFailure()  ||  !mcpartTES ) {
     mLog << MSG::WARNING
          << "No AOD MC truth particle container found in TDS"
          << endreq; 
     return StatusCode::SUCCESS;
  }
  mLog <<MSG::DEBUG << "MC Truth Container Successfully Retrieved" << endreq;
  // just check remove after checking
  //mcpartTES->dump();

  const HepMC::GenEvent* genEvt = mcpartTES->genEvent();
  HepMC::GenEvent::vertex_const_iterator vitr;
  vitr=genEvt->vertices_begin();
// check primary vertex
  std::vector<HepMC::GenParticle> vq_pvrt;
  int num_vbfq=0;
  HepLorentzVector higgs_hlv;  
  for(HepMC::GenVertex::particles_out_const_iterator opitr=(*vitr)->particles_out_const_begin();
          opitr!=(*vitr)->particles_out_const_end();opitr++){
     int pid=(*opitr)->pdg_id();
     if(abs(pid)<7 && !(pid == 0)){
       vq_pvrt.push_back(**opitr);
       num_vbfq++;
     }
     else{
       if(pid=25){
         higgs_hlv = (*opitr)->momentum();
       }
     }
   }
   mLog <<MSG::DEBUG << "primary vertex searched" << endreq;   
   // serch W decay vertex
   int num_w_ele=0;
   int num_w_mu=0;
   int num_w_nu=0;
   int num_w_qq=0;
   int num_w=0;
   HepMC::GenParticle gen_ele,gen_mu,gen_nu,gen_q1,gen_q2;
   
   for(HepMC::GenEvent::vertex_const_iterator vitr=genEvt->vertices_begin();
     vitr!=genEvt->vertices_end(); vitr++){
      if((*vitr)->particles_in_size() == 1){ 
       HepMC::GenVertex::particles_in_const_iterator ipitr=(*vitr)->particles_in_const_begin();
       int pid=(*ipitr)->pdg_id();
       if(pid == 24 || pid == -24){
         for(HepMC::GenVertex::particles_out_const_iterator opitr=
          (*vitr)->particles_out_const_begin();opitr!=(*vitr)->particles_out_const_end();opitr++){
            int apid=abs((*opitr)->pdg_id());
            mLog <<MSG::DEBUG << "pdg_id of out= " << (*opitr)->pdg_id() << endreq;
            mLog <<MSG::DEBUG << "abs pdg_id of out= " << apid << endreq;
            if(apid>0 && apid< 5) apid =1;
            switch(apid){
              case 11:
               gen_ele=**opitr;
               num_w_ele++;
               break;
              case 13:
               gen_mu=**opitr;
               num_w_mu++;
               break;
              case 12:
               gen_nu=**opitr;
               num_w_nu++;
               break;
              case 14:
               gen_nu=**opitr;
               num_w_nu++;
               break;
              case 1:
               if(num_w_qq == 0){
                 gen_q1=**opitr;
                 num_w_qq++;
               }
               else{
                 gen_q2=**opitr;
                 num_w_qq++;
               }
               break;
            } 
         }// w decay product check
         num_w++;
       } //w found
      }// in comming particle should be one
      if(num_w == 2) break;
   } //vertex loop
   //genEvt->print();
//  for (HepMC::GenEvent::particle_const_iterator pitr=genEvt->particles_begin();
//	 pitr!=genEvt->particles_end();
//	 ++pitr) {
//     int id = (*pitr)->pdg_id();
//  }   //mLog <<MSG::DEBUG << "id=" << id << endreq; 
     //check q truth
//       const HepMC::GenVertex* genVtx = (*pitr)->production_vertex();
//       HepPoint3D three=genVtx->point3d();
//       mLog <<MSG::DEBUG << "Vertex= "<< three << endreq;
/*       if(!genVtx){
       for(HepMC::GenVertex::particles_out_const_iterator opiter=genVtx->particles_out_const_begin();
          opiter!=genVtx->particles_out_const_end();opiter++){
         int pid=(*opiter)->pdg_id();
         if(pid == 24 || pid == -24) W_q_flag = true;
       }
       } */
  mLog <<MSG::DEBUG << "num_vbfq= " <<  num_vbfq << " num_w= " <<  num_w << endreq;
  mLog <<MSG::DEBUG << "num_w_ele= " <<  num_w_ele << " num_w_mu= " <<  num_w_mu << endreq;
  mLog <<MSG::DEBUG << "num_w_nu= " <<  num_w_ele << " num_w_qq= " <<  num_w_qq << endreq;

  if( num_w_ele == 1 || num_w_mu ==1){
    double lep_Pt = 0.;
    double nu_Pt= 0.;
    HepLorentzVector lep,nu;
    if( num_w_ele == 1){
       lep = gen_ele.momentum();
       nu = gen_nu.momentum();
       lep_Pt=std::sqrt(lep.px()*lep.px() + lep.py()*lep.py());
       nu_Pt=std::sqrt(nu.px()*nu.px() + nu.py()*nu.py());
    }
    if( num_w_mu == 1){
       lep = gen_mu.momentum();
       nu = gen_nu.momentum();
       lep_Pt=std::sqrt(lep.px()*lep.px() + lep.py()*lep.py());
       nu_Pt=std::sqrt(nu.px()*nu.px() + nu.py()*nu.py());
    }
    double m2 = 2.0*(lep_Pt*nu_Pt - (lep.px()*nu.px() + lep.py()*nu.py()));
    if (m2<0) {
      mLog << MSG::DEBUG << "Warning: m2 is negative." << endreq;
    }
    else {
      double mT = std::sqrt(m2);
      m_h_tranverseMofW->Fill(mT,1.);
      m_h_nuPt->Fill(nu_Pt,1);     
    }
    HepLorentzVector q1,q2,p_w;
    if(num_w_qq == 2 ){
       q1 = gen_q1.momentum();
       q2 = gen_q2.momentum();
       p_w = q1 + q2;
       //double mqqb = p_w.invariantMass();
       double mqqb=0;
       double m2 = p_w.e()*p_w.e()-(p_w.px()*p_w.px()+p_w.py()*p_w.py()+p_w.pz()*p_w.pz());
       if(m2<0){
         mLog << MSG::DEBUG << "Warning: m2 is negative." << endreq;
       }else{
         mqqb = std::sqrt(m2);
       }
       m_h_qqMass->Fill( mqqb,1.);
    }
  }
//
  const ElectronContainer* elecTES = 0;
  sc=m_storeGate->retrieve( elecTES, m_electronContainerName);
  if( sc.isFailure()  ||  !elecTES ) {
     mLog << MSG::WARNING
          << "No AOD electron container found in TDS"
          << endreq; 
     return StatusCode::FAILURE;
  }  
  mLog << MSG::DEBUG << "ElectronContainer successfully retrieved" << endreq;

  /** iterators over the electron container */ 
  ElectronContainer::const_iterator elecItr  = elecTES->begin();
  ElectronContainer::const_iterator elecItrE = elecTES->end();

  for (; elecItr != elecItrE; ++elecItr) {

    /** check for the author of the electron */
    bool author = (*elecItr)->author() == egammaParameters::AuthorElectron;
    if ( !author || (*elecItr)->pt() < m_etElecCut ) continue;

    m_aan_size++;

    /** fill histograms */
    m_h_elecpt->Fill( (*elecItr)->pt(), 1.);
    m_h_eleceta->Fill( (*elecItr)->eta(), 1.);

    /** fill Athena-Aware NTuple */
    m_aan_eta->push_back((*elecItr)->eta());
    m_aan_pt->push_back((*elecItr)->pt());

    /** find a match to this electron in the MC truth container
        the index and deltaR are returned */
    int index = -1;
    double deltaRMatch;
    if( (*elecItr)->trackParticle() && (*elecItr)->pt()> m_etElecCut ) {
       const TruthParticleContainer * truthContainer = mcpartTES;
       bool findAMatch = m_analysisTools->matchR((*elecItr), truthContainer, 
		       index, deltaRMatch, (*elecItr)->pdgId());
       if (findAMatch) {
          deltaRMatch = (deltaRMatch > m_maxDeltaR) ? m_maxDeltaR : deltaRMatch;

          /** check for good match */
          if ( deltaRMatch < m_deltaRMatchCut) {
             const TruthParticle*  electronMCMatch = (*mcpartTES)[index]; 
             double res = (*elecItr)->pt() / electronMCMatch->pt();
             m_aan_elecetres->push_back(res);
          }
       }
    }    
  }

  /// remove the above lines if not needed for your analysis
  /// they are there to exercice Distributed Analysis 
  ////////////////// The above to be remove if not needed /////////////////////////////////

  mLog << MSG::DEBUG << "electronSkeleton() succeeded" << endreq;
  		
  return StatusCode::SUCCESS;
}