/****************************************************************\
| This file is part of the PAX example files for CDF             |
| -------------------------------------------------------------- | 
|   Author of the last change: $Author: erdmann $                 
|   Date of last change:       $Date: 2003/11/18 15:13:17 $   
|   Revision number:           $Revision: 1.3 $        
\****************************************************************/
//_____________________________________________________________________________
// PAX analysis module: below are examples of what 
// one can do in the interactive session
//_____________________________________________________________________________

#include "TCanvas.h"
#include "TF1.h"

#include "TPAXAnaModule.hh"

#include "PAX/PAXcollision.hh"
#include "PAX/PAXexperimentclass.hh"
#include "PAX/PAXfile.hh"
#include "PAX/PAXfourvector.hh"
#include "PAX/PAXid.hh"
#include "PAX/PAXinterpret.hh"
#include "PAX/PAXiterator.hh"
#include "PAX/PAXmap.hh"
#include "PAX/PAXphysicsmap.hh"
#include "PAX/PAXprintlevel.hh"
#include "PAX/PAXvertex.hh"

#include <iostream>

//_____________________________________________________________________________
TPAXAnaModule::TPAXAnaModule(const char* name, const char* title):
  TStnModule(name,title){
}

//_____________________________________________________________________________
TPAXAnaModule::~TPAXAnaModule() {
}

//_____________________________________________________________________________
void TPAXAnaModule::BookHistograms() {
  char name [200];
  char title[200];

  Delete("hist");

  // book histograms 

  sprintf(name, "%s_calo_et",GetName());
  sprintf(title,"Calo: Et");
  fHist.fCaloEt = new TH1F(name, title, 29,1,30);
  AddHistogram(fHist.fCaloEt);

  sprintf(name, "%s_sum_et",GetName());
  sprintf(title,"Calo: Sum(Et)");
  fHist.fSumEt = new TH1F(name, title, 10,0,300);
  AddHistogram(fHist.fSumEt);

  sprintf(name, "%s_W_pt",GetName());
  sprintf(title,"W: pt");
  fHist.fWpt = new TH1F(name, title, 8,0,100);
  AddHistogram(fHist.fWpt);

  sprintf(name, "%s_W_eta",GetName());
  sprintf(title,"W: #eta");
  fHist.fWeta = new TH1F(name, title, 8,-8,8);
  AddHistogram(fHist.fWeta);
}


//_____________________________________________________________________________
int TPAXAnaModule::BeginJob() {

  // register the data block

  RegisterDataBlock( "CalDataBlock" , "TCalDataBlock"    , &fCalData      );
  RegisterDataBlock( "VertexBlock"  , "TStnVertexBlock"  , &fVertexData   );
  RegisterDataBlock( "ElectronBlock", "TStnElectronBlock", &fElectronData );
  RegisterDataBlock( "MetBlock"     , "TStnMetBlock"     , &fMetData      );

  // book histograms

  BookHistograms();

  // open file to write PAX event interpretations

  pax_file_output = new PaxFile( "pax.root", PaxFile::Write );

  //--------------------------------------------------------------------------
  // disable module if branch we're interested in doesn't exist
  //--------------------------------------------------------------------------

  if (! fCalData) fEnabled = 0;

  return 0;
}

//_____________________________________________________________________________
int TPAXAnaModule::BeginRun() {
  return 0;
}   

