#include "MyAODAna/MyZmumu.h"
#include "GaudiKernel/MsgStream.h"

#include "StoreGate/StoreGateSvc.h"

#include "HepMC/GenEvent.h"
#include "HepMC/GenVertex.h"
#include "HepMC/GenParticle.h" 


#include "ParticleEvent/TruthParticle.h"

#include "ParticleEvent/TruthParticleContainer.h"

#include "ParticleEvent/MuonContainer.h"

#include <math.h>

#include "GaudiKernel/IHistogramSvc.h"
#include "AIDA/IHistogram1D.h"

#include "GaudiKernel/INTupleSvc.h"
#include "GaudiKernel/SmartDataPtr.h"  /* NTupleFilePtr */


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

MyZmumu::MyZmumu(const std::string& name, ISvcLocator* pSvcLocator) :
Algorithm(name, pSvcLocator)
{

// Part 2: Properties go here


}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

StatusCode MyZmumu::initialize(){
  
  // Part 1: Get the messaging service, print where you are
  MsgStream mLog(msgSvc(), name());
  mLog << MSG::INFO << "initialize()" << endreq;
  
  StatusCode sc = service("StoreGateSvc", m_storeGate);
  if (sc.isFailure()) {
    mLog << MSG::ERROR
	 << "Unable to retrieve pointer to StoreGateSvc"
	<< endreq;
    return sc;
  }

  // Book histogram
  m_hptreso = histoSvc()-> book("pt",1,"Pt resolution",100,-4000.,4000.);
  
  if (0 == m_hptreso) {
    mLog << MSG::ERROR << " ERROR booking histogram" << endreq;
    return StatusCode::FAILURE;
  }
  
  
  //Book ntuple
  // Create a column-wise ntuple
  NTuple::Directory *col=0;
  NTupleFilePtr file1(ntupleSvc(), "/NTUPLES/FILE1");
  
  if ( 0 != file1 )    {
    if ( !(col=ntupleSvc()->createDirectory("/NTUPLES/FILE1/COL")) )   {
      mLog << MSG::ERROR << "Cannot create directory /NTUPLES/FILE1/COL" << endreq;  
    }
    
  } else {
    mLog << MSG::ERROR << "Cannot access file /NTUPLES/FILE1" << endreq;
    return StatusCode::FAILURE;
  }  
  
  NTuplePtr nt1(ntupleSvc(), "/NTUPLES/FILE1/COL/10");
  // Check if already booked
  if ( 0 == nt1 )    {
    nt1 = ntupleSvc()->book (col, 10, CLID_ColumnWiseTuple, "Muons");
    if ( 0 != nt1 ) {
      p_nt1 = nt1;
      sc = nt1->addItem ("Nmu",m_nmu,0,10);
      if ( sc == StatusCode::FAILURE ) {
	mLog << MSG::ERROR << "could not add item to col wise ntuple" 
	     << endreq;
      }
      sc = nt1->addIndexedItem ("Charge",m_nmu,m_charge);
      sc = nt1->addIndexedItem ("Pt",m_nmu,m_pt);
      sc = nt1->addIndexedItem ("Eta",m_nmu,m_eta);
      sc = nt1->addIndexedItem ("Chargef",m_nmu,m_chargef);
      sc = nt1->addIndexedItem ("Ptf",m_nmu,m_ptf);
      sc = nt1->addIndexedItem ("Etaf",m_nmu,m_etaf);
    } else {
      mLog << MSG::ERROR  << "could not book Ntuple" << endreq;
    }
  } else {
    mLog << MSG::ERROR << "Ntuple is already booked" << endreq;
  }
  
  
  
  return StatusCode::SUCCESS;
}

// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

StatusCode MyZmumu::execute() {
  
  // Part 1: Get the messaging service, print where you are
  MsgStream mLog(msgSvc(), name());
  mLog << MSG::INFO << "execute()" << endreq;
  
  // Part 2: Print out the different levels of messages
  mLog << MSG::INFO << "Your new package and algorithm are successfully installed" << endreq; 
  /// get the MC truth particle AOD container from StoreGate
  StatusCode sc = mymctruth();
  if (sc.isFailure()) {
    mLog << MSG::ERROR << "The method mymctruth() failed" << endreq;
    return StatusCode::SUCCESS;
  }

  /// get the Muon AOD container from StoreGate
  sc = mymuonevt();
  if (sc.isFailure()) {
    mLog << MSG::ERROR << "The method mymuonevt() failed" << endreq;
    return StatusCode::SUCCESS;
  }
  
  m_nmu=0;
  for (int it=0; it<2;it++){
    mLog << MSG::INFO << " muone\n "
	 <<" carica MC "<<muone[it].charge
	 <<" carica "<<muone[it].chargef
	 <<"\n Pt MC "<<muone[it].pt
	 <<" Pt "<<muone[it].ptf
	 <<"\n eta MC "<<muone[it].eta
	 <<" eta "<<muone[it].etaf
	 << endreq;
    m_hptreso->fill(muone[it].ptf-muone[it].pt,1.);
    m_chargef[m_nmu]=muone[it].chargef;
    m_charge[m_nmu]=muone[it].charge;
    m_ptf[m_nmu]=muone[it].ptf;
    m_pt[m_nmu]=muone[it].pt;
    m_etaf[m_nmu]=muone[it].etaf;
    m_eta[m_nmu]=muone[it].eta;
    m_nmu = m_nmu+1;
  }

  // Fill ntuple
  if (!(ntupleSvc()->writeRecord(p_nt1)).isSuccess()) {
    mLog << MSG::ERROR << "problems writing ntuple record" << endreq;
  }

  
  return StatusCode::SUCCESS;
}

