//==================================================================
//
//	TConvEleCand.cc
//
//	Version 0.0:      P. Koehn    5/5/01
//	Version 0.9:      R. Hughes    7/23/01
//
//==================================================================
#if !defined (__CINT__) || defined (__MAKECINT__)
#include "TF1.h"
#include <iostream>
#include <iomanip>
#include <algorithm>

#include "Stntuple/loop/TStnAna.hh"
#include "Stntuple/obj/TStnHeaderBlock.hh"
#endif
#include "TConvEleCand.hh"

//_____________________________________________________________________________
TConvEleCand::TConvEleCand(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();
}

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

void TConvEleCand::init() {
//
  NElectronMax = 20;
  totEvents = 0;
  cnvEvents = 0;
  SetMaxDeltaLam  (0.03);
  SetMaxSeparation(0.3);
//
// Default is to make no histograms
  enableEleHist(false);
}

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

//_____________________________________________________________________________
int TConvEleCand::BeginJob() {

  printf("----- Begin job: ---- %s\n",GetName());
  totEvents = 0;
  cnvEvents = 0;

// register the Electron data block
  pEleBlock = (TStnElectronBlock*) fAna->RegisterDataBlock("ElectronBlock","TStnElectronBlock");
  if (! pEleBlock) {
    printf(" >>> branch *** %s *** doesn't exist \n","EleBlock");
    fEnabled = 0;
  }

// register the Electron data block
  pTrkBlock = (TStnTrackBlock*) fAna->RegisterDataBlock("TrackBlock","TStnTrackBlock");
  if (! pTrkBlock) {
    printf(" >>> branch *** %s *** doesn't exist \n","TrkBlock");
    fEnabled = 0;
  }
//
// If requested, book the histograms
  if (isEnableHist()) BookHistograms();

  return 0;
}

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

