#include "modules/Tracking.h"

#include "classes/DelphesClasses.h"
#include "classes/DelphesFactory.h"
#include "classes/DelphesFormula.h"

#include "ExRootAnalysis/ExRootResult.h"
#include "ExRootAnalysis/ExRootFilter.h"
#include "ExRootAnalysis/ExRootClassifier.h"

#include "TMath.h"
#include "TString.h"
#include "TFormula.h"
#include "TRandom3.h"
#include "TObjArray.h"
#include "TDatabasePDG.h"
#include "TLorentzVector.h"

#include <algorithm> 
#include <stdexcept>
#include <iostream>
#include <sstream>

using namespace std;

//------------------------------------------------------------------------------

Tracking::Tracking() :
  fItInputArray(0)
{
}

//------------------------------------------------------------------------------

Tracking::~Tracking()
{
}

//------------------------------------------------------------------------------

void Tracking::Init()
{
  ExRootConfParam paramEff, paramEffEta, paramEffPt, paramEffVal, paramEffPart;
  ExRootConfParam paramPtRes, paramPtResEta, paramPtResPt, paramPtResVal, paramPtResPart;
  
  Int_t i, size, ref=0, pid=0;
  
  fTrkMaxEta = GetDouble("TrackMaxEta", 2.5);
  fTrkMinPt  = GetDouble("TrackMinPt" , 1);
 
 
 // fTrkEff = GetDouble("TrackingEfficiency", 90); //for a flat efficiency
  
  paramEff   = GetParam("TrackingEfficiency");
  paramPtRes = GetParam("TrackingPtResolution");
  
  size = paramEff.GetSize();
  //default tracking efficiency all are set to 90% by default
  TInt EtaInt = make_pair(0.0,50);
  TInt PtInt  = make_pair(0.0,10000.0);
  TPtEtaInt PtEtaInt = make_pair(EtaInt,PtInt);
  TTrkMap TrkEffMap; 
  TrkEffMap[PtEtaInt]=50;    //only one entry in this simple map 
  
  fTrkEffPartMap[11]  = TrkEffMap;  // electrons
  fTrkEffPartMap[13]  = TrkEffMap;  // muons
  fTrkEffPartMap[999] = TrkEffMap; // hadrons (pid 999 does not exist ...)
  
  // read efficiencies from cfg file
  for(i = 0; i < size/4; ++i)
  {
    paramEffPart = paramEff[i*4];
    paramEffEta  = paramEff[i*4+1];
    paramEffPt   = paramEff[i*4+2];
    paramEffVal  = paramEff[i*4+3];
    pid = paramEffPart[0].GetInt();
   
    if(ref!=pid) TrkEffMap.clear();
    ref = (ref==pid)? ref : pid;
    
    EtaInt    = make_pair(paramEffEta[0].GetDouble(),paramEffEta[1].GetDouble());
    PtInt     = make_pair(paramEffPt[0].GetDouble(),paramEffPt[1].GetDouble());
    PtEtaInt  = make_pair(EtaInt,PtInt);
    TrkEffMap[PtEtaInt] = paramEffVal.GetDouble();
    fTrkEffPartMap[pid] = TrkEffMap;
   } 
 
 
  size = paramPtRes.GetSize();
  //default tracking pt res all are set to 90% by default
  TTrkMap TrkPtResMap; 
  TrkPtResMap[PtEtaInt]=1.;    //only one entry in this simple map 
  
  fTrkPtResPartMap[11]  = TrkPtResMap;  // electrons
  fTrkPtResPartMap[13]  = TrkPtResMap;  // muons
  fTrkPtResPartMap[999] = TrkPtResMap; // hadrons (pid 999 does not exist ...)
  
  // read resolutions from cfg file
  for(i = 0; i < size/4; ++i)
  {
    paramPtResPart = paramPtRes[i*4];
    paramPtResEta  = paramPtRes[i*4+1];
    paramPtResPt   = paramPtRes[i*4+2];
    paramPtResVal  = paramPtRes[i*4+3];
    pid = paramPtResPart[0].GetInt();
   
    if(ref!=pid) TrkPtResMap.clear();
    ref = (ref==pid)? ref : pid;
    
    EtaInt    = make_pair(paramPtResEta[0].GetDouble(),paramPtResEta[1].GetDouble());
    PtInt     = make_pair(paramPtResPt[0].GetDouble(),paramPtResPt[1].GetDouble());
    PtEtaInt  = make_pair(EtaInt,PtInt);
    TrkPtResMap[PtEtaInt] = paramPtResVal.GetDouble();
    fTrkPtResPartMap[pid] = TrkPtResMap;
   } 
 
 
 
  // import array with output from filter/classifier module

  fInputArray = ImportArray(GetString("InputArray", "ParticlePropagator/candidates"));
  fItInputArray = fInputArray->MakeIterator();

  // create output arrays

  fOutputArray = ExportArray(GetString("TrackOutputArray", "tracks"));

}

