#include "Stntuple/jet/TJetUtil.hh"
#include "Stntuple/obj/TStnEvent.hh"
#include "Stntuple/obj/TStnJetBlock.hh"
#include "Stntuple/obj/TStnHeaderBlock.hh"
#include "Stntuple/obj/TGenpBlock.hh"
#include "Stntuple/obj/TStnVertexBlock.hh"
#include "Stntuple/alg/TStntuple.hh"
#include "JetUser/JetEnergyCorrections.hh" 


//_____________________________________________________________________________
Int_t TJetUtil::FillJetCorr(TStnEvent* event, char* jetBlockName) {

  TStnHeaderBlock* fHeaderBlock =  
    (TStnHeaderBlock*)event->UnpackDataBlock("HeaderBlock"); 

  TStnJetBlock* fJetBlock = 
    (TStnJetBlock*)event->UnpackDataBlock(jetBlockName); 
  if(!fJetBlock) return 1;

  int run = fHeaderBlock->RunNumber(); 
  int evt = fHeaderBlock->EventNumber(); 
  bool qMc = fHeaderBlock->McFlag();  

  int coneSize=0;  // 0=0.4, 1=0.7 and 2=1.0 
  float fConeSize = fJetBlock->ConeSize();
  if(fabs(fConeSize - 0.4)<0.1) {
    coneSize = 0;
  } else if(fabs(fConeSize - 0.7)<0.1) {
    coneSize = 1;
  } else if(fabs(fConeSize - 1.0)<0.1) {
    coneSize = 2;
  } else {
    return 2;
  }

  TStnVertexBlock* fZVertexBlock = 
    (TStnVertexBlock*)event->UnpackDataBlock("ZVertexBlock"); 
  if(!fZVertexBlock) return 3;

  int nzVertex = fZVertexBlock->NVertices(12);

  int version,syscode,level,imode,nv;

  // gen5 and gen6
  version = 5;  // for data and MC
  imode = (qMc ? 0 : 1);
  syscode=0;   // systemtic code: 0=none

  // full corrections
  level = 7;
  nv = nzVertex;

  JetEnergyCorrections jcorFull(
	   "JetCorrections","JetCorrections",level,nv,coneSize,
	   version,syscode,run,imode); 

  // corrections for met
  level = 5;
  nv = 1;  

  JetEnergyCorrections jcorMet(
	   "JetCorrections","JetCorrections",level,nv,coneSize,
	   version,syscode,run,imode); 
  
  for(int i=0; i<fJetBlock->NJets(); i++) {
    TStnJet* jet = fJetBlock->Jet(i);

    float em = jet->Emfr(); 
    float &emFraction = em; 
    float DetJet = jet->DetEta(); 
    float pt = jet->Momentum()->Pt();
    
    float scaleFactor;
    scaleFactor = jcorFull.doEnergyCorrections(pt,emFraction,DetJet); 
    jet->SetEtCorrM(scaleFactor);// number = 1+eps

    scaleFactor = jcorMet.doEnergyCorrections(pt,emFraction,DetJet); 
    jet->SetEtCorr(scaleFactor);// number = 1+eps
  }

  return 0;
}


//_____________________________________________________________________________
Int_t TJetUtil::FillJetDr(TStnEvent* event, char* jetBlockName,
		       float etaCut, float etCut, float etCorrCut) {

  TStnJetBlock* fJetBlock = 
    (TStnJetBlock*)event->UnpackDataBlock(jetBlockName); 
  if(!fJetBlock) return 1;

  bool useCorrEt = (etCorrCut>0.0);

  //find rank and closest jet
  for(int i=0; i<fJetBlock->NJets(); i++) {
    TStnJet* jeti = fJetBlock->Jet(i);
    jeti->SetDrMin(999.0);
    jeti->SetRank(999);
    int rank = 0;
    float drMin = 999.0;
    //printf("jet %i %f \n ",i,jeti->Et()*jeti->Corfm());
    if(fabs(jeti->DetEta())<etaCut && jeti->Et()>etCut 
       && jeti->Et()*jeti->Corfm()>etCorrCut) {

      for(int j=0; j<fJetBlock->NJets(); j++) {
	TStnJet* jetj = fJetBlock->Jet(j);
	if(i!=j) {
	  if(fabs(jetj->DetEta())<etaCut && jetj->Et()>etCut 
	     && jetj->Et()*jetj->Corfm()>etCorrCut) {

	    float dr = TStntuple::DeltaR(
		  jeti->Phi(),jeti->DetEta(),
		  jetj->Phi(),jetj->DetEta());
	    //printf("%i %i dr=%f\n ",i,j,dr);
	    if(dr<drMin) drMin = dr;
	  }
	  if(useCorrEt) {
	    if(jetj->Et()*jetj->Corfm()>jeti->Et()*jeti->Corfm()) rank++;
	  } else {
	    if(jetj->Et()>jeti->Et()) rank++;
	  }
	} // i!=j
      } // j
      //printf("drmin=%f\n",drMin);
      jeti->SetDrMin(drMin);
      jeti->SetRank(rank);
    }
  }  // i

  return 0;
}


//_____________________________________________________________________________
Int_t TJetUtil::FillJetType(TStnEvent* event, char* jetBlockName) {

  TStnHeaderBlock* fHeaderBlock =  
    (TStnHeaderBlock*)event->UnpackDataBlock("HeaderBlock"); 

  TStnJetBlock* fJetBlock = 
    (TStnJetBlock*)event->UnpackDataBlock(jetBlockName); 
  if(!fJetBlock) return 1;

  TGenpBlock* fGenpBlock =  
    (TGenpBlock*)event->UnpackDataBlock("GenpBlock"); 
  if(!fGenpBlock) return 1;

  Int_t ngen = fGenpBlock->NParticles();
  TLorentzVector v0;

  for (int j=0; j<fJetBlock->NJets(); j++) {
    TStnJet* jet = fJetBlock->Jet(j);

    float etaj = jet->Eta();
    float etj = jet->Et()*jet->Corfm();
    float phij = TVector2::Phi_0_2pi(jet->Phi());

    float etMax = 0;
    float et;
    int type = 0;
    int ipt = 0;

    for(int i=0; i<ngen; i++) {
      TGenParticle* p = fGenpBlock->Particle(i);
      if(p) {
	int pdg = abs(p->GetPdgCode());
	if(pdg<=5 || pdg==11 || pdg==13 || pdg==15 || pdg==21 || pdg==22) {
	  //if(true) {
	  p->Momentum(v0);
	  float etp = v0.Pt();
	  float etap = (etp<1.e-3? 999.0 : v0.PseudoRapidity() );
	  float phip = TVector2::Phi_0_2pi(v0.Phi());
	  float dr = TStntuple::DeltaR(phij,etaj,phip,etap);
	  if(dr<0.4 && etp>etMax) {
	    type = pdg;
	    ipt = i;
	    et = etp;
	    etMax = et;
	  }
	}
      }   // if(p)
    }   //i loop to ngen
    jet->SetPartType(type);
    jet->SetGenp(ipt);
  } // jet loop

  return 0;
}