//_____________________________________________________________________________
void TConvEleCand::BookHistograms() {
//
// Delete the old histos
//   GetHistogramList()->Delete();
   histFile->cd(histDir);    // make the "dir" directory the current directory
//
// Now define histos
   EleWSEtHist = new TH1F("EleWSEtHist","EleWSCand Et",50,0,100);
   EleWSEoverPHist = new TH1F("EleWSEoverPHist","EleWSCand E over P",50,0,5);
   EleWSDelxHist = new TH1F("EleWSDelxHist","EleWSCand DeltaX",50,-20,20);
   EleWSDelzHist = new TH1F("EleWSDelzHist","EleWSCand DeltaZ",50,-20,20);
   EleWSLshrHist = new TH1F("EleWSLshrHist","EleWSCand Lshr",50,0.0,5.0);
   EleWSHademHist = new TH1F("EleWSHademHist","EleWSCand HadEm",50,0,1.0);
   EleWSIsolHist = new TH1F("EleWSIsolHist","EleWSCand Iso",50,0,2);
   EleWSTisolHist = new TH1F("EleWSTisolHist","EleWSCand Track Iso",50,0,100);
   EleWSChisHist = new TH1F("EleWSChisHist","EleWSCand ChiStrip",50,0,100);
   EleWSDelZvrtHist = new TH1F("EleWSDelZvrtHist","EleWSCand Delta(Zvtx-Zele)",50,-10,10);
   EleWSZvrtHist = new TH1F("EleWSZvrtHist","EleWSCand Event Zvertex",50,-100,100);
   EleWSDetHist = new TH1F("EleWSDetHist","EleWSCand Detector",3,0,2);

   EleRSEtHist = new TH1F("EleRSEtHist","EleRSCand Et",50,0,100);
   EleRSEoverPHist = new TH1F("EleRSEoverPHist","EleRSCand E over P",50,0,5);
   EleRSDelxHist = new TH1F("EleRSDelxHist","EleRSCand DeltaX",50,-20,20);
   EleRSDelzHist = new TH1F("EleRSDelzHist","EleRSCand DeltaZ",50,-20,20);
   EleRSLshrHist = new TH1F("EleRSLshrHist","EleRSCand Lshr",50,0.0,5.0);
   EleRSHademHist = new TH1F("EleRSHademHist","EleRSCand HadEm",50,0,1.0);
   EleRSIsolHist = new TH1F("EleRSIsolHist","EleRSCand Iso",50,0,2);
   EleRSTisolHist = new TH1F("EleRSTisolHist","EleRSCand Track Iso",50,0,100);
   EleRSChisHist = new TH1F("EleRSChisHist","EleRSCand ChiStrip",50,0,100);
   EleRSDelZvrtHist = new TH1F("EleRSDelZvrtHist","EleRSCand Delta(Zvtx-Zele)",50,-10,10);
   EleRSZvrtHist = new TH1F("EleRSZvrtHist","EleRSCand Event Zvertex",50,-100,100);
   EleRSDetHist = new TH1F("EleRSDetHist","EleRSCand Detector",3,0,2);

   fHist.fDeltaLamRS = new TH1F("fDeltaLamRS","delta(lambda), RS",250,-2.5,2.5);
   fHist.fDeltaLamWS = new TH1F("fDeltaLamWS","delta(lambda), WS",250,-2.5,2.5);
   fHist.fSeparationRS = new TH1F("fSeparationRS","delta(lambda), RS",250,-25,25);
   fHist.fSeparationWS = new TH1F("fSeparationWS","delta(lambda), WS",250,-25,25);
   fHist.fRConvRS = new TH1F("fConvRS","delta(lambda), RS",500,0,100);
   fHist.fRConvWS = new TH1F("fConvWS","delta(lambda), WS",500,0,100);

   fHist.fmassRS = new TH1F("fmassRS","right sign pair mass",50,0.0,1.0);
   fHist.fmassWS = new TH1F("fmassWS","wrong sign pair mass",50,0.0,1.0);

   fHist.fEpVsTrackPt = new TH2F("fEpVsTrackPt", "EoverP vs Track Pt, RS",
   				50,0.0,10.0,40,0.0,4.0);

   fHist.fYConvVsXConvRS = new TH2F("fYConvVsXConvRS", "Conv x vs y, RS",100,-50,50,100,-50,50);
   fHist.fYConvVsXConvWS = new TH2F("fYConvVsXConvWS", "Conv x vs y, WS",100,-50,50,100,-50,50);



}

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

  double  r1, r2, x1, x2, y1, y2, ff, dd, sep, dt, conv_sign;
  double  xconv, yconv, rconv, dx, dy, sqr;
  int     ncand, ntrk;


  if (fDebugLevel > 0) {
        printf(" ---- ientry = %7i\n",ientry);
  }
//
// How many events have we processed
  totEvents++;
  ncand = 0;
  bool foundConv = false;
//
// Get the header info
  TStnHeaderBlock* header = GetHeaderBlock();
  int iev = header->EventNumber();
  int irun = header->RunNumber();
//
// Electron Selection
//===================
  pEleBlock->GetEntry(ientry);
  pTrkBlock->GetEntry(ientry);

  int nelectrons = pEleBlock->NElectrons();
  NconvElectron = 0;
  ntrk = pTrkBlock->NTracks();
  
