#include "modules/MadGraphShowerLeptonSelector.h"

#include "modules/MadGraphParticleClassifier.h"

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

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

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

#include <iostream>
#include <set>

using namespace std;

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

MadGraphShowerLeptonSelector::MadGraphShowerLeptonSelector() :
  fFilter(0), fClassifier(0)
{
}

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

MadGraphShowerLeptonSelector::~MadGraphShowerLeptonSelector()
{
}

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

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

  Int_t i, j, status, pid, sizeParam, sizeParticles;

  // import ROOT tree branch

  fBranchParticle = UseBranch("GenParticle");

  // create classifier and filter

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

  // read particle status from configuration file and setup classifier

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

  for(i = 0; i < sizeParam; ++i)
  {
    status = param[i].GetInt();
    fClassifier->InsertParticleStatus(status);
  }

  // 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 MadGraphShowerLeptonSelector::Finish()
{ 
  if(fFilter) delete fFilter;
  if(fClassifier) delete fClassifier;
}

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

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

  Int_t category, entry;
  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();

    entry = 1;

    TIter itArray(array);

    while(particle = static_cast<ExRootGenParticle*>(itArray.Next()))
    {
      if(particle->M1 < 0) continue;
      
      parent = static_cast<ExRootGenParticle*>(fBranchParticle->At(particle->M1));

      if(parent == 0 || parent->Status != 3) continue;

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

      candidate = factory->NewCandidate();

      candidate->SetP4(momentum);
      candidate->SetName(Form("%s_{%d}", name.Data(), entry ));

      fOutputArray->Add(candidate);

      ++entry;
    }
  }
}
//------------------------------------------------------------------------------