#include "modules/MadGraphIsolatedLeptonFinder.h"

#include "modules/MadGraphParticleClassifier.h"

#include "ExRootAnalysis/ExRootClasses.h"
#include "ExRootAnalysis/ExRootFilter.h"

#include "ExRootAnalysis/ExRootFactory.h"
#include "ExRootAnalysis/ExRootCandidate.h"

#include "TString.h"
#include "TClonesArray.h"
#include "TLorentzVector.h"

#include <map>
#include <set>
#include <deque>

#include <iostream>

using namespace std;

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

Double_t MadGraphIsolatedLeptonFinder::GetMinDeltaR(ExRootGenParticle *lepton)
{
  Double_t distMin = 1.0e6;
  Double_t pt1, pt2, dist;
  TLorentzVector vector1, vector2;
  ExRootGenParticle *particle;
 
  vector1.SetPxPyPzE(lepton->Px, lepton->Py, lepton->Pz, lepton->E);
  pt1 = vector1.Pt();

  fItParticle->Reset();
  while(particle = static_cast<ExRootGenParticle*>(fItParticle->Next()))
  {
    if(particle->Status != 1) continue;

    vector2.SetPxPyPzE(particle->Px, particle->Py, particle->Pz, particle->E);
    pt2 = vector2.Pt();

    dist = (pt1 == 0.0 || pt2 == 0.0 ? 1.0e6 : vector1.DeltaR(vector2));

    if(pt2 > fMinPT && dist < distMin)
    {
      distMin = dist;
    }
  }

  return distMin;
}

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

MadGraphIsolatedLeptonFinder::MadGraphIsolatedLeptonFinder() :
  fItParticle(0)
{
}

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

MadGraphIsolatedLeptonFinder::~MadGraphIsolatedLeptonFinder()
{
  if(fItParticle) delete fItParticle;
}

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

void MadGraphIsolatedLeptonFinder::Init()
{
  TString className;
  ExRootConfParam param, classParticles;

  Int_t i, j, pid, sizeParam, sizeParticles;

  // import ROOT tree branch

  fBranchParticle = UseBranch("GenParticle");

  fItParticle = fBranchParticle->MakeIterator();

  // create classifier and filter

  fMinPT = GetDouble("MinPT", 1.0);
  fMinDeltaR = GetDouble("MinDR", 0.1);

  fClassifier = new MadGraphParticleClassifier();
  fFilter = new ExRootFilter(fBranchParticle);

  // read particle classes from configuration file and setup classifier

  param = GetParam("ClassParticles");
  sizeParam = param.GetSize();

  for(i = 0; i < sizeParam/2; ++i)
  {
    className = param[i*2].GetString();
    classParticles = param[i*2 + 1];
    sizeParticles = classParticles.GetSize();

    for(j = 0; j < sizeParticles; ++j)
    {
      pid = classParticles[j].GetInt();
      fClassifier->InsertClassPID(className, pid);
    }
  }

  fClassifier->SetExtendable(kFALSE);

  // create output arrays

  fOutputArray = ExportArray("candidates");
}

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

void MadGraphIsolatedLeptonFinder::Finish()
{
  if(fFilter) delete fFilter;
  if(fClassifier) delete fClassifier;
}

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

void MadGraphIsolatedLeptonFinder::Process()
{
  TObjArray *array = 0;
  ExRootGenParticle *particle = 0;
  ExRootCandidate *candidate = 0;
  ExRootFactory *factory = GetFactory();

  Int_t category;
  TString name;
  TLorentzVector momentum;

  fFilter->Reset();

  // make filter classify particles and fill all subarrays
  // at this point classifier creates additional/missing classes
  fFilter->GetSubArray(fClassifier, 0);

  // loop over all classes and export class names and classified particles
  for(category = 0; category < fClassifier->GetMaxCategories(); ++category)
  {
    array = fFilter->GetSubArray(fClassifier, category);
    name = fClassifier->GetCategoryClassName(category);

    if(array == 0) continue;

    // sort particles by PT
    ExRootGenParticle::fgCompare = ExRootComparePT<ExRootGenParticle>::Instance();
    array->Sort();

    TIter itArray(array);

    while(particle = static_cast<ExRootGenParticle*>(itArray.Next()))
    {
      if(GetMinDeltaR(particle) < fMinDeltaR) continue;

      momentum.SetPxPyPzE(particle->Px, particle->Py, particle->Pz, particle->E);

      candidate = factory->NewCandidate();

      candidate->SetP4(momentum);
      candidate->SetName(name);

      fOutputArray->Add(candidate);
    }
  }
}