#include "modules/MadGraphClassFilter.h"

#include "modules/MadGraphParticleClassifier.h"

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

#include "ExRootAnalysis/ExRootFilter.h"

#include "TClonesArray.h"

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

#include <iostream>

using namespace std;

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

MadGraphClassFilter::MadGraphClassFilter()
{
}

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

MadGraphClassFilter::~MadGraphClassFilter()
{
}

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

void MadGraphClassFilter::Init()
{
  TString className;
  ExRootConfParam param, classParticles;
  
  Bool_t extendable;
  Int_t i, j, status, pid, sizeParam, sizeParticles;

  // import ROOT tree branch

  fBranchParticle = UseBranch("Particle");

  // 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);
    }
  }

  // allow classifier to create additional classes for particles with unknown PID

  extendable = GetBool("AllowExtendableClasses", kTRUE);

  fClassifier->SetExtendable(extendable);

  // create output arrays

  fOutputArray = ExportArray("particles");
}

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

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

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

void MadGraphClassFilter::Process()
{
  TObjArray *subarray;
  Int_t category;

  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)
  {
    subarray = fFilter->GetSubArray(fClassifier, category);
    if(subarray)
    {
      subarray->SetName(fClassifier->GetCategoryClassName(category));
      fOutputArray->Add(subarray);
  
      // sort particles by PT
      ExRootLHEFParticle::fgCompare = ExRootComparePT<ExRootLHEFParticle>::Instance();
      subarray->Sort();
    }
  }
}

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