/** \class EnergySmearing
 *
 *  Performs energy resolution smearing.
 *
 *  $Date: 2012-11-19 14:05:19 +0100 (Mon, 19 Nov 2012) $
 *  $Revision: 829 $
 *
 *
 *  \author P. Demin - UCL, Louvain-la-Neuve
 *
 */

#include "modules/EnergySmearing.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;

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

EnergySmearing::EnergySmearing() :
  fItInputArray(0)
{
}

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

EnergySmearing::~EnergySmearing()
{
}

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

void EnergySmearing::Init()
{
  map< Int_t, DelphesFormula * >::iterator itFormulaMap;
  ExRootConfParam param;
  DelphesFormula *formula;
  Int_t i, size;

  // read resolution formulas
  param = GetParam("ResolutionFormula");
  size = param.GetSize();
  
  fFormulaMap.clear();
  for(i = 0; i < size/2; ++i)
  {
    formula = new DelphesFormula;
    formula->Compile(param[i*2 + 1].GetString());

    fFormulaMap[param[i*2].GetInt()] = formula;
  }

  // set default Resolution formula
  itFormulaMap = fFormulaMap.find(0);
  if(itFormulaMap == fFormulaMap.end())
  {
    formula = new DelphesFormula;
    formula->Compile("0.0");

    fFormulaMap[0] = formula;
  }

  // import array with output from filter/classifier module

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

  // create output arrays

  fOutputArray = ExportArray(GetString("OutputArray", "candidates"));
}

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

void EnergySmearing::Finish()
{
  map< Int_t, DelphesFormula * >::iterator itFormulaMap;
  DelphesFormula *formula;
  
  if(fItInputArray) delete fItInputArray;
  
  for(itFormulaMap = fFormulaMap.begin(); itFormulaMap != fFormulaMap.end(); ++itFormulaMap)
  {
    formula = itFormulaMap->second;
    if(formula) delete formula;
  }
}

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

void EnergySmearing::Process()
{
  Candidate *candidate, *mother;
  Double_t energy, eta, phi;
  Int_t pdgCode; 
  map< Int_t, DelphesFormula * >::iterator itEfficiencyMap, itFormulaMap;
  DelphesFormula *formula;

  fItInputArray->Reset();
  while((candidate = static_cast<Candidate*>(fItInputArray->Next())))
  {
    const TLorentzVector &candidatePosition = candidate->Position;
    const TLorentzVector &candidateMomentum = candidate->Momentum;
    pdgCode = TMath::Abs(candidate->PID);
    eta = candidatePosition.Eta();
    phi = candidatePosition.Phi();
    energy = candidateMomentum.E();

    // find smearing formula
    itFormulaMap = fFormulaMap.find(pdgCode);
    if(itFormulaMap == fFormulaMap.end())
    {
      itFormulaMap = fFormulaMap.find(0);
    }
    formula = itFormulaMap->second;

    // apply smearing formula
    energy = gRandom->Gaus(energy, formula->Eval(0.0, eta, 0.0, energy));
    
    if(energy <= 0.0) continue;

    mother = candidate;
    candidate = static_cast<Candidate*>(candidate->Clone());
    candidate->Momentum.SetE(energy);
    candidate->Mother = mother;
    
    fOutputArray->Add(candidate);
  }
}

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