// Z->ee actual data, tight electrons.  Not a fully debugged code.  Please replace this with a fully debugged code.
// Lynette Wehmer's code, August 22, 2003.  This uses the accepted good run list, with no silicon, GoodRunFcnNoSi.C, which
// is useful for the Winter 02-03 and Summer (up to May 03) data, and 
// needs to be loaded in root before this program.  For example, in root:
// root [0] .L GoodRunFcnNoSi.C+
// root [1] .x load.C     
// root [2] .L zee_tight.C     
// root [3] zsearch(&tc) 
// where load.C is the file that calls up the files containing the data.  
// This is for use with the UCNtuple 4.11.1

#include <iostream>
#include <math.h>
#include "TChain.h"
#include "TTree.h"
#include "TFile.h"
#include "TLeaf.h"
#include "TBranch.h"
#include "TObjString.h"
#include "TClonesArray.h"
#include "UCNtuple/myuber_t.hh"

#include "TH1F.h"


void declarestructs(TTree* tree, myuber_t* uber);

// Function definitions. Implementations in GoodRunFcn.C, GoodRunFcnNoSi.C.  These are good run lists.
bool good_run_nosi(int run);
//Baseline track quality cut
bool GoodTrackQ(mytrack_t& mytrack, int i);
//Baseline electron cut
bool StandardElectronQ(myelectron_t& myelectron, int i, mytrack_t& mytrack);
//Inv mass calculation.
Float_t invmass(Float_t et1, Float_t theta1, Float_t phi1, Float_t et2, Float_t theta2, Float_t phi2);


// To save the compiled program into the following file:
TFile *Zfile = new TFile("Zee_v1_junk.root", "RECREATE");

