//////////////////////////////////////////////////////////////////////////
//                                                                      //
// ATLFast ElectronMaker class.                                         //
//                                                                      //
//////////////////////////////////////////////////////////////////////////

#include <TMCParticle.h>
#include <TRandom.h>

#include "ATLFast.h"
#include "ATLFMCMaker.h"
#include "ATLFClusterMaker.h"
#include "ATLFElectronMaker.h"
#include "ATLFCluster.h"
#include "ATLFElectron.h"

const Float_t kPI   = TMath::Pi();
const Float_t k2PI  = 2*kPI;
const Int_t kMAXELE = 100;

ClassImp(ATLFElectronMaker)

//_____________________________________________________________________________
 ATLFElectronMaker::ATLFElectronMaker()
{
   m_Nelectrons = 0;
}

//_____________________________________________________________________________
 ATLFElectronMaker::ATLFElectronMaker(const char *name, const char *title)
                 :ATLFMaker(name,title)
{
   SetMinPT();
   SetMaxEta();
   SetRconeMatch();
   SetRconeSep();
   SetRconeDep();
   SetMaxEdep();

   m_Fruits     = new TClonesArray("ATLFElectron",2,kFALSE);
   m_BranchName = "Electrons";
   m_Nelectrons = 0;
   Save();
}

//_____________________________________________________________________________
 ATLFElectronMaker::~ATLFElectronMaker()
{
   //dummy
}

//_____________________________________________________________________________
 void  ATLFElectronMaker::AddElectron(Int_t code, Int_t mcparticle, Int_t mother, Float_t eta, Float_t phi, Float_t pt)
{
//            Add a new electron  to the list of electrons

 //Note the use of the "new with placement" to create a new electron object.
 //This complex "new" works in the following way:
 //   electrons[i] is the value of the pointer for electron number i in the TClonesArray
 //   if it is zero, then a new electron must be generated. This typically
 //   will happen only at the first events
 //   If it is not zero, then the already existing object is overwritten
 //   by the new electron parameters.
 // This technique should save a huge amount of time otherwise spent
 // in the operators new and delete.

   TClonesArray &electrons = *(TClonesArray*)m_Fruits;
   new(electrons[m_Nelectrons++]) ATLFElectron(code,mcparticle,mother,eta,phi,pt);
}

//_____________________________________________________________________________
 void ATLFElectronMaker::Clear(Option_t *option)
{
//    Reset Electron Maker

   m_Nelectrons = 0;
   ATLFMaker::Clear(option);
}


//_____________________________________________________________________________
 void ATLFElectronMaker::Init()
{
//  Create histograms
   m_Mult          = new TH1F("EleMult","//electron multiplicity",10,0,10);
   m_MultHard      = new TH1F("EleMultHard","electron multiplicity hard",10,0,10);
   m_MultHardIsol  = new TH1F("EleMultHardIsol","electron multiplicity hard + isol",10,0,10);
   m_PT            = new TH1F("ElePT","(pT- pTcru) vers E electron ISOLATED",50,0,100);
   m_Eta           = new TH1F("EleEta","(eta-etacru) vers E electron ISOLATED",50,0,100);
   m_Phi           = new TH1F("ElePhi","(phi-phicru) vers E electron ISOLATED",50,0,100);
   m_Counter       = new TH1F("EleCounter","counter vers E electron ISOLATED",50,0,100);
   m_Mass2e        = new TH1F("EleMass2e","ee mass",50,80,100);
   m_Mass2esubst   = new TH1F("EleMass2esubst","ee mass subtract width",50,80,100);
   m_Mass4e        = new TH1F("EleMass4e","eeee mass",50,120,140);
}

//_____________________________________________________________________________
 void ATLFElectronMaker::Finish()
{
// Function called by ATLFast::Finish at the end of the job
//   Normalize histograms

   m_PT->Divide(m_Counter);
   m_Eta->Divide(m_Counter);
   m_Phi->Divide(m_Counter);

}