//_____________________________________________________________________________
int TPAXAnaModule::Event(int ientry) {

  if ( ientry%100 == 0) 
    cout<<"----------------- event "<<ientry<<" -----------------"<<endl;

  //-----------------------------------------------------------
  // 0. fill event interpretation with calorimeter data
  //-----------------------------------------------------------
  
  // create empty event interpretation

  PaxEventInterpret* cal_towers = new PaxEventInterpret; 

  // fill event interpretation with calorimeter data

  FillCalorimeter( cal_towers, ientry );

  //-----------------------------------------------------------
  // 1. example for accessing information: for first event
  //    print all information available to event interpretation
  //-----------------------------------------------------------

  if ( ientry==0 )
    {
      cout << "************************** long print of calorimeter information"
	   << ", only for first event **************************" << endl;

      cal_towers->Print();
    }

  //----------------------------------------------------------------
  // 2. example for accessing information:
  //    print calorimeter transverse energies & histogram them
  //----------------------------------------------------------------

  if ( ientry==0 )
    cout << "************************** print calorimeter transverse momenta"
	 << ", only for first event **************************" << endl;

  // iterator for PaxFourVectors

  PaxIterator< PaxFourVector* >* iter = cal_towers->fourvector->GetIterator();

  // loop over all calorimeter fourvectors

  for ( iter->First(); !iter->IsDone(); iter->Next() )
    {
      // access transverse momentum of calorimeter fourvector

      double pt = iter->CurrentItem()->Pt();

      // fill histogram
      
      fHist.fCaloEt->Fill( pt );                       
  
      // for first event: print

      if ( ientry==0 )
	cout  << "name = " << iter->CurrentItem()->GetPaxName() 
	      << "  pt = " << iter->CurrentItem()->Pt() << " GeV" << endl;
    }

  //------------------------------------------------------------
  // 3. example for accessing information: 
  //    access vertices & print their positions, 
  //    get the fourvectors related to the vertices, print them,
  //    histogram the sum of the transverse energies
  //------------------------------------------------------------

  if ( ientry==0 )
    cout << "************************** print vertices and related calorimeter cells"
	 << ", only for first event **************************" << endl;

  // sum up transverse energies

  double sum_Et = 0;

  // get iterator for vertices from event interpretation

  PaxIterator<PaxVertex* >*  iter_vx = cal_towers->vertex->GetIterator();

  // loop over all vertices

  for (iter_vx->First(); !iter_vx->IsDone(); iter_vx->Next() )
    {
      if ( ientry==0 )
	cout << "name = " << iter_vx->CurrentItem()->GetPaxName() 
	     << "   x = " << iter_vx->CurrentItem()->X() 
	     << "   y = " << iter_vx->CurrentItem()->Y() 
	     << "   z = " << iter_vx->CurrentItem()->Z() 
	     << endl;

      // sum up transverse energies for the calorimeter towers 
      // associated to the primary vertex

      if ( iter_vx->CurrentItem()->GetPaxName() == "primary vertex" )
	{
	  // get iterator for fourvectors via the outgoing fourvector relations
	  // of the vertices of event interpretation
      
	  PaxIterator<PaxFourVector* >* iter_fv = iter_vx->CurrentItem()
	    ->outgoing_fourvector_relations->GetIterator();

	  // loop over the calorimeter fourvectors

	  for (iter_fv->First(); !iter_fv->IsDone(); iter_fv->Next() )
	    {
	      // sum transverse energies

	      sum_Et += iter_fv->CurrentItem()->Pt();

	      // print information on calorimeter fourvectors for first event

	      if ( ientry==0 )
		cout << "name = " << iter_fv->CurrentItem()->getPaxName() 
		     << " pt = "              << iter_fv->CurrentItem()->Pt() 
		     << " pseudo-rapidity = " << iter_fv->CurrentItem()->Eta() 
		     << " phi = "             << iter_fv->CurrentItem()->Phi() 
		     << endl;
	    }
	}
    }

  // fill histogram
      
  fHist.fSumEt->Fill( sum_Et );                       

  //-------------------------------------------------------------------
  // 4. example for merging information of two event interpretations:
  // calorimeter towers and electrons are merged without double
  // counting of energy, using the eta-phi-code EtaPhi() of the 
  // calorimeter towers
  //-------------------------------------------------------------------

  // create empty event interpretation

  PaxEventInterpret* electrons = new PaxEventInterpret; 

  // fill event interpretation with electron data

  FillElectrons( electrons, ientry );

  // create empty event interpretation

  PaxEventInterpret* calorimeter_with_electrons = new PaxEventInterpret; 

  // merge electron information with calorimeter tower information

  int number_of_electrons = MergeElectrons( cal_towers, electrons, 
					    calorimeter_with_electrons );

  // print information on the number of succesfully merged electrons

  cout << "number of merged electrons = " << number_of_electrons 
       << " out of " << electrons->fourvector->Size() 
       << " electrons found in the electron block" << endl;

  //-------------------------------------------------------------------
  // 5. example analysis: reconstruct W -> electron + neutrino
  //-------------------------------------------------------------------

  // create two empty event interpretations for two possible W solutions

  PaxEventInterpret* W1 = new PaxEventInterpret;
  PaxEventInterpret* W2 = new PaxEventInterpret;

  // select highest-pt electron for W reconstruction

  double         pt_max        =  0;
  PaxFourVector* electron      =  0;
  string         electron_name = "electron";

  // get iterator for fourvectors 
  // PaxIterator< PaxFourVector* >* iter : has been defined above

  iter = calorimeter_with_electrons->fourvector->GetIterator();

  // loop over merged event interpretation

  for ( iter->First(); !iter->IsDone(); iter->Next() )
    {    
      // select only highest-pt electron

      if (    iter->CurrentItem()->GetPaxName() == electron_name
	   && iter->CurrentItem()->Pt()         >  pt_max )
	{
	  electron = iter->CurrentItem();
	  pt_max   = iter->CurrentItem()->Pt();
	}
    }

  // make W if event has an electron

  if ( electron )
    {
      // get missing transverse momenta from Stntuple

      double px = 0;
      double py = 0;
  
      if (fMetData)
	{
	  fMetData->GetEntry(ientry);
	  
	  px = fMetData->MetX(1);
	  py = fMetData->MetY(1);
	}
      else
	{
	  cout << "TPAXAnaModule::Event : could not find fMetData" << endl;
	}
      
      // mass of W boson
  
      double W_mass = 80.42;

      // make two event interpretations for both W solutions

      int number_of_W = MakeW( calorimeter_with_electrons, W1, W2, 
			       electron, px, py, W_mass );

      // find the solution with the smaller |pseudorapidity|
      
      double W_eta_min = 0;
      double W_pt      = 0;
      
      cout << "************************** decay trees of " << number_of_W 
	   << " reconstructed W bosons **************************" << endl;
      
      if ( !W1->IsEmpty() ) 
	{
	  // PaxIterator< PaxFourVector* >* iter : has been defined above
	  
	  iter = W1->fourvector->GetIterator();
	  
	  for ( iter->First(); !iter->IsDone(); iter->Next() )
	    {    
	      if ( iter->CurrentItem()->GetPaxName() == "W" )
		{
		  // remember W pseudorapidity, pt
		  
		  W_eta_min = iter->CurrentItem()->Eta();
		  W_pt      = iter->CurrentItem()->Pt();
		  
		  // print the decay tree of the W boson
		  
		  iter->CurrentItem()->PrintDecayTree(PAX_PTSHORT);
		}
	    }
	}
      
      if ( !W2->IsEmpty() ) 
	{
	  iter = W2->fourvector->GetIterator();
	  
	  for ( iter->First(); !iter->IsDone(); iter->Next() )
	    {    
	      if ( iter->CurrentItem()->GetPaxName() == "W" )
		{
		  // remember this W pseudorapidity, if smaller than the
		  // above solution 
		  
		  if ( abs( iter->CurrentItem()->Eta() ) < abs( W_eta_min ) )
		    {
		      W_eta_min = iter->CurrentItem()->Eta();
		      W_pt      = iter->CurrentItem()->Pt();
		    }
		  
		  // print decay tree
		  
		  iter->CurrentItem()->PrintDecayTree(PAX_PTSHORT);
		}
	    }
	}
      
      // fill histograms
      
      fHist.fWpt ->Fill( W_pt );                       
      fHist.fWeta->Fill( W_eta_min );                       
    }

  //------------------------------------------------------------------------
  // 6. write W event interpretations into file for later analysis
  //------------------------------------------------------------------------

  if ( !W1->IsEmpty() ) W1->Store(pax_file_output);
  if ( !W2->IsEmpty() ) W2->Store(pax_file_output);

  // end of event: write it

  pax_file_output->WriteEvent();

  //--------------------------------------------------
  // 7. clean up memory
  //--------------------------------------------------

  delete cal_towers;
  delete electrons;      
  delete calorimeter_with_electrons;
  delete W1;
  delete W2;

  // for every pax object, e.g., vertex, fourvector, we assign
  // a unique identifier which you can check using GetPaxId(),
  // we recommend to reset the id generator for each event:

  PaxIdGenerator::Instance()->ResetIntKey();

  return 0; 		       
}