//  printf("Ntracks %d \n",ntrk);
  for (int i=0; i<nelectrons; i++) {
		TStnElectron* electron = pEleBlock->Electron(i);

		if (electron->EoverP() < 4.0) {
//
// Get the electron track pointer
			int trkele_index = electron->TrackNumber();
//
//    printf("Electron Et %f trkind %d pt %f phi %f ztr %f\n",
//    	electron->Et(),electron->fTrind,electron->fPt,
//	electron->fPhi,electron->fZtrk);
// Loop over the other tracks
         trkele_index = -1;
         for(int j=1; j<ntrk; j++) {
            TStnTrack& t = *pTrkBlock->Track(j);
            if (t.fPt == electron->fPt) trkele_index = j;
//            if (fabs(t.fPt) > 10.0) {
//   	         printf("track pt %f phi %f charge %f z %f \n",
//		            t.fPt,t.fPhi0,t.Charge(),t.fZ0);
//            }
         }
         if (trkele_index > ntrk) break;
         if (trkele_index < 0) break;
         TStnTrack& t1 = *pTrkBlock->Track(trkele_index);
//
	      bool rightSign = false;
   	   bool wrongSign = false;
//
// Make sure the electron track is reasonable
			if ((t1.NAxSeg() > 1) && (t1.NStSeg() > 1) && (t1.fPt > 10.0)) {

// Loop over the other tracks
		      for(int j=1; j<ntrk; j++) {
               if (j == trkele_index)  goto NEXT_T2;
               TStnTrack& t2 = *pTrkBlock->Track(j);

               if ((t2.NAxSeg() <= 1) || (t2.NStSeg() <= 1)
                      || (fabs(t2.fPt) > 10.0)) goto NEXT_T2;

               dt        = t1.Lam0()-t2.Lam0();
               conv_sign = t1.Charge()+t2.Charge();

               if (fabs(dt) < 0.25) {

                  r1 = 1.0/fabs(2*t1.C0());
                  ff = t1.Charge()*r1+t1.D0();
                  x1 = -ff*sin(t1.Phi0());
                  y1 =  ff*cos(t1.Phi0());

                  r2 = 1.0/fabs(2*t2.C0());
                  ff = t2.Charge()*r2+t2.D0();
                  x2 = -ff*sin(t2.Phi0());
                  y2 =  ff*cos(t2.Phi0());

                  dd  = sqrt((x1-x2)*(x1-x2) + (y1-y2)*(y1-y2));

                  sep = dd-r1-r2;

                  dx  = x2-x1;
                  dy  = y2-y1;
                  sqr = sqrt(dx*dx+dy*dy);

                  xconv = x1+dx/sqr*(r1+sep/2);
                  yconv = y1+dy/sqr*(r1+sep/2);

                  rconv = sqrt(xconv*xconv+yconv*yconv);

//-----------------------------------------------------------------------------
// take some "arbitrary" cut on delta(lambda) and plot distribution for 
// separation in XY
//-----------------------------------------------------------------------------
                  if (fabs(dt) < 0.1) {
         	         if (isEnableHist()) {
            		      if (conv_sign == 0) 
                  	      fHist.fSeparationRS->Fill(sep);
            		      else                
                  	      fHist.fSeparationWS->Fill(sep);
            	      }
                  }
//-----------------------------------------------------------------------------
// now for the track pairs with small dR(xy) plot distribution for
// delta(Lambda)
//-----------------------------------------------------------------------------
                  if (fabs(sep) < 0.5) {
                     if (isEnableHist()) {
               	      if (conv_sign == 0) 
                  	      fHist.fDeltaLamRS->Fill(dt);
               	      else                
                  	      fHist.fDeltaLamWS->Fill(dt);
                     }
                  }
//-----------------------------------------------------------------------------
// count number of candidate conversion pairs per event
//-----------------------------------------------------------------------------
                  if ( (fabs(sep) < fMaxSeparation) && (fabs(dt) < fMaxDeltaLam)) {
            	      ncand += 1;
//
// For these candidates, calculate the mass at the converion radius
                     float mass = massTracks(rconv,
                                 t1.fPt*t1.Charge(),t1.fPhi0,t1.fCot,t1.fD0,t1.fZ0,
                                 t2.fPt*t2.Charge(),t2.fPhi0,t2.fCot,t2.fD0,t2.fZ0);
                     if (isEnableHist()) {
                        if (conv_sign == 0) {
                 		      rightSign = true;
                  	      fHist.fmassRS->Fill(mass);
                  	      fHist.fRConvRS->Fill(rconv);
                  	      fHist.fYConvVsXConvRS->Fill(xconv,yconv);
                  	      if (rconv > 35.0) 
                           	fHist.fEpVsTrackPt->Fill(t2.fPt,electron->EoverP());
                        } else {
                  	      wrongSign = true;
                  	      fHist.fmassWS->Fill(mass);
                  	      fHist.fRConvWS->Fill(rconv);
                  	      fHist.fYConvVsXConvWS->Fill(xconv,yconv);
                        }
                     }
                  }
               }
               NEXT_T2:;
            }  // loop over other tracks in the events
//
// Did we find a conversion electron?
            if (ncand > 0) {
               NconvElectronList[i] = true;
               foundConv = true;
        	      if (NconvElectron < NElectronMax) {
                  convElectron[NconvElectron] = electron;
                  NconvElectron++;
//
// If histograms are enabled, then plot for the good candidates
	               if (rightSign && !wrongSign) {
   	                  EleRSEtHist->Fill(electron->Et());
	                  EleRSEoverPHist->Fill(electron->EoverP());
	                  EleRSHademHist->Fill(electron->HadEm());
	                  EleRSIsolHist->Fill(electron->Iso()/electron->Et());
	                  EleRSTisolHist->Fill(electron->Tiso());
	                  EleRSDetHist->Fill(electron->DetectorCode());
	                  EleRSDelxHist->Fill(electron->DelX());
	                  EleRSDelzHist->Fill(electron->DelZ());
	                  EleRSLshrHist->Fill(electron->Lshr());
	                  EleRSChisHist->Fill(electron->Chi2Strip());
	                  EleRSDelZvrtHist->Fill(electron->fZtrk - electron->fZv);
	                  EleRSZvrtHist->Fill(electron->fZv);
	               } else if (!rightSign && wrongSign) {
   	                  EleWSEtHist->Fill(electron->Et());
	                  EleWSEoverPHist->Fill(electron->EoverP());
	                  EleWSHademHist->Fill(electron->HadEm());
	                  EleWSIsolHist->Fill(electron->Iso()/electron->Et());
	                  EleWSTisolHist->Fill(electron->Tiso());
	                  EleWSDetHist->Fill(electron->DetectorCode());
	                  EleWSDelxHist->Fill(electron->DelX());
	                  EleWSDelzHist->Fill(electron->DelZ());
	                  EleWSLshrHist->Fill(electron->Lshr());
	                  EleWSChisHist->Fill(electron->Chi2Strip());
	                  EleWSDelZvrtHist->Fill(electron->fZtrk - electron->fZv);
	                  EleWSZvrtHist->Fill(electron->fZv);
	               }
	      
//
               }
            } else {
              NconvElectronList[i] = false;       
            }
         }  // cuts on electron track
//
		}	// cuts on electrons
   }  // end loop over electrons
