//==================================================================
//
//	TZBosFnd.cc
//
//	Version 0.0:      P. Koehn    5/5/01
//
//==================================================================
#include "TF1.h"
#include <iostream>
#include <iomanip>
#include <algorithm>

#include "Stntuple/loop/TStnAna.hh"
#include "Stntuple/obj/TStnHeaderBlock.hh"
#include "TEleCand.hh"
#include "TMuoCand.hh"
#include "TJetCand.hh"

#include "TZBosFnd.hh"

//_____________________________________________________________________________
TZBosFnd::TZBosFnd(TFile* file,const char* name, const char* title):
  TStnModule(name,title){
//
// create a subdirectory "dir" in this file
   histFile = file;
   histDir = name;
   dir = histFile->mkdir(histDir);
//
// Initialize cuts
	init();
}
//_____________________________________________________________________________
TZBosFnd::TZBosFnd(const char* name, const char* title):
  TStnModule(name,title){
//
// create a subdirectory "dir" in this file
   histFile = new TFile("TEleCand.hbk","RECREATE","test file");
   histDir = name;
   dir = histFile->mkdir(histDir);
//
// Initialize cuts
	init();
}

void TZBosFnd::init() {
//
  totEvents = 0;
  passEvents = 0;
//
// Default is to make histograms
  enableHist(true);
//
// Default to not veto Z's
  setZveto(0);

// Default mass window
   setLowMassCut(75.);
   setHighMassCut(105.);

}

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

//_____________________________________________________________________________
void TZBosFnd::BookHistograms() {

//
// Delete the old histos
//   GetHistogramList()->Delete();
   histFile->cd(histDir);    // make the "dir" directory the current directory

//-----------------------------------------------------------------------------
//  electron histograms
//-----------------------------------------------------------------------------
   fHist.fEt      = new TH1F("fEt","Et of Z Legs",100,0.,100.);
   fHist.fEt1vEt2 = new TH2F("fEt1vEt2","Et1 vs Et2 of Z Legs",100,0.,100.,100,0.,100.);
   fHist.fDelPhi  = new TH1F("fDelPhi","Delta Phi of Legs",100,0.,M_PI);
   fHist.fDelR    = new TH1F("fDelR","Delta R of Z Legs",100,0.,5.);

//-----------------------------------------------------------------------------
//  jet histograms
//-----------------------------------------------------------------------------

//-----------------------------------------------------------------------------
//  muon histograms
//-----------------------------------------------------------------------------


//-----------------------------------------------------------------------------
//  missing Et histograms
//-----------------------------------------------------------------------------

//-----------------------------------------------------------------------------
//  other histograms
//-----------------------------------------------------------------------------
  fHist.fPairMass = new TH1F("fPairMass","Pair Mass",100,0.,200.);
  fHist.fZMass    = new TH1F("fZMass","Z Mass",50,LowMassCut-10.,HighMassCut+10.);
  fHist.fZPt	  = new TH1F("fZPt","Z Boson Pt",100,0.,50.);
  fHist.fZEta	  = new TH1F("fZEta","Z Boson Eta",100,-5.,5.);
  fHist.fZPhi	  = new TH1F("fZPhi","Z Boson Phi",100,0.,2*M_PI);
  fHist.fdType    = new TH1F("fdType","Z Decay Type",60,-0.5,19.5);
}


//_____________________________________________________________________________
int TZBosFnd::BeginJob() {

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

  totEvents = 0;
  passEvents = 0;
// register the data blocks
  pEleBlock = (TStnElectronBlock*) fAna->RegisterDataBlock("ElectronBlock","TStnElectronBlock");

// book histograms
  if (isEnableHist()) BookHistograms();

  if (! pEleBlock) {
    printf(" >>> branch *** %s *** doesn't exist \n","EleBlock");
    fEnabled = 0;
  }
  
// register the Muon data block
  pMuonBlock = (TStnMuonBlock*) fAna->RegisterDataBlock("MuonBlock","TStnMuonBlock");
  if (! pMuonBlock) {
    printf(" >>> branch *** %s *** doesn't exist \n","MuonBlock");
    fEnabled = 0;
  }
  
  printf("Cut Settings:\n");
  printf("  ZVeto = %i \n",Zveto);
  printf("  Mass Window %f < M < %f \n",LowMassCut,HighMassCut);

  return 0;

}


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

