/************************************************************************
*  WGRAD_HistModule                                                     *
*  ----------------                                                     *
*                                                                       *
*  D. Waters                                                            *
*  21.8.99                                                              *
*************************************************************************/

//-----------------------
// This Class's Header --
//-----------------------
#include "Electroweak/EW_MC_Validation.hh"
//-------------
// C Headers --
//-------------
#include <assert.h>
 //------------------------------- 
// Collaborating Class Headers --
//---------------------------

#include "BaBar/Cdf.hh"
#ifdef USE_CDFEDM2
#include "Edm/ConstHandle.hh"
#include "Edm/EventRecord.hh"
#include "SimulationObjects/HEPG_StorableBank.hh"
#else
#include "Banks/HEPG_Bank.hh"
#include "Trybos/TRY_Abstract_Record.hh"
#endif // USE_CDFEDM2

#include "CLHEP/Matrix/Vector.h"
#include "CLHEP/Vector/LorentzVector.h"
#include "HepTuple/HepHist1D.h"
#include "Trybos/TRY_Abstract_Record.hh"

// -------------------------------------------------
// -- Declarations of variables defined elsewhere --
// -------------------------------------------------

//-----------------------------------------------------------------------
// Local Macros, Typedefs, Structures, Unions and Forward Declarations --
//-----------------------------------------------------------------------

static const char rcsid[] = " ";

//----------------
// Constructors --
//----------------
EW_MC_Validation::EW_MC_Validation(const char* const theName, 
				   const char* const theDescription )
  : HepHistModule( theName, theDescription)
{
}

//--------------
// Destructor --
//--------------
EW_MC_Validation::~EW_MC_Validation( )
{
}

//--------------
// Operations --
//--------------

AppResult EW_MC_Validation::bookHistograms( )
{
  HepFileManager* manager = fileManager( );
  assert(0 != manager);

  _ChargedLeptonPt = &manager->hist1D("Charged Lepton pT (GeV)", 200 ,0.0,200.0,100);
  assert(0 != _ChargedLeptonPt);

  _NeutrinoPt = &manager->hist1D("Neutrino pT (GeV)", 200,0.0,200.0,101);
  assert(0 != _NeutrinoPt);

  _ChargedLeptonPairMass = &manager->hist1D("Charged Lepton Pair Mass (GeV)",1000,0.0,1000.0,102);
  assert(0 != _ChargedLeptonPairMass);

  _ChargedLeptonNeutrinoMt = &manager->hist1D("Charged Lepton Neutrino Mt (GeV)",1000,0.0,1000.0,103);
  assert(0 != _ChargedLeptonNeutrinoMt);

  _FinalStateMultiplicity = &manager->hist1D("Final State Multiplicity",500,0.0,500.0,104);
  assert(0 != _FinalStateMultiplicity);

  return AppResult::OK;
}

