//////////////////////////////////////////////////////////////////////////
//                                                                      //
// ATLFast ClusterMaker class.                                          //
//                                                                      //
//////////////////////////////////////////////////////////////////////////

#include <TRandom.h>

#include <TMCParticle.h>

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

const Float_t kPI  = TMath::Pi();
const Float_t k2PI = 2*kPI;
const Int_t kMAXCELLS = 2000;
const Int_t kMAXCLU   = 1000;

ClassImp(ATLFClusterMaker)

//_____________________________________________________________________________
 ATLFClusterMaker::ATLFClusterMaker()
{
   m_CellIndex = 0;
   m_CellCount = 0;
   m_CellFlag  = 0;
   m_CellEta   = 0;
   m_CellPhi   = 0;
   m_CellSumPT = 0;
}

//_____________________________________________________________________________
 ATLFClusterMaker::ATLFClusterMaker(const char *name, const char *title)
                 :ATLFMaker(name,title)
{
   SetMinETCalo();
   SetDeltaRBarrel1();
   SetDeltaRForward1();
   SetDeltaRBarrel2();
   SetDeltaRForward2();
   SetEtaCoverage();
   SetMinETInitiator();
   SetMinPTBfield();
   SetMinETCell();
   SetBarrelForwardEta();
   SetGranBarrelEta();
   SetGranBarrelPhi();
   SetCellsEsharing();

   m_CellIndex  = new Int_t[kMAXCELLS];
   m_CellCount  = new Int_t[kMAXCELLS];
   m_CellFlag   = new Int_t[kMAXCELLS];
   m_CellEta    = new Float_t[kMAXCELLS];
   m_CellPhi    = new Float_t[kMAXCELLS];
   m_CellSumPT  = new Float_t[kMAXCELLS];

   m_Fruits     = new TClonesArray("ATLFCluster",5,kFALSE);
   m_BranchName = "Clusters";
   m_Nclusters  = 0;
   m_Ncells     = 0;
//   Save();   //by default clustera are not saved
}

//_____________________________________________________________________________
 ATLFClusterMaker::~ATLFClusterMaker()
{
   delete [] m_CellIndex;
   delete [] m_CellCount;
   delete [] m_CellFlag;
   delete [] m_CellEta;
   delete [] m_CellPhi;
   delete [] m_CellSumPT;
}

//_____________________________________________________________________________
 void  ATLFClusterMaker::AddCluster(Int_t code, Int_t ncells, Int_t nparticles, Int_t use, Float_t eta0, Float_t phi0,
                             Float_t eta, Float_t phi, Float_t et)
{
//            Add a new cluster to the list of clusters

 //Note the use of the "new with placement" to create a new cluster object.
 //This complex "new" works in the following way:
 //   clusters[i] is the value of the pointer for cluster number i in the TClonesArray
 //   if it is zero, then a new cluster 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 cluster parameters.
 // This technique should save a huge amount of time otherwise spent
 // in the operators new and delete.

   TClonesArray &clusters = *(TClonesArray*)m_Fruits;
   new(clusters[m_Nclusters++]) ATLFCluster(code, ncells,nparticles, use, eta0,phi0,eta,phi,et);
}

//_____________________________________________________________________________
 void ATLFClusterMaker::Clear(Option_t *option)
{
//    Reset Cluster Maker

   m_Nclusters  = 0;
   m_Ncells     = 0;
   ATLFMaker::Clear(option);
}


//_____________________________________________________________________________
 void ATLFClusterMaker::Init()
{
// Create histograms
   m_Multiplicity = new TH1F("CluMultiplicity","clusters multiplicity",10,0,10);
   m_DeltaPhi     = new TH1F("CluDeltaPhi","B-field delta_phi vers p_T",50,0,50);
   m_Counters     = new TH1F("CluCounters","B-field counters vers p_T",50,0,50);
   m_PhiClu       = new TH1F("CluPhiClu ","delta_phi clu-bari_centre_particle",50,-0.5,0.5);
   m_EtaClu       = new TH1F("CluEtaClu","delta_eta clu--bari_centre_particle",50,-0.5,0.5);
   m_DeltaR       = new TH1F("CluDeltaR","delta_r clu--bari_centre_particle",50,0,0.5);
}

