#include "modules/MadGraphKtJetFinder.h"

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

#include "KtJet/KtEvent.h"
#include "KtJet/KtLorentzVector.h"

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

#include <iostream>
#include <vector>

using namespace std;
using namespace KtJet;

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

MadGraphKtJetFinder::MadGraphKtJetFinder() :
  fItInputArray(0)
{

}

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

MadGraphKtJetFinder::~MadGraphKtJetFinder()
{

}

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

void MadGraphKtJetFinder::Init()
{

  // define KtJet algorithm

  fCollisionType = GetInt("CollisionType", 4); // PP
  fDistanceScheme = GetInt("DistanceScheme", 1); // Angular
  fRecombinationScheme = GetInt("RecombinationScheme", 1); // E
  fParameterR = GetDouble("ParameterR", 1.0);

  fIsExclusive = GetBool("Exclusive", false);
  fECut = GetDouble("ECut", -1.0);
  fDCut = GetDouble("DCut", 900.0);

  fParticleNumberMin = GetInt("ParticleNumberMin", 1);

  // import array with output from filter/classifier module

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

  // create output arrays

  fOutputArrayCandidates = ExportArray("candidates");
  fOutputArrayMatching = ExportArray("matching");

}

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

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

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

void MadGraphKtJetFinder::Process()
{
  ExRootCandidate *candidate;
  ExRootMatching* matching;
  TLorentzVector momentum;
  KtEvent *event = 0;
  Int_t entry;

  if(fInputArray->GetEntriesFast() < fParticleNumberMin) return;

  ExRootFactory *factory = GetFactory();

  fTowersList.clear();

  // loop over all particles in event and select stable ones
  fItInputArray->Reset();
  while((candidate = static_cast<ExRootCandidate*>(fItInputArray->Next())))
  {
     momentum = candidate->GetP4();
     fTowersList.push_back(KtLorentzVector(momentum.Px(), momentum.Py(),
                                           momentum.Pz(), momentum.E()));
  }

  // construct jets from a list of stable particles
  if(fIsExclusive)
  {
    event = new KtEvent(fTowersList, fCollisionType, fDistanceScheme,
                        fRecombinationScheme);

    if(fECut > 0.0) event->setECut(fECut);
    event->findJetsY(fDCut);
  }
  else
  {
    event = new KtEvent(fTowersList, fCollisionType, fDistanceScheme,
                        fRecombinationScheme, fParameterR);

  }

  fJetsList.clear();
  fJetsList = event->getJetsPt();

  // loop over all jets and export them
  vector<KtLorentzVector>::iterator itJet;
  for(itJet = fJetsList.begin(), entry = 1; itJet != fJetsList.end(); ++itJet, ++entry)
  {
    momentum.SetPxPyPzE(itJet->px(), itJet->py(), itJet->pz(), itJet->e());

    candidate = factory->NewCandidate();

    candidate->SetP4(momentum);
    candidate->SetName(Form("jet_{%d}", entry ));

    fOutputArrayCandidates->Add(candidate);
  }

  for(entry = 0; entry < fTowersList.size(); ++entry)
  {
    matching = factory->New<ExRootMatching>();

    matching->DMerge = event->getDMerge(entry);
    matching->YMerge = event->getYMerge(entry);

    fOutputArrayMatching->Add(matching);
  }

  if(event) delete event;
}