//
// If at least one good electron was found....
   if (foundConv) {
     cnvEvents++;
     if (fDebugLevel > 0) {
        printf(" Conversion electron run %d event %d\n",irun,iev);
     }
   }

  return 0;		       
}

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

  printf("----- end job: ---- %s\n",GetName());
  printf("  Events processed %d with good conv electrons %d \n",
  	totEvents,cnvEvents);
  return 0;
}


/**
* Calculates the x,y,z position of a charged track passing through the
* COT, at a given radius from the origin.
* @param radius is the radius from x=0,y=0 to determine the x,y,z
* @param curvature is the half curvature for the charged track
* @param PhiZeroIn is the phi at the origin for the charged track
* @param cotan is the  cot(theta) for the charged track
* @param D_O is the impact parameter for the charged track
* @param Z_O is the z position at the distance of closest approach
*            to x=0,y=0
* @return Returns xyz position in a 3-element array.
*/
void TConvEleCand::track_Prop(float radius, float curvature, 
			float PhiZeroIn, float cotan, float D_O, float Z_0, float* xyz) {
//
//		float xyz[] = new float[3];
//
// calculate some constants on entry
//		float epsilon = 1.0;        // epsilon is outgoing so is 1
//	printf("in track_prop rad %f d0 %f \n",radius,D_O);
	float epsilon = 1.0;
	float B = curvature* sqrt((radius*radius-D_O*D_O)/(1.0+2.0*curvature*D_O));
	float U_O = cos(PhiZeroIn);
	float V_O = sin(PhiZeroIn);
	float rho = 2.0*curvature;
	float XO = -D_O*V_O;
	float YO = D_O*U_O;			

	float cosPhiZero = cos(PhiZeroIn);
	float sinPhiZero = sin(PhiZeroIn);

	float Radius_of_Arc = fabs(1.0/(2*curvature)); 	
	float X_C = -(1/(2*curvature)+D_O)*sinPhiZero;  // X_C calculated here
	float Y_C =  (1/(2*curvature)+D_O)*cosPhiZero;  // Y_C calculated here
//			
	xyz[0] = XO + (U_O*2.0*epsilon*B*sqrt(1.0-B*B)-V_O*2.0*B*B)/rho;
	xyz[1] = YO + (V_O*2.0*epsilon*B*sqrt(1.0-B*B)+U_O*2.0*B*B)/rho;
	xyz[2] = Z_0 + ((cotan)*asin(B))/curvature;

//   direction vector (NOT Px,Py,Pz)

	float HALFPI = 0.5*3.14159;
   float A1 = 0.5/curvature;
   float A2 = 1.0;
   if (curvature<0.0) A2 = -1.0;
   float A3 = A2*(A1+D_O);
   float PSI= PhiZeroIn + A2*HALFPI;
   A1 = xyz[0] - A3*cos(PSI);
   A3 = xyz[1] - A3*sin(PSI);
   PSI= atan2(A3,A1) + A2*HALFPI;
   A1 = 1.0/sqrt(1.0+cotan*cotan);
   xyz[3] = cos(PSI)*A1;
   xyz[4] = sin(PSI)*A1;
   xyz[5] = cotan*A1;

}// end method track_Pop
   