//*********************** mctruth method to be called event by event

StatusCode MyZmumu::mymctruth() {
  
  // Part 1: Get the messaging service, print where you are
  MsgStream mLog(msgSvc(), name());
  mLog << MSG::INFO << "mymctruth()" << endreq;
  
  const TruthParticleContainer*  mcParts = 0;
  StatusCode sc=m_storeGate->retrieve( mcParts,"SpclMC" );
  if( sc.isFailure()  ||  !mcParts ) {
    mLog << MSG::WARNING
	<< "No AOD MC truth particle container found in TDS"
	<< endreq; 
    return StatusCode::FAILURE;
  }
  mLog <<MSG::DEBUG << "MC Truth Container Successfully Retrieved" << endreq;
  
  //  TruthParticleContained eredita da DataVector.h
  unsigned int iPart = 0;
  for ( TruthParticleContainer::const_iterator itr = mcParts->begin();
	itr != mcParts->end();
	++itr, ++iPart ) {
    if ((*itr)->pdgId()==PDG::Z0){
      for ( unsigned int iChild = 0; iChild != (*itr)->nDecay(); ++iChild ){
	const TruthParticle * child = (*itr)->getChild( iChild );
	if (child->pdgId() == PDG::mu_plus) 	    
	  {muone[0].charge=+1; 
	  muone[0].pt=pow(child->px()*child->px()+child->py()*child->py(),0.5);
	  muone[0].eta=-0.5*log((child->e()-child->pz())/(child->e()+child->pz()));
	  }
	else if (child->pdgId() == PDG::mu_minus) 	  
	  {muone[1].charge=-1;  
	  muone[1].pt=pow(child->px()*child->px()+child->py()*child->py(),0.5);
	  muone[1].eta=-0.5*log((child->e()-child->pz())/(child->e()+child->pz()));
	  }
      }
    }
    
//    mLog << MSG::INFO << "Part " << iPart
//	<< " PDG-ID: " << (*itr)->pdgId()
//	<< " nChildren: " << (*itr)->nDecay()
//	<< " status: " << (*itr)->getGenParticle()->status()
//	<< " bc: " << (*itr)->getGenParticle()->barcode()
//	<< endreq;
//    for ( unsigned int iChild = 0; iChild != (*itr)->nDecay(); ++iChild ){
//      const TruthParticle * child = (*itr)->getChild( iChild );
//      if ( 0 != child ) {
//	mLog << MSG::INFO
//	    << "\tchild: " << iChild
//	    << "\tPDGID: " << child->pdgId()
//	    << " status: " << child->getGenParticle()->status()
//	    << " bc: "     << child->getGenParticle()->barcode()
//	    << " bc Parents: " << child->nParents() << " [ ";
//	for ( unsigned int iMoth = 0; iMoth != child->nParents(); ++iMoth ) {
//	  mLog << child->getGenMother(iMoth)->barcode() << " ";
//	}
//	mLog << "]" << endreq;
//      } else {
//	mLog << MSG::WARNING << "Wrong pointer to child !!" << endreq;
//      }
//    }//> loop over children
  }//> end loop over TruthParticles
  
  return StatusCode::SUCCESS;
  
}

//*********************** mctruth method to be called event by event

StatusCode MyZmumu::mymuonevt() {
  
  // Part 1: Get the messaging service, print where you are
  MsgStream mLog(msgSvc(), name());
  mLog << MSG::INFO << "mymuonevt()" << endreq;
  
  const MuonContainer*  muon = 0;
  StatusCode sc=m_storeGate->retrieve( muon,"MuidMuonCollection" );
  if( sc.isFailure()  ||  !muon ) {
    mLog << MSG::WARNING
	<< "No MuidMuonCollection particle container found in AOD"
	<< endreq; 
    return StatusCode::FAILURE;
  }
  mLog <<MSG::DEBUG << "MuidMuonCollection Container Successfully Retrieved" << endreq;
  
  //  TruthParticleContained eredita da DataVector.h
  unsigned int iPart = 0;
  for ( MuonContainer::const_iterator itr = muon->begin();
	itr != muon->end();
	++itr, ++iPart ) {
    mLog <<MSG::DEBUG 
	 << "Ipart muon " << iPart
	 << " Pt " << (*itr)->pt()
	 << " eta " << (*itr)->eta()
	 << " PDGID " << (*itr)->pdgId() 
	 << endreq;
    if ((*itr)->pdgId()==PDG::mu_plus) 
      {muone[0].chargef=static_cast<double>((*itr)->charge());
      muone[0].ptf=(*itr)->pt();
      muone[0].etaf=(*itr)->eta();
      }
    else if ((*itr)->pdgId()==PDG::mu_minus)
      {muone[1].chargef=static_cast<double>((*itr)->charge());
      muone[1].ptf=(*itr)->pt();
      muone[1].etaf=(*itr)->eta();}
  }//> end loop over muon
  
  return StatusCode::SUCCESS;
  
}
// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *

StatusCode MyZmumu::finalize() {

// Part 1: Get the messaging service, print where you are
MsgStream mLog(msgSvc(), name());
mLog << MSG::INFO << "finalize()" << endreq;

return StatusCode::SUCCESS;
}