AppResult EW_MC_Validation::fillHistograms( AbsEvent* anEvent )
{
  int finalStateParticles = 0;
  int numberOfChargedLeptons = 0;
  int numberOfNeutrinos = 0;
  HepLorentzVector lepton1, lepton2;

  //---------------------------------------------------------------------------//
  // Find the HEPG bank.                                                       //
  //---------------------------------------------------------------------------//
#ifdef USE_CDFEDM2
  EventRecord::ConstIterator hepgIter(anEvent,"HEPG_StorableBank");
#else
  TRY_Record_Iter_Same hepgIter(anEvent,"HEPG");
#endif // USE_CDFEDM2
  if (!hepgIter.is_valid())
    {
      std::cerr << " EW_MC_Validation : ERROR - HEPG Bank does not exist " << std::endl;
      return AppResult::OK;
    }
#ifdef USE_CDFEDM2
  const HEPG_StorableBank& hepg = *(ConstHandle<HEPG_StorableBank>(hepgIter));
#else
  HEPG_Bank hepg(hepgIter);
#endif // USE_CDFEDM2
  //---------------------------------------------------------------------------//
  // Iterate over the entries in the bank.                                     //
  //---------------------------------------------------------------------------//
#ifdef USE_CDFEDM2
  for (HEPG_StorableBank::particleIter partIter(hepg); partIter.is_valid(); ++partIter)
#else
  for (HEPG_Bank::particleIter partIter(hepg); partIter.is_valid(); ++partIter)
#endif // USE_CDFEDM2
    {
      //---------------------------------------------------------------------------//
      // Look for "initial entry" charged leptons and neutrinos. WGRAD does not    //
      // currently store documentation HEPG entries that would indicate that the   //
      // leptons are daughters of a W boson.                                       //
      //---------------------------------------------------------------------------//
      bool initialEntry = false;
      if (hepg.jmo1hep(partIter) == 0) initialEntry = true;

      //---------------------------------------------------------------------------//
      // Look for W and Z daughters.                                               //
      //---------------------------------------------------------------------------//
      // Remember to subtract 1 !!!
      int parentCode = abs(hepg.idhep(hepg.jmo1hep(partIter)-1)); 
      bool W_or_Z_Daughter = false;
      if (parentCode == 23 || parentCode == 24) W_or_Z_Daughter = true;

      //---------------------------------------------------------------------------//
      // Look for leptons.                                                         //
      //---------------------------------------------------------------------------//
      int partCode = abs(hepg.idhep(partIter));
      bool chargedLepton = false;
      bool neutrino = false;
      if (partCode == 11 || partCode == 13 || partCode == 15) chargedLepton = true;
      if (partCode == 12 || partCode == 14 || partCode == 16) neutrino = true;
 
      //---------------------------------------------------------------------------//
      // Look for all final state particles.                                       //
      //---------------------------------------------------------------------------//
      if (hepg.istdhep(partIter) == 1) 
	{
	  finalStateParticles++;
	}

      std::cout << " Daughter : " << hepg.jda1hep(partIter)-1 << std::endl;


      if ((initialEntry || W_or_Z_Daughter) && (chargedLepton || neutrino))
	{
	  //---------------------------------------------------------------------------//
	  // This is a lepton from W or Z decay.                                       //
	  //---------------------------------------------------------------------------//
	  HepLorentzVector lepton(hepg.Px(partIter),
				  hepg.Py(partIter),
				  hepg.Pz(partIter),
				  hepg.E(partIter));

	  //---------------------------------------------------------------------------//
	  // Fill histograms with individual particle information.                     //
	  //---------------------------------------------------------------------------//
	  if (chargedLepton)
	    {
	      _ChargedLeptonPt->accumulate(lepton.perp());
	    }
	  else if (neutrino)
	    {
	      _NeutrinoPt->accumulate(lepton.perp());
	    }
	    
	  //---------------------------------------------------------------------------//
	  // Store the first two leptons to calculate masses etc.                      //
	  //---------------------------------------------------------------------------//
	  if (chargedLepton) numberOfChargedLeptons++;
	  if (neutrino) numberOfNeutrinos++;
	  int leptonCount = (numberOfChargedLeptons+numberOfNeutrinos);
	  if (leptonCount == 1) lepton1 = lepton;
	  if (leptonCount == 2) lepton2 = lepton;
	}
    }
  
  if ((numberOfChargedLeptons+numberOfNeutrinos) == 2)
    {
      //---------------------------------------------------------------------------//
      // Combine lepton 4-momenta.                                                 //
      // Calculate invariant mass for Z events, transverse mass for W events.      //
      //---------------------------------------------------------------------------//
      HepLorentzVector gaugeBoson = lepton1 + lepton2;
      if (numberOfChargedLeptons == 2)
	{
	  _ChargedLeptonPairMass->accumulate(gaugeBoson.m());
	}
      else if (numberOfChargedLeptons == 1)
	{
	  _ChargedLeptonNeutrinoMt->accumulate(sqrt(2.0*(lepton1.perp()*lepton2.perp() - 
							 lepton1.px()*lepton2.px() - 
							 lepton1.py()*lepton2.py())));
	}
      else if (numberOfChargedLeptons == 0)
	{
	  // Fill Z->vv histograms ?
	}
    }

  //---------------------------------------------------------------------------//
  // Fill histograms with global event information.                            // 
  //---------------------------------------------------------------------------//
  _FinalStateMultiplicity->accumulate(finalStateParticles);

  return AppResult::OK;
  
}

AppModule*
EW_MC_Validation::clone(const char* cloneName)
{
  return new EW_MC_Validation(cloneName,"this module is a clone EW_MC_Validation");
}

const char * 
EW_MC_Validation::rcsId( ) const
{
   return rcsid;
}