#include "modules/MadGraphMatchingTreeWriter.h"


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

#include "ExRootAnalysis/ExRootCandidate.h"

#include "TClonesArray.h"

#include "TH1.h"
#include "TH2.h"
#include "TString.h"
#include "TCanvas.h"
#include "TLorentzVector.h"

#include <iostream>

using namespace std;

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

MadGraphMatchingTreeWriter::MadGraphMatchingTreeWriter()
{
}

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

MadGraphMatchingTreeWriter::~MadGraphMatchingTreeWriter()
{
}

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

void MadGraphMatchingTreeWriter::Init()
{
  fJetPTMin = GetDouble("JetPTMin", 20.0);
  fJetEtaMax = GetDouble("JetEtaMax", 4.5);

  // import array with output from filter/classifier/jetfinder modules

  fInputArrayPartonJets = ImportArray(GetString("InputArrayPartonJets", "partonjetfinder/candidates"));
  fItInputArrayPartonJets = fInputArrayPartonJets->MakeIterator();

  fInputArrayHadronJets = ImportArray(GetString("InputArrayHadronJets", "hadronjetfinder/candidates"));
  fItInputArrayHadronJets = fInputArrayHadronJets->MakeIterator();

  fInputArrayMatching = ImportArray(GetString("InputArrayMatching", "partonjetfinder/matching"));
  fItInputArrayMatching = fInputArrayMatching->MakeIterator();

  fInputArrayPartons = ImportArray(GetString("InputArrayPartons", "initstateselection/candidates"));
  fItInputArrayPartons = fInputArrayPartons->MakeIterator();

  fBranchPartonJets = NewBranch("PartonJet", ExRootGenJet::Class());
  fBranchHadronJets = NewBranch("HadronJet", ExRootGenJet::Class());
  fBranchMatching = NewBranch("Match", ExRootMatching::Class());
  fBranchPartons = NewBranch("Parton", ExRootGenParticle::Class());
}

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

void MadGraphMatchingTreeWriter::Finish()
{
  if(fItInputArrayPartonJets) delete fItInputArrayPartonJets;
  if(fItInputArrayHadronJets) delete fItInputArrayHadronJets;
  if(fItInputArrayMatching) delete fItInputArrayMatching;
  if(fItInputArrayPartons) delete fItInputArrayPartons;
}

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

void MadGraphMatchingTreeWriter::Process()
{
  ExRootCandidate *candidate = 0;
  ExRootGenParticle *entryParton = 0;
  ExRootMatching *matching = 0, *entryMatching = 0;
  ExRootGenJet *entryJet = 0;
  Double_t pt, signPz, eta, rapidity;

  // loop over all parton jets
  fItInputArrayPartonJets->Reset();
  while((candidate = static_cast<ExRootCandidate*>(fItInputArrayPartonJets->Next())))
  {
    const TLorentzVector &momentum = candidate->GetP4();

    pt = momentum.Pt();
    signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
    eta = (pt == 0.0 ? signPz*999.9 : momentum.Eta());
    rapidity = (pt == 0.0 ? signPz*999.9 : momentum.Rapidity());

    if(pt < fJetPTMin) continue;
    if(TMath::Abs(eta) > fJetEtaMax) continue;

    entryJet = static_cast<ExRootGenJet*>(fBranchPartonJets->NewEntry());

    entryJet->E = momentum.E();
    entryJet->Px = momentum.Px();
    entryJet->Py = momentum.Py();
    entryJet->Pz = momentum.Pz();

    entryJet->Eta = eta;
    entryJet->Phi = momentum.Phi();
    entryJet->PT = pt;

    entryJet->Rapidity = rapidity;

    entryJet->Mass = momentum.M();
  }

  // loop over all hadron jets
  fItInputArrayHadronJets->Reset();
  while((candidate = static_cast<ExRootCandidate*>(fItInputArrayHadronJets->Next())))
  {
    const TLorentzVector &momentum = candidate->GetP4();

    pt = momentum.Pt();
    signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
    eta = (pt == 0.0 ? signPz*999.9 : momentum.Eta());
    rapidity = (pt == 0.0 ? signPz*999.9 : momentum.Rapidity());

    if(pt < fJetPTMin) continue;
    if(TMath::Abs(eta) > fJetEtaMax) continue;

    entryJet = static_cast<ExRootGenJet*>(fBranchHadronJets->NewEntry());

    entryJet->E = momentum.E();
    entryJet->Px = momentum.Px();
    entryJet->Py = momentum.Py();
    entryJet->Pz = momentum.Pz();

    entryJet->Eta = eta;
    entryJet->Phi = momentum.Phi();
    entryJet->PT = pt;

    entryJet->Rapidity = rapidity;

    entryJet->Mass = momentum.M();
  }

  // loop over all matching
  fItInputArrayMatching->Reset();
  while((matching = static_cast<ExRootMatching*>(fItInputArrayMatching->Next())))
  {
    entryMatching = static_cast<ExRootMatching*>(fBranchMatching->NewEntry());

    entryMatching->DMerge = matching->DMerge;
    entryMatching->YMerge = matching->YMerge;
  }

  // loop over all partons
  fItInputArrayPartons->Reset();
  while((candidate = static_cast<ExRootCandidate*>(fItInputArrayPartons->Next())))
  {
    const TLorentzVector &momentum = candidate->GetP4();

    entryParton = static_cast<ExRootGenParticle*>(fBranchPartons->NewEntry());

    pt = momentum.Pt();
    signPz = (momentum.Pz() >= 0.0) ? 1.0 : -1.0;
    eta = (pt == 0.0 ? signPz*999.9 : momentum.Eta());
    rapidity = (pt == 0.0 ? signPz*999.9 : momentum.Rapidity());

    entryParton->PID = candidate->GetType()->PdgCode();

    entryParton->E = momentum.E();
    entryParton->Px = momentum.Px();
    entryParton->Py = momentum.Py();
    entryParton->Pz = momentum.Pz();

    entryParton->Eta = eta;
    entryParton->Phi = momentum.Phi();
    entryParton->PT = pt;

    entryParton->Rapidity = rapidity;
  }
}

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