//------------------------------------------------------------------------------

void Tracking::Finish()
{
}

//------------------------------------------------------------------------------

void Tracking::Process()
{ 
  Candidate *candidate, *mother;
  TLorentzVector candidatePosition, candidateMomentum;
  Double_t x, y, z, px, py, pz, pt, pt2, eta, phi, e, m, q;
  Int_t pid; 
  TRandom3 * grandom = new TRandom3(0);
  
  fItInputArray->Reset();
  while((candidate = static_cast<Candidate*>(fItInputArray->Next())))
  { 	      
    candidatePosition = candidate->Position;
    candidateMomentum = candidate->Momentum;
    x = candidatePosition.X()*1.0E-3;
    y = candidatePosition.Y()*1.0E-3;
    z = candidatePosition.Z()*1.0E-3;
    m = candidate->Mass;
    q = candidate->Charge;
    pid = TMath::Abs(candidate->PID);
    px = candidateMomentum.Px();
    py = candidateMomentum.Py();
    pz = candidateMomentum.Pz();
    pt = candidateMomentum.Pt();
    pt2 = candidateMomentum.Perp2();
    e   = candidateMomentum.E();
    eta = candidateMomentum.Eta();
    phi = candidateMomentum.Phi();
    
    // check that particle is charged, and that it passes standard acceptance cuts
    if(q==0 || TMath::Abs(eta)>fTrkMaxEta || pt<fTrkMinPt ) continue;
    
    // check that track is lucky 
    if( grandom->Uniform(0,100) > Value(pid,pt,eta,fTrkEffPartMap) ) continue;
    
    //smearing pt according to tracking resolution
    Double_t sPT = grandom->Gaus(pt, Value(pid,pt,eta,fTrkPtResPartMap)*pt*0.01 );
    if(sPT <=0) continue;
    else
     { 
     mother = candidate;
     candidate = static_cast<Candidate*>(candidate->Clone());
     candidate->Momentum.SetPtEtaPhiE(sPT,eta,phi,sPT*cosh(eta));
     candidate->Mother = mother;
     }  
    
    fOutputArray->Add(candidate);
    }
 delete grandom;
}

//------------------------------------------------------------------------------

Double_t Tracking::Value(Int_t pid, Double_t pt, Double_t eta, TTrkPartMap TrkPartMap)
{
 
 Double_t eff=0;
 if(pid!=11 &&pid!=13) pid=999;
 TTrkMap TrkEffMap = TrkPartMap[pid];
 
 for(It_tm it = TrkEffMap.begin(); it != TrkEffMap.end(); it++) 
 {
  TPtEtaInt inter = it->first;  
  TInt eta_int = inter.first;
  TInt pt_int  = inter.second;
  Double_t etamin = eta_int.first;
  Double_t etamax = eta_int.second;
  Double_t ptmin = pt_int.first;
  Double_t ptmax = pt_int.second;
  Double_t EFF= it->second;
  
 // cout<<"PID: "<<pid<<" ,ETA: "<<eta<<", PT: "<<pt<<",eff:  "<<EFF<<endl;
  
  Bool_t ok = pt>ptmin && pt<ptmax && TMath::Abs(eta)>etamin && TMath::Abs(eta)<etamax;
  if(ok)eff=EFF;
  if(ok)break;
 }
return eff;

}