//_____________________________________________________________________________
 Int_t ATLFElectronMaker::Make()
{
//.............................................
//  This function searches for isolated electrons, by scanning through
//  the list of MC particles. If an electron is found, its energy is
//  smeared using function RESELE. Different energy-smearing for low luminosity
//  and high luminosity can be invoked.
//  Isolated electrons are stored in The TClonesArray of electrons
//  and the energy clusters associated with them are removed.
//  Non-isolated electrons are not considered further.
//.............................................

   Int_t elCode[kMAXELE],    elPart[kMAXELE],  elMother[kMAXELE];
   Float_t elEta[kMAXELE],   elPhi[kMAXELE],   elPT[kMAXELE];

   Int_t sortindex[kMAXELE];
   Int_t nele   = 0;
   Int_t i, k, ic, ij, KF;

//........................................................
//... Get pointers to MCparticles arrays and ClonesArray
//........................................................
   ATLFMCMaker *mcarlo = gATLFast->MCMaker();
   TClonesArray *particles = mcarlo->Fruits();
   TMCParticle *part, *parent;
   Float_t *pPT  = mcarlo->PT();
   Float_t *pEta = mcarlo->Eta();
   Float_t *pPhi = mcarlo->Phi();
   Int_t   *Index= mcarlo->Index();
   

//........................................................
//... Get pointers to clustermaker arrays and ClonesArray
//........................................................
   ATLFClusterMaker *clustermaker = gATLFast->ClusterMaker();
   TClonesArray *clusters = clustermaker->Fruits();
   ATLFCluster *cluster;
   Float_t *cellEta   = clustermaker->CellEta();
   Float_t *cellPhi   = clustermaker->CellPhi();
   Float_t *cellSumPT = clustermaker->CellSumPT();
   Int_t nclusters    = clustermaker->Nclusters();
   Int_t ncells       = clustermaker->Ncells();

//........................................................
//....look for isolated electrons, sort clusters common
//........................................................
   Bool_t isolated;
   Float_t ptcru, etacru, phicru;
   Float_t sigma, pxele, pyele, pzele, eeele;
   Float_t pt, eta, phi, atheta, abseta, deta, dphi;
   Float_t adphi, dr, ddr, edep;
   Float_t theta = 0;
   Float_t ene;
   Int_t lclu;
   Int_t nstop = mcarlo->Nstop();
   Float_t RconeSep2   = m_RconeSep*m_RconeSep;
   Float_t RconeDep2   = m_RconeDep*m_RconeDep;
   Float_t RconeMatch2 = m_RconeMatch*m_RconeMatch;

   for (Int_t ind=0;ind<mcarlo->Nstables();ind++) {
      i = Index[ind];
      if (i < mcarlo->Nstart()) continue;
      part = (TMCParticle*)particles->UncheckedAt(i);
      KF   = part->GetKF();
      if (TMath::Abs(KF) != 11) continue;
                 //-- analyse  electrons
      isolated = kTRUE;
      lclu   = -1;
      pt     = pPT[i];
      eta    = pEta[i];
      phi    = pPhi[i];
      ptcru  = pt;
      etacru = eta;
      phicru = phi;
                 //.... apply smearing
      if (gATLFast->Smearing() == 1) {
          ene    = part->GetEnergy();
          if (ene <= 0) continue;
          sigma = ElectronResolution(gATLFast->Luminosity(),ene, eta, pt); 
          pxele = part->GetPx()*(1+sigma);
          pyele = part->GetPy()*(1+sigma);
          pzele = part->GetPz()*(1+sigma);
          eeele = part->GetEnergy()*(1+sigma);
          if (TMath::Abs(pzele/eeele) < 1) {
             theta = TMath::ACos(pzele/eeele);
          } else {
             if (pzele > 0) theta = 0;
             if (pzele < 0) theta = kPI;
          }
          atheta = TMath::Abs(TMath::Tan(0.5*theta));
          if (atheta < 0.0001) atheta = 0.0001;
          eta = -TMath::Log(atheta);
          phi = mcarlo->Angle(pxele,pyele);
          pt  = TMath::Sqrt(pxele*pxele + pyele*pyele);
      }
      if (pt < m_MinPT) continue;
      abseta = TMath::Abs(eta);
      if (abseta > m_MaxEta) continue;

           //........................................................
           //....mark electron cluster
           //........................................................
      dr = 100;
      for (k=0;k<nclusters;k++) {
         cluster = (ATLFCluster*)clusters->UncheckedAt(k);
         deta    = eta - cluster->Eta();
         dphi    = phi - cluster->Phi();
         adphi   = TMath::Abs(dphi) - k2PI;
         if (TMath::Abs(dphi) > kPI) ddr = deta*deta +adphi*adphi;
         else                        ddr = deta*deta + dphi*dphi;
         if (ddr < dr) {lclu = k; dr = ddr;}
      }
      if (dr > RconeMatch2) lclu = -1;
           //........................................................
           //....check if electron isolated from clusters
           //........................................................
      for (k=0;k<nclusters;k++) {
         if (k == lclu) {
            ddr = 100;
         } else {
            cluster = (ATLFCluster*)clusters->UncheckedAt(k);
            deta    = eta - cluster->Eta();
            dphi    = phi - cluster->Phi();
            adphi   = TMath::Abs(dphi) - k2PI;
            if (TMath::Abs(dphi) > kPI) ddr = deta*deta +adphi*adphi;
            else                        ddr = deta*deta + dphi*dphi;
         }
         if (ddr < RconeSep2) isolated = kFALSE;
      }

           //........................................................
           //....check on energy deposition of cells EDEP in cone m_MaxEdep
           //........................................................
      edep = 0;
      for (ic=0;ic<ncells;ic++) {
         deta  = eta - cellEta[ic];
         dphi  = phi - cellPhi[ic];
         adphi = TMath::Abs(dphi) - k2PI;
         if (TMath::Abs(dphi) > kPI)  ddr = deta*deta + adphi*adphi;
         else                         ddr = deta*deta + dphi*dphi;
         if (ddr < RconeDep2) edep += cellSumPT[ic];
      }
      if (edep-pt > m_MaxEdep) isolated = kFALSE;

      parent = (TMCParticle*)particles->UncheckedAt(part->GetParent()-1);
      if (isolated) { // fill isolated electron structure

         if (lclu >= 0) {nclusters--; clustermaker->RemoveCluster(lclu);}

         elCode[nele]   = KF;
         elPart[nele]   = i; 
         elMother[nele] = parent->GetKF();
         elEta[nele]    = eta;
         elPhi[nele]    = phi;
         elPT[nele]     = pt;
         nele++;
           //........................................................
           //.....monitor energy smearing for isolated electrons
           //........................................................
         m_PT->Fill(part->GetEnergy(), TMath::Abs(pt-ptcru));
         m_Eta->Fill(part->GetEnergy(),TMath::Abs(eta-etacru));
         m_Phi->Fill(part->GetEnergy(),TMath::Abs(phi-phicru));
         m_Counter->Fill(part->GetEnergy());
     }
 }  // end of loop on primary particles


    m_Mult->Fill(nele+0.01);

//........................................................
//.....arrange isolated electrons in falling E_T sequence
//........................................................

   Int_t eleCode[kMAXELE],    elePart[kMAXELE],  eleMother[kMAXELE];
   Float_t eleEta[kMAXELE],   elePhi[kMAXELE],   elePT[kMAXELE];

   gATLFast->SortDown(nele, elPT, sortindex);

   for (i=0;i<nele;i++) {
      ij = sortindex[i];
      eleCode[i]   = elCode[ij];
      elePart[i]   = elPart[ij];
      eleMother[i] = elMother[ij];
      eleEta[i]    = elEta[ij] ;
      elePhi[i]    = elPhi[ij];
      elePT[i]     = elPT[ij];
   }

   Float_t pxele1, pyele1, pzele1, eeele1;
   Float_t pxele2, pyele2, pzele2, eeele2;
   Float_t px4el,  py4el,  pz4el,  ee4el;
   Float_t aelel, ammc, ammcc;

//........................................................
//......check mass resolution
//........................................................
   if (nele == 2) {
      pxele1 = elePT[0]*TMath::Cos(elePhi[0]);
      pyele1 = elePT[0]*TMath::Sin(elePhi[0]);
      pzele1 = elePT[0]*TMath::SinH(eleEta[0]);
      eeele1 = elePT[0]*TMath::CosH(eleEta[0]);
      pxele2 = elePT[1]*TMath::Cos(elePhi[1]);
      pyele2 = elePT[1]*TMath::Sin(elePhi[1]);
      pzele2 = elePT[1]*TMath::SinH(eleEta[1]);
      eeele2 = elePT[1]*TMath::CosH(eleEta[1]);
      aelel  = (eeele1+eeele2)*(eeele1+eeele2) - (pxele1+pxele2)*(pxele1+pxele2)
              -(pyele1+pyele2)*(pyele1+pyele2) - (pzele1+pzele2)*(pzele1+pzele2);
      if (aelel > 0) aelel = TMath::Sqrt(aelel);
      m_Mass2e->Fill(aelel);
           //.....substract natural width
      for (k=0;k<nstop;k++) {
         part   = (TMCParticle*)particles->UncheckedAt(k);
         parent = (TMCParticle*)particles->UncheckedAt(part->GetParent()-1);
         KF     = part->GetKF();
         if (KF == 11 && parent->GetKF() == 23) {
            ammc  = 91.173;
            ammcc = parent->GetMass();
            m_Mass2esubst->Fill(aelel-(ammcc-ammc));
         }
      }
   }
//........................................................
//........mass of 4electrons
//........................................................
   if (nele == 4) {
      px4el = py4el = pz4el = ee4el = 0;
      for (k=0;k<4;k++) {
          px4el += elePT[k]*TMath::Cos(elePhi[k]);
          py4el += elePT[k]*TMath::Sin(elePhi[k]);
          pz4el += elePT[k]*TMath::SinH(eleEta[k]);
          ee4el += elePT[k]*TMath::CosH(eleEta[k]);
      }
      aelel = ee4el*ee4el -px4el*px4el -py4el*py4el - pz4el*pz4el;
      if (aelel > 0) aelel = TMath::Sqrt(aelel);
      m_Mass4e->Fill(aelel);
   }

//........................................................
//....cross-check with partons
//........................................................
   Int_t iele    = 0;
   Int_t ieleiso = 0;
   Float_t ener;
   Int_t KF2;
   TMCParticle *part2;
   for (i=0;i<nstop;i++) {
      part = (TMCParticle*)particles->UncheckedAt(i);
      KF   = part->GetKF();
      if (TMath::Abs(KF) == 11) {
         abseta   = TMath::Abs(pEta[i]);
         ener     = 0;
         isolated = kTRUE;
         for (k=0;k<nstop;k++) {
            part2  = (TMCParticle*)particles->UncheckedAt(k);
            KF2    = TMath::Abs(part2->GetKF());
            if ((KF2 <= 21) && (i != k) && (KF2 != 12) && (KF2 != 14) && (KF2 !=16)) {
               deta  = pEta[i] - pEta[k];
               dphi  = pPhi[i] - pPhi[k];
               adphi = TMath::Abs(dphi) - k2PI;
               if (TMath::Abs(dphi) > kPI) ddr = deta*deta + adphi*adphi;
               else                        ddr = deta*deta + dphi*dphi;
               if (ddr < RconeSep2 && pPT[k] > clustermaker->MinETCalo()) isolated = kFALSE;
               if (ddr < RconeDep2 && pPT[k] < clustermaker->MinETCalo()) ener += pPT[k];
            }
         }
         if (ener > m_MaxEdep) isolated = kFALSE;
         if (abseta < m_MaxEta && pPT[i] > m_MinPT) {
            iele++;
            if (isolated)  ieleiso++;
         }
      }
   }
   m_MultHard->Fill(iele+0.01);
   m_MultHardIsol->Fill(ieleiso+0.01);

//........................................................
//....Store electrons in Root ClonesArray
//........................................................
   m_Nelectrons = 0;
   for (k=0;k<nele;k++) {
      AddElectron(eleCode[k],
                  elePart[k],
                  eleMother[k],
                  eleEta[k],
                  elePhi[k],
                  elePT[k]);
   }
   return 0;
}

