#include "modules/MadGraphMatchingAnalysis.h"


#include "ExRootAnalysis/ExRootResult.h"
#include "ExRootAnalysis/ExRootClasses.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;

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

MadGraphMatchingAnalysis::MadGraphMatchingAnalysis()
{
}

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

MadGraphMatchingAnalysis::~MadGraphMatchingAnalysis()
{
}

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

void MadGraphMatchingAnalysis::Init()
{
  fOutputFileName = GetString("OutputFile", "matching_plots.root");

  // import array with output from filter/classifier module

  fInputArrayCandidates = ImportArray(GetString("InputArrayCandidates", "jetfinder/candidates"));

  fInputArrayMatching = ImportArray(GetString("InputArrayMatching", "jetfinder/matching"));

}

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

void MadGraphMatchingAnalysis::Finish()
{
  GetPlots()->Write(fOutputFileName);

  GetPlots()->GetCanvas()->SetLogy(1);
  GetPlots()->Print();
  GetPlots()->GetCanvas()->SetLogy(0);
}

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

void MadGraphMatchingAnalysis::Process()
{
  ExRootCandidate *candidate = 0;
  ExRootMatching *matching = 0;
  MatchingHistograms *histograms = 0;
  Int_t maxEntry, entry;
  TString name;

  Double_t pt, signPz, rapidity;

  // loop over jets
  maxEntry = fInputArrayCandidates->GetEntriesFast();
  for(entry = 0; entry < maxEntry; ++entry)
  {
    candidate = static_cast<ExRootCandidate*>(fInputArrayCandidates->At(entry));

    const TLorentzVector &momentum = candidate->GetP4();

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

  // fill matching histograms
  maxEntry = fInputArrayMatching->GetEntriesFast();
  for(entry = 0; entry < maxEntry; ++entry)
  {
    matching = static_cast<ExRootMatching*>(fInputArrayMatching->At(entry));

    name = Form("q_{%d #rightarrow %d}", entry + 1, entry);
    histograms = GetMatchingHistograms(name);

    histograms->fYMerge->Fill(TMath::Log10(matching->DMerge));
  }

}

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

void MadGraphMatchingAnalysis::BookMatchingHistograms(MadGraphMatchingAnalysis::MatchingHistograms *histograms,
                                                      const char *name, const char *title)
{
  ExRootResult *result = GetPlots();
  histograms->fYMerge = result->AddHist1D(name, title, title,
                                          "entries", 60, 0.0, 10.0, 0, 1);
  histograms->fYMerge->SetStats(kTRUE);

}

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

MadGraphMatchingAnalysis::MatchingHistograms *
MadGraphMatchingAnalysis::GetMatchingHistograms(const char *name)
{
  map<TString, MatchingHistograms *>::iterator itMatchingHistogramsMap;
  MatchingHistograms *histograms = 0;
  TString newName = Form("%s", name);
  newName.ReplaceAll("{", "");
  newName.ReplaceAll("}", "");
  newName.ReplaceAll("^", "");
  newName.ReplaceAll("#bar", "anti_");
  newName.ReplaceAll("#rightarrow", "to");
  newName.ReplaceAll("#", "");
  newName.ReplaceAll(" ", "_");
  TString title = Form("%s", name);
  itMatchingHistogramsMap = fMatchingHistogramsMap.find(newName);
  if(itMatchingHistogramsMap == fMatchingHistogramsMap.end())
  {
    histograms = new MatchingHistograms;

    BookMatchingHistograms(histograms, newName, title);

    fMatchingHistogramsMap[newName] = histograms;
  }
  else
  {
    histograms = itMatchingHistogramsMap->second;
  }
  return histograms;
}

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