#include "modules/MadGraphAnalysis.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;

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

MadGraphAnalysis::MadGraphAnalysis()
{
}

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

MadGraphAnalysis::~MadGraphAnalysis()
{
}

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

void MadGraphAnalysis::Init()
{
  fOutputFileName = GetString("OutputFile", "pythia_plots.root");

  // import array with output from filter/classifier module

  fInputArray = ImportArray(GetString("InputArray", "merger/candidates"));

  fIsUnWeighted = GetBool("IsUnWeighted", kFALSE);

  // import ROOT tree branch
  
  if(fIsUnWeighted)
  {
    fBranchEvent = UseBranch("Event");
  }
  else
  {
    fBranchEvent = 0;
  }

}

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

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

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

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

void MadGraphAnalysis::Process()
{
  ExRootCandidate *candidate1 = 0, *candidate2 = 0;
  ParticleHistograms *histogramsParticle = 0;
  PairHistograms *histogramsPair = 0;
  Int_t maxEntry, entry1, entry2;

  Double_t weight = 1.0;
  Double_t pt1, pt2, dr, rapidity, signPz;

  ExRootLHEFEvent *eventInfo = 0;
  if(fIsUnWeighted && fBranchEvent && fBranchEvent->GetEntriesFast() == 1)
  {
    eventInfo = static_cast<ExRootLHEFEvent*>(fBranchEvent->At(0));

    weight = eventInfo->Weight;
  }

  // fill histograms
  maxEntry = fInputArray->GetEntriesFast();
  for(entry1 = 0; entry1 < maxEntry; ++entry1)
  {
    candidate1 = static_cast<ExRootCandidate*>(fInputArray->At(entry1));

    const TLorentzVector &vector1 = candidate1->GetP4();

    pt1 = vector1.Pt();
    signPz = (vector1.Pz() >= 0.0) ? 1.0 : -1.0;
    rapidity = (pt1 == 0.0 ? signPz*999.9 : vector1.Rapidity());

    // fill histograms for single particles

    histogramsParticle = GetParticleHistograms(candidate1->GetName());

    histogramsParticle->fParticlePt->Fill(pt1, weight);

    histogramsParticle->fParticleRapidity->Fill(rapidity, weight);
    
    // skip pairs of resonances
    if(candidate1->IsResonance()) continue;

    for(entry2 = entry1 + 1; entry2 < maxEntry; ++entry2)
    {
      candidate2 = static_cast<ExRootCandidate*>(fInputArray->At(entry2));

      // skip pairs of resonances
      if(candidate2->IsResonance()) continue;

      const TLorentzVector &vector2 = candidate2->GetP4();

      pt2 = vector2.Pt();
      dr = (pt1 == 0.0 || pt2 == 0.0 ? 999.9 : vector1.DeltaR(vector2));

      // fill histograms for pairs of particles

      histogramsPair = GetPairHistograms(candidate1->GetName(), candidate2->GetName());

      histogramsPair->fPairDeltaR->Fill(dr, weight);
      histogramsPair->fPairMass->Fill((vector1 + vector2).M(), weight);
    }
  }
}

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

void MadGraphAnalysis::BookParticleHistograms(MadGraphAnalysis::ParticleHistograms *histograms,
                                                    const char *name, const char *title)
{
  ExRootResult *result = GetPlots();
  histograms->fParticlePt = result->AddHist1D(Form("pt_%s", name),
                                              Form("P_{T}(%s)", title),
                                              Form("P_{T}(%s), GeV/c", title),
                                              "pb/bin",
                                              60, 0.0, 300.0, 0, 1);
  histograms->fParticlePt->SetStats(kTRUE);

  histograms->fParticleRapidity = result->AddHist1D(Form("y_%s", name),
                                                    Form("y(%s)", title),
                                                    Form("y(%s)", title),
                                                    "pb/bin",
                                                    100, -5.0, 5.0, 0, 0);
  histograms->fParticleRapidity->SetStats(kTRUE);

}

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

void MadGraphAnalysis::BookPairHistograms(MadGraphAnalysis::PairHistograms *histograms,
                                                const char *name, const char *title)
{
  ExRootResult *result = GetPlots();
  histograms->fPairDeltaR = result->AddHist1D(Form("dr_%s", name),
                                              Form("#DeltaR(%s)", title),
                                              Form("#DeltaR(%s)", title),
                                              "pb/bin",
                                              70, 0.0, 7.0, 0, 1);
  histograms->fPairDeltaR->SetStats(kTRUE);

  histograms->fPairMass = result->AddHist1D(Form("mass_%s", name),
                                            Form("M_{inv}(%s)", title),
                                            Form("M_{inv}(%s), GeV/c^{2}", title),
                                            "pb/bin",
                                            120, 0.0, 600.0, 0, 1);
  histograms->fPairMass->SetStats(kTRUE);

}

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

MadGraphAnalysis::ParticleHistograms *
MadGraphAnalysis::GetParticleHistograms(const char *candName)
{
  map<TString, ParticleHistograms *>::iterator itParticleHistogramsMap;
  ParticleHistograms *histograms = 0;
  TString name = Form("%s", candName);
  name.ReplaceAll("{", "");
  name.ReplaceAll("}", "");
  name.ReplaceAll("^", "");
  name.ReplaceAll("#bar", "anti_");
  name.ReplaceAll("#", "");
  TString title = Form("%s", candName);
  itParticleHistogramsMap = fParticleHistogramsMap.find(name);
  if(itParticleHistogramsMap == fParticleHistogramsMap.end())
  {
    histograms = new ParticleHistograms;

    BookParticleHistograms(histograms, name, title);

    fParticleHistogramsMap[name] = histograms;
  }
  else
  {
    histograms = itParticleHistogramsMap->second;
  }
  return histograms;
}

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

MadGraphAnalysis::PairHistograms *
MadGraphAnalysis::GetPairHistograms(const char *candName1, const char *candName2)
{
  map<TString, PairHistograms *>::iterator itPairHistogramsMap;
  PairHistograms *histograms = 0;
  TString name = Form("%s_%s", candName1, candName2);
  name.ReplaceAll("{", "");
  name.ReplaceAll("}", "");
  name.ReplaceAll("^", "");
  name.ReplaceAll("#bar", "anti_");
  name.ReplaceAll("#", "");
  TString title = Form("%s, %s", candName1, candName2);
  itPairHistogramsMap = fPairHistogramsMap.find(name);
  if(itPairHistogramsMap == fPairHistogramsMap.end())
  {
    histograms = new PairHistograms;

    BookPairHistograms(histograms, name, title);

    fPairHistogramsMap[name] = histograms;
  }
  else
  {
    histograms = itPairHistogramsMap->second;
  }
  return histograms;
}