//==================================
//       Event(int ientry)
//==================================
int TZBosFnd::Event(int ientry) {

  bool pass   = false;
  if (fDebugLevel > 0) printf(" ---- ientry = %7i\n",ientry);
  totEvents++;
//
// Get the header info
  TStnHeaderBlock* header = GetHeaderBlock();
  int iev = header->EventNumber();
  int irun = header->RunNumber();
//
//
// Get access to the candidate modules (ele,muo,jet):
//---------------------------------------------------
  TStnAna * TestAna  = TStnModule::GetAna();
  TEleCand* eleCand  = (TEleCand*) TestAna->TStnAna::GetModule("EleC");
  TEleCand* LeleCand = (TEleCand*) TestAna->TStnAna::GetModule("EleLoose");
  TMuoCand* muoCand  = (TMuoCand*) TestAna->TStnAna::GetModule("MuoC");
  TMuoCand* LmuoCand = (TMuoCand*) TestAna->TStnAna::GetModule("MuoLoose");

// Set the event flag
   Zboson = false;
   leg1 = -1;
   leg2 = -1;
   decayType = -1; 

//
// Get the basic event info for the electrons
  int Nele = pEleBlock->NElectrons();
  int nelectrons = eleCand->getNgoodElectron();
  int Lnelectrons = LeleCand->getNgoodElectron();
  if (fDebugLevel > 0) printf("Elec Cand %i %i \n",nelectrons,Lnelectrons);
   
// Check if we have at least one tight electron and 
// that tight+loose is more than 2.
    if(nelectrons>0 & (nelectrons+Lnelectrons)>1) {

// Loop over the electron objects
       for(int nele1=0;  nele1<Nele-1; nele1++) {
	  for(int nele2=nele1+1; nele2<Nele; nele2++) {

// For this combination check that at least one is tight
	     if( (eleCand->IsagoodElectron(nele1) &  eleCand->IsagoodElectron(nele2)) |	
	         (eleCand->IsagoodElectron(nele1) & LeleCand->IsagoodElectron(nele2)) |
	        (LeleCand->IsagoodElectron(nele1) &  eleCand->IsagoodElectron(nele2)) ) {

		   if (fDebugLevel > 0) printf("Checking an Electron Pair");

// Get first & second electron
 		   TStnElectron* ele1 = pEleBlock->Electron(nele1);
		   TStnElectron* ele2 = pEleBlock->Electron(nele2);

// Reconstruct the Z boson...take the first lep. and add the second.
		   const TLorentzVector* p1 = (TLorentzVector*)ele1->Momentum();
		   const TLorentzVector* p2 = (TLorentzVector*)ele2->Momentum();
		   TLorentzVector *pair = new TLorentzVector();
		   pair->operator+=(*p1);
		   pair->operator+=(*p2);
		   fHist.fPairMass->Fill(pair->M());	
			   
// Check if this is inside the selection window
		   if(pair->M()>LowMassCut & pair->M()<HighMassCut) {
		        if(Zboson) printf("Warning: Event with 2 Z Candidates %i %i \n",irun,iev);
		        Zboson = true;
			Zmomentum = pair;
			leg1 = nele1;
			leg2 = nele2;
			decayType = (ele1->DetectorCode())+(ele2->DetectorCode());;			
                   } 
	     }
	  }
       }
    }   


//
// Get the basic event info for the Muons
  int Nmuo = pMuonBlock->NMuons();
  int nmuons = muoCand->getNgoodMuon();
  int Lnmuons = LmuoCand->getNgoodMuon();
  if (fDebugLevel > 0) printf("Muon Cand %i %i \n",nmuons,Lnmuons);


// Check if we have at least one tight muon and 
// that tight+loose is more than 2.
    if(nmuons>0 & (nmuons+Lnmuons)>1) {

// Loop over the muon objects
       for(int nmuo1=0;  nmuo1<Nmuo-1; nmuo1++) {
	  for(int nmuo2=nmuo1+1; nmuo2<Nmuo; nmuo2++) {

// For this combination check that at least one is tight
	     if( (muoCand->IsagoodMuon(nmuo1) &  muoCand->IsagoodMuon(nmuo2)) |	
	         (muoCand->IsagoodMuon(nmuo1) & LmuoCand->IsagoodMuon(nmuo2)) |
	        (LmuoCand->IsagoodMuon(nmuo1) &  muoCand->IsagoodMuon(nmuo2)) ) {

		   if (fDebugLevel > 0) printf("Checking a Muon Pair");

// Get first & second electron
 		   TStnMuon* muo1 = pMuonBlock->Muon(nmuo1);
		   TStnMuon* muo2 = pMuonBlock->Muon(nmuo2);

// Reconstruct the Z boson...take the first lep. and add the second.
		   const TLorentzVector* p1 = (TLorentzVector*)muo1->Momentum();
		   const TLorentzVector* p2 = (TLorentzVector*)muo2->Momentum();
		   TLorentzVector *pair = new TLorentzVector();
		   pair->operator+=(*p1);
		   pair->operator+=(*p2);
		   fHist.fPairMass->Fill(pair->M());	
			   
// Check if this is inside the selection window
		   if(pair->M()>LowMassCut & pair->M()<HighMassCut) {
		        if(Zboson) printf("Warning: Event with 2 Z Candidates %i %i \n",irun,iev);
			Zmomentum = pair;
			leg1 = nmuo1;
			leg2 = nmuo2;
			decayType = GetMuonCode(muo1->Detector(),muo2->Detector());			
			Zboson = true;
                   }		    
	     }
	  }
       }
    }   


// Make some plots for the found Z boson
  if(Zboson) {
  
// Z mass, Pt, eta, phi
	fHist.fZMass->Fill(Zmomentum->M());
	fHist.fZPt->Fill(Zmomentum->Pt());
	fHist.fZEta->Fill(Zmomentum->Eta());
	fHist.fZPhi->Fill(Zmomentum->Phi()+M_PI);

// Plot the quantities associated with each leg of the Z
	TLorentzVector* lep1;
	TLorentzVector* lep2;
	if(decayType<9) {
	    TStnElectron* ele1 = pEleBlock->Electron(leg1);
	    lep1 = (TLorentzVector*)ele1->Momentum();
	    TStnElectron* ele2 = pEleBlock->Electron(leg2);
	    lep2 = (TLorentzVector*)ele2->Momentum();
	} else if (decayType>9) {
	    TStnMuon* muo1 = pMuonBlock->Muon(leg1);
	    lep1 = (TLorentzVector*)muo1->Momentum();
	    TStnMuon* muo2 = pMuonBlock->Muon(leg2);
	    lep2 = (TLorentzVector*)muo2->Momentum();	
	}

// Et	
	fHist.fEt->Fill(lep1->Pt());
	fHist.fEt->Fill(lep2->Pt());
	fHist.fEt1vEt2->Fill(lep1->Pt(),lep2->Pt());
// Delta Phi & eta	
	double delphi = fabs(lep1->DeltaPhi(*lep2));
	fHist.fDelPhi->Fill(delphi);
	double delR = fabs(lep1->DeltaR(*lep2));
	fHist.fDelR->Fill(delR);
	fHist.fdType->Fill(decayType);		
  
  }



//  printf("njets %d nele %d \n",njets,nelectrons);
//
// Lets make our cuts for this event
//  Zveto=1 means filter out Z's, Zveto=-1 means filter to KEEP Z's, Zveto=0 no filter
  if( (Zveto==1 & !Zboson) | (Zveto==-1 & Zboson) | Zveto==0 ) {
    pass = true;
    passEvents++;
  }

//=================================================
// filter on events passing all cuts
//=================================================
  if (pass) {
     SetPassed(1);
  } else {
     SetPassed(0);
  }

//================================================
// Histogram the Electrons in events passing cuts=
//================================================
  if (pass) {
    if (fDebugLevel > 0) {
       printf(" Good event run %d event %d\n",irun,iev);
    }

//
  }
  if (ientry%1000==0)  printf("  Events processed %d passing %d \n",
  	totEvents,passEvents);
  return 0;		       
}


