#include "modules/MadGraphClassMerger.h"

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

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

#include <iostream>

using namespace std;

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

MadGraphClassMerger::MadGraphClassMerger() :
  fItInputArray(0)
{
}

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

MadGraphClassMerger::~MadGraphClassMerger()
{
}

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

void MadGraphClassMerger::Init()
{
  // import array with output from filter/classifier module

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

  // create output array

  fOutputArray = ExportArray("candidates");
}

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

void MadGraphClassMerger::Finish()
{
  if(fItInputArray) delete fItInputArray;
}

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

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

  Int_t entry;
  TString name;
  TLorentzVector momentum;

  fItInputArray->Reset();
  while(array = static_cast<TObjArray*>(fItInputArray->Next()))
  {
    TIter itArray(array);
    name = array->GetName();
    entry = 1;

    while(particle = static_cast<ExRootLHEFParticle*>(itArray.Next()))
    {
      momentum.SetPxPyPzE(particle->Px, particle->Py, particle->Pz, particle->E);

      candidate = factory->NewCandidate();

      candidate->SetP4(momentum);
      candidate->SetName(Form("%s_{%d}", name.Data(), entry ));
      if(particle->Status != 1) candidate->SetResonance(kTRUE);

      fOutputArray->Add(candidate);
      ++entry;
    }
  }
}