////////////////////////////////////////////////////////////////////////////////// 
/// 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 double mW = 80.403*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);
  
  addBranch("TrueValuesSize",m_aan_t_size, "TrueValuesSize/I");
  //addBranch("TrueTransMass", m_aan_t_tranverseMofW,"TrueTransMass/D");
  addBranch("TrueTransMass", m_aan_t_tranverseMofW);  
  addBranch("TrueNuPt", m_aan_t_nuPt);
  addBranch("TrueQQMass", m_aan_t_qqMass);
  addBranch("TrueLNuMass", m_aan_t_lnuMass);
  addBranch("TrueHMass", m_aan_t_hmass);
  /// 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);


  m_h_lnu = new TH1F("lnu_mass","mass_lnu",50,0.,150.*GeV);
  sc = m_thistSvc->regHist("/AANT/Truth/lnu_mass",m_h_lnu);

  m_h_hmass = new TH1F("higgs_mass","mass_h",50,100.,300.*GeV);
  sc = m_thistSvc->regHist("/AANT/Truth/higgs_mass",m_h_hmass);

  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();
// this is dangerous. always 0s are filled.
  m_aan_t_size=0;
  //m_aan_t_tranverseMofW=0.;
  m_aan_t_tranverseMofW->clear();
  m_aan_t_nuPt->clear();
  m_aan_t_qqMass->clear();
  m_aan_t_lnuMass->clear();
  m_aan_t_hmass->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 hig;  
//  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){
//         hig = (*opitr)->momentum();
//     double hmass_gen=sqrt(hig.e()*hig.e()-(hig.px()*hig.px()+hig.py()*hig.py()+hig.pz()*hig.pz()));
//         m_h_hmass_gen->Fill( hmass_gen,1.);
//       }
//     }
//   }
//   mLog <<MSG::DEBUG << "primary vertex searched" << endreq;   
   // serch W decay vertex
   int num_w_lep=0;
   int num_w_nu=0;
   int num_w_q=0;
   int num_w=0;
   int num_w_decay=0;
   std::vector<HepMC::GenParticle> gen_higgs;
   std::vector<HepMC::GenParticle> gen_w;
   std::vector<HepMC::GenParticle> gen_lep;
   std::vector<HepMC::GenParticle> gen_nu;
   std::vector<HepMC::GenParticle> gen_q;
   std::vector<int> bar_h_w;
   std::vector<int> bar_w_lep;
   std::vector<int> bar_w_nu;
   std::vector<int> bar_w_q;
   int lepSign=1;
//   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();
       int istat=(*ipitr)->status();
       //higgs vertex check
       if(pid == 25 && istat == 2){
          for(HepMC::GenVertex::particles_out_const_iterator opitr=(*vitr)->particles_out_const_begin();opitr!=(*vitr)->particles_out_const_end();opitr++){
             gen_w.push_back(**opitr);
             bar_h_w.push_back((*ipitr)->barcode());
             num_w++;
          }
       }
       //w deday vertex check
       if((pid == 24 || pid == -24) && istat == 2){ // exclude status == 3
         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_lep.push_back(**opitr);
               bar_w_lep.push_back((*ipitr)->barcode());
               num_w_lep++;
               if((*opitr)->pdg_id() < 0) lepSign=-1*lepSign;
               break;
              case 13:
               gen_lep.push_back(**opitr);
               bar_w_lep.push_back((*ipitr)->barcode());
               num_w_lep++;
               if((*opitr)->pdg_id() < 0) lepSign=-1*lepSign;
               break;
              case 12:
               gen_nu.push_back(**opitr);
               bar_w_nu.push_back((*ipitr)->barcode());
               num_w_nu++;
               break;
              case 14:
               gen_nu.push_back(**opitr);
               bar_w_nu.push_back((*ipitr)->barcode());
               num_w_nu++;
               break;
              case 1:
               gen_q.push_back(**opitr);
               bar_w_q.push_back((*ipitr)->barcode());
                 num_w_q++;
               break;
            } 
         }// w decay product check
         num_w_decay++;
       } //w 
      }// in comming particle should be one
      if(num_w_decay == 4) 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_w_decay= " <<  num_w_decay << " num_w= " <<  num_w << endreq;
  mLog <<MSG::DEBUG << "num_w_lep= " <<  num_w_lep << endreq;
  mLog <<MSG::DEBUG << "num_w_nu= " <<  num_w_nu << " num_w_q= " <<  num_w_q << endreq;
  if( num_w_lep == 2 && lepSign ==1 && num_w_decay == 4){
    double lep_Pt = 0.;
    double nu_Pt= 0.;
    HepLorentzVector lep,nu;
    for(int ilep =0;ilep < 2; ilep++){  // long lepton for loop
      int bar_h_lep =0; // higss barcode of the lepton
      lep = gen_lep[ilep].momentum(); // filled in pair
      nu = gen_nu[ilep].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());
      for(int iw=0;iw<4;iw++){
        if( gen_w[iw].barcode() == bar_w_lep[ilep]){
          bar_h_lep=bar_h_w[iw];
          break;
        } 
      } 
      double mT;
      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 {
        mT = std::sqrt(m2);
        m_h_tranverseMofW->Fill(mT,1.);
        m_h_nuPt->Fill(nu_Pt,1);
        m_aan_t_nuPt->push_back(nu_Pt);  // it is temporaly treatment. no observable
        m_aan_t_tranverseMofW->push_back(mT); //AAN vector filling
      }
      HepLorentzVector w_lnu;
      double mlnu=0.;
      w_lnu = lep + nu;
      mlnu = sqrt(w_lnu.e()*w_lnu.e()-(w_lnu.px()*w_lnu.px()+w_lnu.py()*w_lnu.py()+w_lnu.pz()*w_lnu.pz()) ); 
      m_aan_t_lnuMass->push_back(mlnu);
      m_h_lnu->Fill(mlnu,1. );
      HepLorentzVector q1,q2,p_w,p_h;
      bool right_h_found = false;
      int iq_found = 0;
      if(num_w_q == 4 ){
         for(int iq=0; iq < 4 ; iq++){ // search right q jet
           for(int iw=0;iw<4;iw++){
              if( gen_w[iw].barcode() == bar_w_q[iq]){
                if(bar_h_lep == bar_h_w[iw]) right_h_found = true;
                break;
              }
           }
           if(right_h_found){
               iq_found = iq;
               break;//break with right iq
           }
         }
         if(iq_found < 3 && right_h_found ){
           q1 = gen_q[iq_found].momentum();
           q2 = gen_q[iq_found+1].momentum(); // should be pair
           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.);
           m_aan_t_qqMass->push_back(mqqb);
         }
      }
      if(num_w_q == 4 && (num_w_lep == 2) && right_h_found ){
         m_aan_t_size++;
         //check result if gen w mass known
         p_h = w_lnu + p_w;
         double mH=sqrt(p_h.e()*p_h.e()-(p_h.px()*p_h.px()+p_h.py()*p_h.py()+p_h.pz()*p_h.pz()));
         m_h_hmass->Fill( mH,1.);
         m_aan_t_hmass->push_back(mH);
      } // qq jet 4 and 2 lepton if
    } // end of 2 lepton for loop
  } // 2 same sign lepton and 4 w if 
  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;
}