float TConvEleCand::massTracks(float radius, 
						float pt1, float phi1, float cotan1, float d1, float z1, 
						float pt2, float phi2, float cotan2, float d2, float z2)
{
	float konst = 0.00014989*14.116;
	float c1 = konst / pt1;
	float c2 = konst / pt2;
   float xyz_tr1[6];
   float xyz_tr2[6];

   float px1 = fabs(pt1)*cos(phi1);
   float py1 = fabs(pt1)*sin(phi1);
   float pz1 = fabs(pt1)*cotan1;
//   printf("\ntr1 px %f py %f pz %f\n",px1,py1,pz1);
	track_Prop(radius, c1, phi1, cotan1, d1, z1, xyz_tr1);
//	printf("x %f y %f z %f ux %f uy %f uz %f \n",
//   		xyz_tr1[0],xyz_tr1[1],xyz_tr1[2],xyz_tr1[3],xyz_tr1[4],xyz_tr1[5]);

   float px2 = fabs(pt2)*cos(phi2);
   float py2 = fabs(pt2)*sin(phi2);
   float pz2 = fabs(pt2)*cotan2;
//   printf("tr2 px %f py %f pz %f\n",px2,py2,pz2);
	track_Prop(radius, c2, phi2, cotan2, d2, z2, xyz_tr2);
//	printf("x %f y %f z %f ux %f uy %f uz %f \n",
//   		xyz_tr2[0],xyz_tr2[1],xyz_tr2[2],xyz_tr2[3],xyz_tr2[4],xyz_tr2[5]);

   float ptot1 = fabs(pt1)*sqrt(1 + cotan1*cotan1);
   float ptot2 = fabs(pt2)*sqrt(1 + cotan2*cotan2);
   float px = ptot1*xyz_tr1[3] + ptot2*xyz_tr2[3];
   float py = ptot1*xyz_tr1[4] + ptot2*xyz_tr2[4];
   float pz = ptot1*xyz_tr1[5] + ptot2*xyz_tr2[5];
   float e = ptot1 + ptot2;
   float mass = e*e - px*px - py*py - pz*pz;
	float rad1 = sqrt(xyz_tr1[0]*xyz_tr1[0] + xyz_tr1[1]*xyz_tr1[1]);
	float rad2 = sqrt(xyz_tr2[0]*xyz_tr2[0] + xyz_tr2[1]*xyz_tr2[1]);
//	printf("radius %f rad1 %f rad2 %f pt1 %f ptot1 %f mass %f \n",
//   	radius,rad1,rad2,pt1,ptot1,mass);
   if (mass > 0.0) {
   	mass = sqrt(mass);
   } else {
   	mass = -1.0;
   }
   return mass;
}