//_____________________________________________________________________________
 Float_t ATLFClusterMaker::FieldPhi(Float_t pt, Float_t eta, Float_t theta)
{
//     parametrizes effects of the B-field
//     shift in the phi position of the charged particle
//     parametrization from L. Poggioli

   Float_t def;

   if (TMath::Abs(eta) < 1.4) {
      def = -150.              * 0.006/pt/2.0;
   } else {
      Float_t abstantheta = TMath::Abs(TMath::Tan(theta));
      def = -350. * abstantheta* 0.006/pt/2.0;
   }
   return def;

}

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

   m_DeltaPhi->Divide(m_Counters);
}

//_____________________________________________________________________________
 Float_t ATLFClusterMaker::HadronResolution(Int_t lumi, Float_t ene, Float_t eta, Float_t pt, Float_t rcone)
{
//     parametrizes smearing for hadronic energy deposition
//     parametrization from L. Poggioli
//     pile-up added as in HADCALO

  Float_t epileup = 0;
  Float_t aa, bb, cc, dd, sigma;
  Float_t sqrtene = TMath::Sqrt(ene);
  Float_t abseta  = TMath::Abs(eta);

  if(rcone <= 0.4)                epileup = 0.4;
  if(rcone == 0.4)                epileup = 7.5;
  if(rcone > 0.4 && rcone <= 0.5) epileup = 12.0;
  if(rcone > 0.5 && rcone <= 0.7) epileup = 18.0;
  if(rcone > 0.7)                 epileup = 20.0;

  if(lumi <= 1) {
     while(1) {
        gRandom->Rannor(aa,bb);
        if(abseta < 3.) sigma = aa*0.5/sqrtene + bb*0.03;
        else            sigma = aa*1.0/sqrtene + bb*0.07;
        if(1.+sigma > .0) return sigma;
     }
  } else if(lumi == 2) {
     while(1) {
        gRandom->Rannor(aa,bb);
        gRandom->Rannor(cc,dd);
        if(abseta < 3.) sigma = aa*0.5/sqrtene + bb*0.03  + cc*epileup/pt;
        else            sigma = aa*1.0/sqrtene + bb*0.07  + cc*epileup/pt;
        if(1.+sigma > .0)return sigma;
     }
  }
  return 0;
}

