#include "modules/MadGraphCDFConeJetFinder.h"

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

#include "CDFCones/JetCluAlgorithm.hh"
#include "CDFCones/MidPointAlgorithm.hh"

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

#include <iostream>
#include <vector>

using namespace std;

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

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

}

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

MadGraphCDFConeJetFinder::~MadGraphCDFConeJetFinder()
{

}

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

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

  double seedThreshold    = GetDouble("SeedThreshold", 1.0);
  double coneRadius       = GetDouble("ConeRadius", 0.5);
  double coneAreaFraction = GetDouble("ConeAreaFraction", 1.0);
  int    maxPairSize      = GetInt("MaxPairSize", 2);
  int    maxIterations    = GetInt("MaxIterations", 100);
  double overlapThreshold = GetDouble("OverlapThreshold", 0.75);

  fJetAlgo = new MidPointAlgorithm(seedThreshold, coneRadius, coneAreaFraction,
                                   maxPairSize, maxIterations, overlapThreshold);

  // import array with output from filter/classifier module

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

  // create output arrays

  fOutputArray = ExportArray("candidates");

}

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

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

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

void MadGraphCDFConeJetFinder::Process()
{
  ExRootCandidate *candidate;
  TLorentzVector momentum;
  LorentzVector jetMomentum;
  Int_t entry;

  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(PhysicsTower(LorentzVector(momentum.Px(), momentum.Py(),
                                                      momentum.Pz(), momentum.E())));
  }

  // construct jets from a list of stable particles
  fJetsList.clear();
  fJetAlgo->run(fTowersList, fJetsList);

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

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