//_____________________________________________________________________________
 void ATLFElectronMaker::PrintInfo()
{
   ATLFMaker::PrintInfo();
}

//_____________________________________________________________________________
 Float_t ATLFElectronMaker::ElectronResolution(Int_t lumi, Float_t ene, Float_t eta, Float_t pt)
{
//     parametrizes electron  resolution for low and high luminosity
//     parametrization from L. Poggioli

  Float_t rpilup = 0;
  Float_t aa, bb, sigph1, sigph, sigpu;
  Float_t sqrtene = TMath::Sqrt(ene);
  Float_t abseta  = TMath::Abs(eta);

  while(1) {
     gRandom->Rannor(aa,bb);
     sigph1 = aa*0.12/sqrtene;
     if (1 + sigph1 <= 0) continue;
     sigph = sigph1;
     if (abseta < 1.4) {
        while(1) {
           gRandom->Rannor(aa,bb);
           sigph1 = aa*0.245/pt + bb*0.007;
           if (1 + sigph1 > 0) break;
        }
     } else {
        while(1) {
           gRandom->Rannor(aa,bb);
           sigph1 = aa* 0.306*((2.4-abseta)+0.228)/ene + bb*0.007;
           if (1 + sigph1 > 0) break;
        }
     }
     sigph += sigph1;
     if (lumi == 2) {
        if (abseta <  0.6)                rpilup = 0.32;
        if (abseta > 0.6 && abseta < 1.4) rpilup = 0.295;
        if (abseta > 1.4)                 rpilup = 0.27;
        while(1) {
           gRandom->Rannor(aa,bb);
           sigpu = aa*rpilup/pt;
           if (1 + sigpu > 0) break;
        }
        sigph += sigpu;
     }
     if (1 + sigph > 0) return sigph;
  }
}