//_____________________________________________________________________________
 Int_t ATLFClusterMaker::Make()
{
//.............................................
//  This function deposits the transverse energies of undecayed particles
//  (except neutrinos, muons and SUSY LSP) in  ETA times PHI cells
//  with default granularity: 0.1x0.1 for ETA<3 and 0.2x0.2 for ETA>3. 
//  This granularity can be modified, but with the limitation that the granularity
//  of cells beyond the barrel/forward transition (ETA=3) is twice as large as 
//  inside. The algorithm then groups these cells into clusters 
//  if the sum of their transverse energy is above a threshold value. 
//  The energies deposited in the cells are stored also in a cell map
//  Cells used for cluster reconstruction  are flagged.
//  The effect of the  solenoidal magnetic field is only parametrized as
//  a shift in the PHI position of charged particles, which is
//  calculated by the function FieldPhi. Charged particles with PT
//  below the default threshold of 0.5GeV are not stored  if the magnetic field 
//  is on.

//  There is the possibility to activate the energy sharing of cells belonging to 
//  overlaping jets (defaults: no sharing).  In that case the corresponding 
//  cell energy is shared according to the relative energies of each overlaping jet.

//  Clusters are arranged in order of decreasing  ET.

//  The reconstructed clusters are added to the list of clusters.
//.............................................

   Float_t rconeb, rconef, ptlrat, abseta;
   Float_t eta, phi, theta;
   eta= phi= theta= 0;
   Float_t detphi, pt, pz, pE, chrg;
   Int_t i, nparticles, ncells, nclusters, kc, ic, icell, ij, k, KF, KS;
   Int_t ieta, iphi, ietph;
   Int_t nbeta = Int_t(2*m_EtaCoverage/m_GranBarrelEta); 
   Int_t nbphi = Int_t(2*kPI/m_GranBarrelPhi);
           
   rconeb = m_DeltaRBarrel1;
   rconef = m_DeltaRForward1;
   if (gATLFast->Mode() == 2) {
      rconeb= m_DeltaRBarrel2;
      rconef= m_DeltaRForward2;
   }
   Float_t rconeb2 = rconeb*rconeb;
   Float_t rconef2 = rconef*rconef;

//........................................................
//... Get pointers to MCparticles arrays and ClonesArray
//........................................................
   ATLFMCMaker *mcarlo = gATLFast->MCMaker();
   TClonesArray *particles = mcarlo->Fruits();
   nparticles = particles->GetEntriesFast();
   TMCParticle *part;
   Float_t *pPT  = mcarlo->PT();
   Float_t *pPT2 = mcarlo->PT2();
   Float_t *pEta = mcarlo->Eta();
   Float_t *pPhi = mcarlo->Phi();
   Int_t   *Index= mcarlo->Index();
   
//........................................................
//...Loop over all particles. Find cell that was hit by given particle. 
//........................................................
   Float_t setacov = TMath::SinH(m_EtaCoverage);
   ptlrat=1./(setacov*setacov);
   Int_t kFLSP = gATLFast->SUSYcodeLSP();

   ncells = 0;
   for (Int_t ind=0;ind<mcarlo->Nstables();ind++) {
      i = Index[ind];
      part = (TMCParticle*)particles->UncheckedAt(i);
      KS = part->GetKS();
      KF = part->GetKF();
      pz = part->GetPz();
      if (pPT2[i] <= ptlrat*pz*pz) continue;
      kc = mcarlo->Compress(KF); 
      if (kc == 0  || kc == 12 || kc == 14 || kc == 16) continue; 
      if (kc == 13 || kc == kFLSP) continue; 
      detphi = 0;
      pt     = pPT[i];
      eta    = pEta[i];
      phi    = pPhi[i];
      pE     = part->GetEnergy();
      if (gATLFast->Bfield() == 1 && mcarlo->Charge(KF) != 0) {
        if (pt < m_MinPTBfield) continue;
        chrg = mcarlo->Charge(KF)/3.;
        if (TMath::Abs(pz/pE) < 1.) {
          theta = TMath::ACos(pz/pE);
        } else {
          if (pz > 0) theta = 0;
          if (pz < 0) theta = kPI;
        }
        detphi= chrg*FieldPhi(pt,eta,theta);
        m_DeltaPhi->Fill(pt,TMath::Abs(detphi));
        m_Counters->Fill(pt);
      } 
      phi   += detphi;
      abseta = TMath::Abs(eta);
      if ( abseta < m_BarrelForwardEta) {
           ieta = Int_t((eta+m_EtaCoverage)/2.0/m_EtaCoverage*nbeta)+1;
           iphi = Int_t((phi+kPI)/2.0/kPI*nbphi)+1;
      } else {
           ieta = 2*Int_t((eta+m_EtaCoverage)/2.0/m_EtaCoverage*nbeta/2.0)+1;
           iphi = 2*Int_t((phi+kPI)/2.0/kPI*nbphi/2.0)+1; 
      }
      ietph = nbphi*ieta + iphi;
//........................................................
//...Add to cell already hit, or book new cell.
//........................................................
      icell = ncells; 
      for (ic=0;ic<ncells;ic++) {
         if (ietph == m_CellIndex[ic]) { icell = ic; break; }
      }
      if (icell == ncells) {  // found new cell
         ncells++;
         m_CellIndex[icell] = ietph;
         m_CellCount[icell] = 1;
         m_CellFlag[icell]  = 2;
         if (abseta < m_BarrelForwardEta) {
            m_CellEta[icell] = 2*m_EtaCoverage*(ieta -1 +0.5)/nbeta - m_EtaCoverage;
            m_CellPhi[icell] = k2PI*(iphi -1 +0.5)/nbphi - kPI;
         } else {
            m_CellEta[icell] = 2*m_EtaCoverage*ieta/nbeta - m_EtaCoverage;
            m_CellPhi[icell] = k2PI*iphi/nbphi - kPI;
         }
         m_CellSumPT[icell]  = pt;
      } else {         // increment already found cell
         m_CellCount[icell]++;
         m_CellSumPT[icell] += pt;
      }
   }  // end of loop on particles      
         
//........................................................
//...Remove cells below threshold. 
//........................................................
   Int_t noldcells = ncells;
   ncells = 0;
   for (ic =0; ic<noldcells;ic++) {
       if (m_CellSumPT[ic] < m_MinETCell) continue;
       m_CellIndex[ncells] = m_CellIndex[ic];
       m_CellCount[ncells] = m_CellCount[ic];
       m_CellFlag[ncells]  = m_CellFlag[ic];
       m_CellEta[ncells]   = m_CellEta[ic];
       m_CellPhi[ncells]   = m_CellPhi[ic];
       ncells++;
   }

//........................................................
//...Find initiator cell: the one with highest pT of not yet used ones. 
//........................................................
   Int_t   sortindex[kMAXCLU];
   Int_t   clusterCells[kMAXCLU], clusterParts[kMAXCLU], clusterFlag[kMAXCLU];
   Float_t clusterEta0[kMAXCLU],  clusterPhi0[kMAXCLU];
   Float_t clusterEta[kMAXCLU],   clusterPhi[kMAXCLU],   clusterET[kMAXCLU];
   Int_t   cluCode[kMAXCLU];
   Int_t   cluCells[kMAXCLU],     cluParts[kMAXCLU],     cluFlag[kMAXCLU];
   Float_t cluEta0[kMAXCLU],      cluPhi0[kMAXCLU];
   Float_t cluEta[kMAXCLU],       cluPhi[kMAXCLU],       cluET[kMAXCLU];
   Float_t ehad[kMAXCLU], iacc[kMAXCLU];
   Float_t etmin = 0;
   Float_t etmax, dr,dphia,phic,deta,dphi,adphi,acluphi,esum,detaphi;
   Float_t etaBarrel;
   Int_t icmax = 0;
   nclusters = 0;
   while(1) {
      etmax = 0;
      for (ic=0;ic<ncells;ic++) {
         if (m_CellFlag[ic] != 2)      continue;
         if (m_CellSumPT[ic] <= etmax) continue;
         icmax = ic;
         eta   = m_CellEta[ic];
         phi   = m_CellPhi[ic];
         etmax = m_CellSumPT[ic];
      }
      if (etmax < m_MinETInitiator) break;
      m_CellFlag[icmax]   = 1;
      cluCode[nclusters]  = 98;
      cluCells[nclusters] = 1;
      cluParts[nclusters] = m_CellCount[icmax];
      cluFlag[nclusters]  = 1;
      cluEta0[nclusters]  = eta;
      cluPhi0[nclusters]  = phi;
      cluEta[nclusters]   = 0;
      cluPhi[nclusters]   = 0;
      cluET[nclusters]    = 0;
 
//........................................................
//...Sum up unused cells within required distance of initiator. 
//........................................................
      abseta    = TMath::Abs(eta);
      etaBarrel = abseta - m_BarrelForwardEta;
      for (ic=0;ic<ncells;ic++) {
         if (m_CellFlag[ic] == 0) continue;
         phic  = m_CellPhi[ic];
         dphia = TMath::Abs(phic - phi);
         if (dphia > kPI)  phic += TMath::Sign(k2PI,phi);
         deta    = m_CellEta[ic] - eta;
         dphi    = phic - phi;
         detaphi = deta*deta+dphi*dphi;
         if (etaBarrel < 0 && detaphi > rconeb2) continue;
         if (etaBarrel > 0 && detaphi > rconef2) continue;
         m_CellFlag[ic]      = -m_CellFlag[ic];
         cluCells[nclusters]++;
         cluParts[nclusters] += m_CellCount[ic];
         cluEta[nclusters]   += m_CellSumPT[ic]*m_CellEta[ic];
         cluPhi[nclusters]   += m_CellSumPT[ic]*phic;
         cluET[nclusters]    += m_CellSumPT[ic];
      }
 
//........................................................
//...Reject cluster below minimum ET, else accept. 
//........................................................
      if (m_CellsEsharing == 0) etmin = m_MinETCalo;
      if (m_CellsEsharing == 1) etmin = m_MinETInitiator;
      if (cluET[nclusters] < etmin) {
         for (ic=0;ic<ncells;ic++) {
            if (m_CellFlag[ic] < 0) m_CellFlag[ic] = -m_CellFlag[ic];
         }
      } else {
         cluEta[nclusters] /= cluET[nclusters];
         cluPhi[nclusters] /= cluET[nclusters];
         acluphi = TMath::Abs(cluPhi[nclusters]);
         if (acluphi > kPI) cluPhi[nclusters] -= TMath::Sign(k2PI,cluPhi[nclusters]);
         for (ic=0;ic<ncells;ic++) {if (m_CellFlag[ic] < 0) m_CellFlag[ic] = 0;}
         nclusters++;
         if (nclusters >= kMAXCLU) {
            Error("Make","Too many clusters=%d",nclusters);
         }
      } 
   }  //end of while

//........................................................
// ...Arrange clusters in falling ET sequence. 
//........................................................

   gATLFast->SortDown(nclusters, cluET, sortindex);

   for (k=0;k<nclusters;k++) {
      ij = sortindex[k];
      clusterCells[k]  = cluCells[ij];
      clusterParts[k]  = cluParts[ij];
      clusterFlag[k]   = cluFlag[ij];
      clusterEta0[k]   = cluEta0[ij];
      clusterPhi0[k]   = cluPhi0[ij];
      clusterEta[k]    = cluEta[ij];
      clusterPhi[k]    = cluPhi[ij];
      clusterET[k]     = cluET[ij];
   }

//.....histogram NCLU
//   printf("nparticles=%d, ncells=%d, nclusters=%dn",nparticles,ncells,nclusters);
   m_Multiplicity->Fill(nclusters+0.01);

//........................................................
//...recalibrate for the energy sharing
//...loop on cells, share energy deposition in cells between clusters
//........................................................
   if (m_CellsEsharing == 1) {
      for (i=0;i<nclusters;i++) ehad[i] = 0;
      for (i=0;i<ncells;i++) {
         eta       = m_CellEta[i];
         phi       = m_CellPhi[i];
         esum      = 0;
         abseta    = TMath::Abs(eta);
         etaBarrel = abseta -m_BarrelForwardEta;
         for (k=0;k<nclusters;k++) {
            iacc[k] = 0;
            deta    = eta - clusterEta0[k];
            dphi    = phi - clusterPhi0[k];
            adphi   = TMath::Abs(dphi) -k2PI;
            if (TMath::Abs(dphi) > kPI) dr = deta*deta + adphi*adphi;
            else                        dr = deta*deta + dphi*dphi;
            if (etaBarrel <= 0 && dr > rconeb2) continue;
            if (etaBarrel >  0 && dr > rconef2) continue;
              //...cluster K contains the running cell
              //...esum is the total energy of jets containing the running cell
            iacc[k] = 1;
            esum += clusterET[k]*TMath::CosH(clusterEta0[k]);
         }

         //...redistribute the energy of the running cell among corresponding clusters
         for (k=0;k<nclusters;k++) {
            if (iacc[k] == 0) continue;
              //...cluster K contains the running cell
              //...the cell ene is shared among the overlaping clusters according to their energy
            ehad[k] += m_CellSumPT[i]*TMath::CosH(m_CellEta[i])*clusterET[k]*TMath::CosH(clusterEta0[k])/esum;
         }
      }

      //...update cluster energies
      for (k=0;k<nclusters;k++) {
         clusterET[k] = ehad[k]/TMath::CosH(clusterEta0[k]);
      }

      //........................................................
      //.....rearrange clusters in falling E_T sequence
      //........................................................

      gATLFast->SortDown(nclusters, clusterET, sortindex);

      for (k=0;k<nclusters;k++) {
         ij = sortindex[k];
         if (clusterET[ij] < m_MinETCalo) { nclusters = k; break; }
         cluCells[k]  = clusterCells[ij];
         cluParts[k]  = clusterParts[ij];
         cluFlag[k]   = clusterFlag[ij];
         cluEta0[k]   = clusterEta0[ij];
         cluPhi0[k]   = clusterPhi0[ij];
         cluEta[k]    = clusterEta[ij];
         cluPhi[k]    = clusterPhi[ij];
         cluET[k]     = clusterET[ij];
      }
   }  // end of loop for energy sharing

//........................................................
//....Store clusters in Root ClonesArray
//........................................................
   m_Ncells    = ncells;
   m_Nclusters = 0;
   for (k=0;k<nclusters;k++) {
      AddCluster(cluCode[k],
                 cluCells[k],
                 cluParts[k],
                 cluFlag[k],
                 cluEta0[k],
                 cluPhi0[k],
                 cluEta[k],
                 cluPhi[k],
                 cluET[k]);
   }

//........................................................
//....reconstruct baricenter of particles
//........................................................
   Float_t etarec, ptrec, phirec;
   for (k=0;k<nclusters;k++) {
      etarec = ptrec = phirec = 0;
      for (i=0;i<nparticles;i++) {
         part = (TMCParticle*)particles->UncheckedAt(i);
         KS = part->GetKS();
         if (KS <= 0  || KS > 10) continue;
         KF = part->GetKF();
         pz = part->GetPz();
         if (pPT2[i] <= ptlrat*pz*pz) continue;
         kc = mcarlo->Compress(KF);
         if (kc == 0  || kc == 12 || kc == 14 || kc == 16) continue;
         if (kc == 13 || kc == kFLSP) continue;
         detphi = 0;
         pt  = pPT[i];
         eta = pEta[i];
         phi = pPhi[i];
         pE  = part->GetEnergy();
         if (gATLFast->Bfield() == 1 && mcarlo->Charge(KF) != 0) {
            if (pt < m_MinPTBfield) continue;
            chrg = mcarlo->Charge(KF)/3.;
            if (TMath::Abs(pz/pE) < 1.) {
               theta = TMath::ACos(pz/pE);
            } else {
              if (pz > 0) theta = 0;
              if (pz < 0) theta = kPI;
            }
            detphi= chrg*FieldPhi(pt,eta,theta);
         } 
         phi   += detphi;
         dphia  = TMath::Abs(cluPhi0[k] - phi);
         if (dphia > kPI) dphia -= k2PI;
         deta      = cluEta0[k] - eta;
         etaBarrel = TMath::Abs(cluEta0[k]) - m_BarrelForwardEta;
         detaphi   = deta*deta + dphia*dphia;
         if (etaBarrel < 0 &&  detaphi > rconeb2) continue;
         if (etaBarrel > 0 &&  detaphi > rconef2) continue;
         ptrec  += pt;
         etarec += eta*pt;
         phirec += phi*pt;
      }

      etarec /= ptrec;
      phirec /= ptrec;
      deta    = etarec - cluEta[k];
      dphi    = phirec - cluPhi[k];
      dr      = TMath::Sqrt(deta*deta + dphi*dphi);
      m_EtaClu->Fill(deta);
      m_PhiClu->Fill(dphi);
      m_DeltaR->Fill(dr);
   }
   return 0;
}

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

//_____________________________________________________________________________
 void ATLFClusterMaker::RemoveCluster(const Int_t cluster)
{
   TClonesArray &clusters = *(TClonesArray*)m_Fruits;
   clusters.RemoveAt(cluster);
   clusters.Compress();
   m_Nclusters--;
}