/// C/C++ includes
#include <iostream>
#include <cstring>

/// ROOT includes
#include "TROOT.h"
#include "TObject.h"
#include "TFile.h"
#include "TKey.h"
#include "TDirectoryFile.h"
// histograms
#include "TNtuple.h"
#include "TH1F.h"
#include "TStyle.h"
#include "TCanvas.h"
//#include "TPad.h"
#include "TString.h"
#include "TPaveStats.h"
#include "TPaveText.h"
#include "TGaxis.h"

void read_dir(TDirectory *dir, TCanvas *canvas);
bool plot(TKey *key, TCanvas *canvas);


int overlay()
{
  TFile *finv = new TFile("~/codebaby/root-files/bbar_inv_iso.root");
  if (finv->IsZombie())
    {
      printf("File does not exist.\n");
      return 0;
    }

  TFile *freg = new TFile("~/codebaby/root-files/bbar_reg_iso.root");
  if (freg->IsZombie())
    {
      printf("File does not exist.\n");
      return 0;
    }

  TCanvas *METOverlay = new TCanvas("METOverlay","Missing Et comparirson b/w b#bar{b} in signal and background region", 800, 500);
  TCanvas *tmp_canvas = new TCanvas("tmp_canvas","Missing Et comparirson b/w b#bar{b} in signal and background region", 800, 500);

  /// Read and plot TNtuple
  // Using these commented out lines instead of the call to read_dir works without problems
  // TNtuple *Wprops_rev = (TNtuple*)(finv->Get("WPrime_0_W_properties_bbmu15X")); 
  // tmp_canvas->cd();
  // Wprops_rev->Draw("MET>>tmp_histo(100,0,100)");

  read_dir(finv, tmp_canvas); /// <-- troubling line
  TH1F* MET_rev_iso = (TH1F*)gPad->GetPrimitive("tmp_histo");

  METOverlay->cd();
  MET_rev_iso->SetNameTitle("MET_rev_iso","Missing Et for b#bar{b} background");
  MET_rev_iso->GetXaxis()->SetRangeUser(0,100);
  MET_rev_iso->SetFillColor(kGreen-7);
  MET_rev_iso->Draw("HIST");
  Double_t rev_integral = MET_rev_iso->Integral();

  gPad->Update();
  TPaveStats *statbox_rev = (TPaveStats*) MET_rev_iso->FindObject("stats");
  statbox_rev->SetY1NDC(0.8);
  statbox_rev->SetY2NDC(0.95);
  statbox_rev->SetFillColor(MET_rev_iso->GetFillColor());
  gPad->Modified();
  METOverlay->Update();

  // TNtuple *Wprops_reg = (TNtuple*)(freg->Get("WPrime_0_W_properties_bbmu15X"));

  // tmp_canvas->cd();
  // Wprops_reg->Draw("MET>>tmp_histo(100,0,100)");

  /// Regular selection
  read_dir(freg, tmp_canvas);
  TH1F* MET_reg_iso = (TH1F*)gPad->GetPrimitive("tmp_histo");

  METOverlay->cd();
  MET_reg_iso->SetNameTitle("MET_reg_iso","Missing Et for b#bar{b} background in signal region");
  MET_reg_iso->GetXaxis()->SetRangeUser(0,100);
  Double_t reg_integral = MET_reg_iso->Integral();
  MET_reg_iso->Scale(rev_integral/reg_integral);
  MET_reg_iso->SetMarkerColor(kBlack);
  MET_reg_iso->Draw("E1 SAMES");
 
  gPad->Update();
  TPaveStats *statbox_reg = (TPaveStats*) MET_reg_iso->FindObject("stats");
  statbox_reg->SetY1NDC(0.5);
  statbox_reg->SetY2NDC(0.65);
  gPad->Modified();
  METOverlay->Update();

  delete tmp_canvas; //METOverlay;

  std::cout << "Done plotting histograms!" << std::endl;
  return 0;
}


void read_dir(TDirectory *dir, TCanvas *canvas)
{
  TDirectory *dirsav = gDirectory;
  TIter next(dir->GetListOfKeys());
  TKey *key;

  while ((key = (TKey*)next()))
    {
      TClass *objClass = TClass::GetClass(key->GetClassName());

      if (objClass && objClass->InheritsFrom(TDirectory::Class()))
	{
	  dir->cd(key->GetName());
	  TDirectory *subdir = gDirectory;
	  read_dir(subdir, canvas);
	  dirsav->cd();
	  continue;
	}

      if (strstr(key->GetName(),"bbmu") && objClass->InheritsFrom(TNtuple::Class()))
      	{
	  plot(key, canvas); /// <-- troubling line
	}
    }
}


bool plot(TKey *key, TCanvas *canvas)
{
  TNtuple *Wprops;

  if (strstr(key->GetName(),"properties"))
    {
      canvas->cd();
      // Wprops = (TNtuple*)(key->ReadObj());
      // Wprops->Draw("MET>>tmp_histo(100,0,100)");
      ((TNtuple*)(key->ReadObj()))->Draw("MET>>tmp_histo(100,0,100)"); /// <-- troubling line
      return true;
    }

  return false;
}