//==================================
//       EndJob()
//==================================
int TZBosFnd::EndJob() {

  printf("----- end job: ---- %s\n",GetName());
  printf("  Events processed %d passed %d \n",
  	totEvents,passEvents);
  return 0;
}

// ====================================
//    GetMuonCode
// ====================================
int TZBosFnd::GetMuonCode(int m1, int m2){

     int muonCode = -1;
     if(m1==1 | m2==1) {
        if(m1==1 && m2==1) muonCode = 10;
	if(m1==2 |  m2==2) muonCode = 11;
	if(m1==3 |  m2==3) muonCode = 12;
	if(m1==4 |  m2==4) muonCode = 13;
     } else if(m1==2 | m2==2) {
	if(m1==2 && m2==2) muonCode = 14;
	if(m1==3 |  m2==3) muonCode = 15;
	if(m1==4 |  m2==4) muonCode = 16;
     } else if(m1==3 | m2==3) {
	if(m1==3 && m2==3) muonCode = 17;
	if(m1==4 |  m2==4) muonCode = 18;
     } else if(m1==4 & m2==4) {
        muonCode = 19;
     }
 
      if(muonCode==-1) printf("TZBosFnd Warning: Illegal Muon Code: Input %i %i \n",m1,m2);
      return muonCode;    
}