#include "modules/MadGraphSISConeJetFinder.h"

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

#include "SISCone/momentum.h"
#include "SISCone/siscone.h"

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

#include <iostream>
#include <vector>

using namespace std;
using namespace siscone;

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

MadGraphSISConeJetFinder::MadGraphSISConeJetFinder() :
  fJetAlgo(0), fItInputArray(0)
{

}

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

MadGraphSISConeJetFinder::~MadGraphSISConeJetFinder()
{

}

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

void MadGraphSISConeJetFinder::Init()
{
  
  // define MidPoint algorithm

  fConeRadius       = GetDouble("ConeRadius", 0.5);
  fConeAreaFraction = GetDouble("ConeAreaFraction", 0.5);
  fMaxIterations    = GetInt("MaxIterations", 100);
  fJetPTMin         = GetDouble("JetPTMin", 20.0);

  fJetAlgo = new Csiscone;

  // import array with output from filter/classifier module

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

  // create output arrays

  fOutputArray = ExportArray("candidates");

}

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

void MadGraphSISConeJetFinder::Finish()
{
  if(fJetAlgo) delete fJetAlgo;
  if(fItInputArray) delete fItInputArray;
}

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

void MadGraphSISConeJetFinder::Process()
{
  ExRootCandidate *candidate;
  TLorentzVector momentum;
  Cmomentum jetMomentum;
  Int_t entry;

  ExRootFactory *factory = GetFactory();

  fParticlesList.clear();

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

  // construct jets from a list of stable particles
  fJetAlgo->compute_jets(fParticlesList, fConeRadius, fConeAreaFraction, fMaxIterations, fJetPTMin);

  // loop over all jets and export them
  vector<Cjet>::iterator itJet;
  for(itJet = fJetAlgo->jets.begin(), entry = 1; itJet != fJetAlgo->jets.end(); ++itJet, ++entry)
  {
    jetMomentum = itJet->v;

    momentum.SetPxPyPzE(jetMomentum.px, jetMomentum.py, jetMomentum.pz, jetMomentum.E);

    candidate = factory->NewCandidate();

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

    fOutputArray->Add(candidate);
  }
}