//_____________________________________________________________________________
void TPAXAnaModule::DisplayEvent() {

  mycanvas = new TCanvas("PAX", "PAX",0,0,800,600);
  mycanvas->Divide(2,2);

  fHist.fCaloEt->SetXTitle("Et [GeV]");
  fHist.fCaloEt->SetYTitle("N");
  fHist.fSumEt ->SetXTitle("Et [GeV]");
  fHist.fSumEt ->SetYTitle("N");
  fHist.fWpt   ->SetXTitle("Et [GeV]");
  fHist.fWpt   ->SetYTitle("N");
  fHist.fWeta  ->SetXTitle("#eta");
  fHist.fWeta  ->SetYTitle("N");

  mycanvas->cd(1);
  fHist.fCaloEt->Draw();
  mycanvas->cd(2);
  fHist.fSumEt ->Draw();
  mycanvas->cd(3);
  fHist.fWpt   ->Draw();
  mycanvas->cd(4);
  fHist.fWeta  ->Draw();

  mycanvas->Modified();
  mycanvas->Update();
}

//_____________________________________________________________________________
int TPAXAnaModule::EndJob() {

  printf("----- end job: ---- %s\n",GetName());

  pax_file_output->Close();
  delete pax_file_output;
  delete mycanvas;

  return 0;
}

//_____________________________________________________________________________