void zsearch(TTree* oldtree) {
  //begin Z search


// forward backward electron asymmetry
 TH1F *Afb = new TH1F("Afb", "Forward Backward Electron Asymmetry Versus Z Mass", 7, 60, 200);
 TH1F *zee = new TH1F("zee", "Z to ee Invariant Z Mass 4.11.1", 100, 0, 200);
 TH1F *zee_winter = new TH1F("zee_winter", "Z to ee Invariant Z Mass (Winter 02-03) 4.11.1", 100, 60, 120);
 TH1F *zee_summer = new TH1F("zee_summer", "Z to ee Invariant Z Mass (Summer 03) 4.11.1", 100, 60, 120);
 TH1F *ZEta = new TH1F("ZEta", "Eta of Z Candidate Electrons", 100, -1.5, 1.5);
 TH1F *runnumber = new TH1F("runnumber", "The Run Numbers of Passing Events", 100,140000,160000);
 TH1F *PEta = new TH1F("PEta", "Number of Positive Muons Versus Eta", 100, -1.5, 1.5);
 TH1F *NEta = new TH1F("NEta", "Number of Negative Muons Versus Eta", 100, -1.5, 1.5);


// finding the first and last run numbers
   int r = 200000;
   int r2 = 0; 

// Counting the number of runs and events
   int runno_new = 0;
   int numruns = 0;
   int eventno_new = 0;
   int numevents = 0;

// forward backward electron asymmetry
   Float_t A[100];
   Float_t F[100];
   Float_t B[100];

for (int o=0; o<100; o++)
   {
     A[o] = 0.0;     
     F[o] = 0.0;
     B[o] = 0.0;
   }


  myuber_t uber;
  Int_t treeno = -1;

  if (! oldtree) return;
//  cout << "Expect to process " << oldtree->GetEntries() << " entries" << endl;
  oldtree->SetBranchStatus("*", 1);
  oldtree->GetEntry(0);
  
#ifdef __CINT__
  gInterpreter->LoadMacro("declarestructs.C");
  //  gInterpreter->LoadMacro("/cdf/data1e/ucntuple/goodrun.C");
#endif

  declarestructs(oldtree, &uber);
  
  int npass = 0;

  for (int i = 0; i < oldtree->GetEntries(); i++) {
    // only way I can get this to be reliable
    oldtree->GetEntry(i);
    if (oldtree->IsA() == TChain::Class()) {
      if (treeno != ((TChain*) oldtree)->GetTreeNumber()) {
        declarestructs(oldtree->GetTree(), &uber);
        oldtree->GetEntry(i);
        treeno = ((TChain*) oldtree)->GetTreeNumber();
      }
    }


// finding the first and last run numbers
   int runno = uber.header.run;

   if (runno < r)
   {
      r = runno;
   }
   
   if (runno > r2)
   {
      r2 = runno;
   }


 // Checking runs to see if they are on the good run list
   if (!(good_run_nosi(runno))) continue;

// Counting the number of runs and events
   int eventno = uber.header.event;
   if (runno_new != runno)
   {
      runno_new = runno;
      numruns++;
   }   

   if (eventno_new != eventno)
   {
      eventno_new = eventno;
      numevents++;
   }



    if (i % 1000 == 0) {
      cout << "entry: " << i << endl;
    }

    int nel = 0;
    int elind[MAXELE];

    if (uber.electron.nele > 0) {
      for (int j = 0; j < uber.electron.nele; j++) {
	bool thiselpasses = StandardElectronQ(uber.electron, j, uber.track);  //tight CEM
	if (!thiselpasses) continue;
	elind[nel] = j;
	nel++;
	
      }
    }

    if (nel >= 2) {
      npass++;
      // cout << "Event passes: " << nel << " di-electrons" << endl; 

      /*  
      Float_t e = uber.electron.e[elind[0]] + uber.electron.e[elind[1]];
      Float_t px = uber.electron.px[elind[0]] + uber.electron.px[elind[1]];
      Float_t py = uber.electron.py[elind[0]]+ uber.electron.py[elind[1]];
      Float_t pz = uber.electron.pz[elind[0]]+ uber.electron.pz[elind[1]];
      Float_t invmass = sqrt(e*e-px*px-py*py-pz*pz);
      */

      // We use the Etcorr for inv. mass calculation.

      float e1_theta = 2.*atan2(exp(-(uber.electron.eveta[elind[0]]) ),1);
      float e1_phi = uber.electron.phi[elind[0]];
      float e1_et = uber.electron.etcorr[elind[0]];

      float e2_theta = 2.*atan2(exp(-(uber.electron.eveta[elind[1]]) ),1);
      float e2_phi = uber.electron.phi[elind[1]];
      float e2_et = uber.electron.etcorr[elind[1]];

      float massll = invmass(e1_et,e1_theta,e1_phi,e2_et,e2_theta,e2_phi);

      zee.Fill(massll);

	if (uber.header.run <= 156487)
	  zee_winter.Fill(massll);
	if (uber.header.run > 156487)
	  zee_summer.Fill(massll);

      ZEta.Fill(uber.electron.eveta[elind[0]]);
      ZEta.Fill(uber.electron.eveta[elind[1]]);
      runnumber->Fill(uber.header.run);


// To plot the number of positive and negative electrons versus Eta

	 if(uber.electron.charge[elind[0]]==1)
	     PEta->Fill(uber.electron.eveta[elind[0]]);
	
	 if(uber.electron.charge[elind[1]]==1)
	     PEta->Fill(uber.electron.eveta[elind[1]]);
	 
	 if(uber.electron.charge[elind[0]]==-1)
	     NEta->Fill(uber.electron.eveta[elind[0]]);
	 
	 if(uber.electron.charge[elind[1]]==-1)
	     NEta->Fill(uber.electron.eveta[elind[1]]);



// To plot asymmetry versus Z mass
     for(int o=1; o<=7; o++)
	{
	if((massll>=Afb->GetBinLowEdge(o))&&(massll<Afb->GetBinLowEdge(o+1)))
	{
	   if((uber.electron.charge[elind[0]]==-1)&&(uber.electron.eveta[elind[0]]>uber.electron.eveta[elind[1]]))
	   F[o-1]++;
	   if((uber.electron.charge[elind[1]]==-1)&&(uber.electron.eveta[elind[0]]>uber.electron.eveta[elind[1]]))
	   B[o-1]++;
	   if((uber.electron.charge[elind[0]]==-1)&&(uber.electron.eveta[elind[0]]<uber.electron.eveta[elind[1]]))
	   B[o-1]++;
	   if((uber.electron.charge[elind[1]]==-1)&&(uber.electron.eveta[elind[0]]<uber.electron.eveta[elind[1]]))
	   F[o-1]++;
	}
	if ((F[o-1]+B[o-1]) != 0)
            A[o-1]=((F[o-1]-B[o-1])/(F[o-1]+B[o-1]));
	Afb->SetBinContent(o,A[o-1]);
        }


   }
  }
  cout << "Number of passing events: " << npass << endl;
   
   zee.SetXTitle("Z Mass (GeV/c^2)");
   zee.SetYTitle("Number of Events");
   zee_winter.SetXTitle("Z Mass (GeV/c^2)");
   zee_winter.SetYTitle("Number of Events");
   zee_summer.SetXTitle("Z Mass (GeV/c^2)");
   zee_summer.SetYTitle("Number of Events");
   Afb.SetXTitle("Z Mass (GeV/c^2)");
   Afb.SetYTitle("Asymmetry (AFB)");
   PEta.SetXTitle("Eta");
   PEta.SetYTitle("Number of Events");
   NEta.SetXTitle("Eta");
   NEta.SetYTitle("Number of Events");
   runnumber.SetXTitle("Run Number");
   runnumber.SetYTitle("Number of Events");

/*
   zee.Draw();
   Afb->Draw();
   gPad->SetTicks(2,2);
   zee->Draw();
   Afb->Draw();
   PEta->Draw();
   NEta->Draw();
   zee->SetLineWidth(2);
   Afb->SetLineWidth(2);
   PEta->SetLineWidth(2);
   NEta->SetLineWidth(2);
   runnumber->SetLineWidth(2);
   c1->Clear();
   c1->Divide(2,2);
   c1->cd(1);
   zee->Draw();
   c1->cd(2);
   Afb->Draw();
   c1->cd(3);
   PEta->SetLineColor(kRed);
   PEta->Draw();
   NEta->Draw("same");
   c1->cd(4);
   runnumber->Draw();
*/


// Counting the number of runs and events
  cout << "   The total number of runs: " << numruns << endl;
  cout << "   The total number of events: " << numevents << endl;

// finding the first and last run numbers
  cout << "   The first run number:  "<<r<<endl;
  cout << "   The last run number:  "<<r2<<endl;

// Forward backward muon asymmetry
  cout << "   The number of forward events for Z's with a mass >60 and <80: " << F[0] << endl;
  cout << "   The number of backward events for Z's with a mass >60 and <80: " << B[0] << endl;
  cout << "   The number of forward events for Z's with a mass >80 and <100: " << F[1] << endl;
  cout << "   The number of backward events for Z's with a mass >80 and <100: " << B[1] << endl;
  cout << "   The number of forward events for Z's with a mass >100 and <120: " << F[2] << endl;
  cout << "   The number of backward events for Z's with a mass >100 and <120: " << B[2] << endl;

  cout << "   The asymmetry for Z's with a mass >60 and <80: " << A[0] << endl;
  cout << "   The asymmetry for Z's with a mass >80 and <100: " << A[1] << endl;
  cout << "   The asymmetry for Z's with a mass >100 and <120: " << A[2] << endl;

/*
  zee_winter->Fit("gaus");
  zee_summer->Fit("gaus");
  float Zmean = zee_winter->GetFunction("gaus")->GetParameter(1);
  float Zrms = zee_winter->GetFunction("gaus")->GetParameter(2);
  float Zmean2 = zee_summer->GetFunction("gaus")->GetParameter(1);
  float Zrms2 = zee_summer->GetFunction("gaus")->GetParameter(2);
  cout << " winter mean from fit " << Zmean
       << ".  winter RMS from fit " << Zrms << endl;
  cout << " summer mean from fit " << Zmean2
       << ".  summer RMS from fit " << Zrms2 << endl;
*/
  Zfile ->Write(); //saving the compiled program
}




//_____________________________________________________________________________

bool StandardElectronQ(myelectron_t& myelectron, int i, mytrack_t& mytrack)
{

  bool ans = false;

  //with fidele==1, you don't need to put electron.detector==0 (CEM)
  bool fiducial = (myelectron.fidele[i]==1);  // |Xces|<21cm, 9<|Zces|<230cm, tower 9 is excluded.

  bool goodTrack = GoodTrackQ(mytrack, myelectron.trkindex[i]);

  bool highpt = ( (myelectron.etcorr[i] > 20.0) &&
                  (myelectron.pttrk[i] > 10.0) );

  if(!highpt)   return(ans);
 
  bool centralEleId = 
    (
     (myelectron.charge[i]*myelectron.delx[i] >-3.0) &&
     (myelectron.charge[i]*myelectron.delx[i] < 1.5) &&
     (fabs(myelectron.delz[i]) < 3.0)
     );

  bool centralEmId =
      (
       ( myelectron.hadem[i]<(0.055+0.00045*myelectron.e[i]) ) &&
       ( myelectron.lshr[i] < 0.2) &&
       ( myelectron.iso4[i]/myelectron.etcorr[i] < 0.10 ) &&
       ( myelectron.chi2strip[i] < 10.0) &&
       ((myelectron.ep[i] < 2.0)||(myelectron.etcorr[i] > 50.0))
       );

        ans = fiducial && goodTrack && highpt && centralEleId && centralEmId ;
  //        ans = fiducial && goodTrack && highpt && centralEmId ;
  
  return(ans);
}
//_____________________________________________________________________________
bool GoodTrackQ(mytrack_t& mytrack, int i)
{

  if(i<0) return(false);

  bool ans = false;

  Int_t num_ax=0, num_axhits=0;
  Int_t num_st=0, num_sthits=0;

  //old method

  /*
    ans =

    (mytrack.naxialhits[i] > 24 && 
     mytrack.nstereohits[i] > 24 && 
     mytrack.naxialsl[i] >= 3 && 
     mytrack.nstereosl[i] >= 3 );

  */
  for (int j = 0; j < 8; j++) {
    Int_t num_hit_sl = (mytrack.slhitcount[i]>>j*4)&0xf;

//    cout << "check " << num_hit_sl << endl;

    if (fmod(j*1.0,2)>0) {
      num_axhits +=num_hit_sl; 
    } else {
      num_sthits +=num_hit_sl; 
    }

    if(num_hit_sl<7) continue;

    if (fmod(j*1.0,2)>0) {
      ++num_ax;}
    else {
      ++num_st;
    }
  }  

  /*
  if( (mytrack.naxialhits[i]!=num_axhits) ||
      (mytrack.nstereohits[i]!=num_sthits) ) {

    cout << "Problem in couting COT hit  " 
	 << setw(5) << mytrack.naxialhits[i]
	 << setw(5) << num_axhits
	 << setw(5) << mytrack.nstereohits[i]
	 << setw(5) << num_sthits << endl;

  }
  */

  ans = (num_ax>=3&&num_st>=3);
  ans &= (fabs(mytrack.z0[i])<60.);

  return(ans);

}
//_____________________________________________________________________________

Float_t invmass(Float_t et1, Float_t theta1, Float_t phi1, Float_t et2, Float_t theta2, Float_t phi2)
{

  //  if(theta1==0.0||theta2==0.0) {
  //    cout << "Warning in theta value " << endl;
  //  }

  Float_t E_1 = et1/sin(theta1);
  Float_t Ex_1 = E_1*sin(theta1)*cos(phi1);
  Float_t Ey_1 = E_1*sin(theta1)*sin(phi1);
  Float_t Ez_1 = E_1*cos(theta1);

  Float_t E_2 = et2/sin(theta2);
  Float_t Ex_2 = E_2*sin(theta2)*cos(phi2);
  Float_t Ey_2 = E_2*sin(theta2)*sin(phi2);
  Float_t Ez_2 = E_2*cos(theta2);
  
  Float_t mass  = sqrt(pow((E_1+E_2),2) - pow((Ex_1+Ex_2),2) 
		     - pow((Ey_1+Ey_2),2) - pow((Ez_1+Ez_2),2));
  
  return mass;

}