#include <iostream>
#include <vector>
#include <cmath>
#include <string>
#include <TH2.h>
#include <TStyle.h>
#include <TCanvas.h>
#include <TNtuple.h>
#include <TTree.h>
#include <TClonesArray.h>
#include <TLorentzVector.h>
#include <TFile.h>
#include <TGraph.h>
#include <TLegend.h>

void recon()
{
  TFile fin1("ttbar-dil.root", "read");

  TTree *fEvents1 = (TTree*) fin1.Get("EventCollection");
  TClonesArray *jetsLV1 = new TClonesArray("TLorentzVector");
  TClonesArray *lepLV1 = new TClonesArray("TLorentzVector");
  // TClonesArray *MCmetLV1 = new TClonesArray("TLorentzVector");
  // TClonesArray *MCjetsLV1 = new TClonesArray("TLorentzVector");
  // TClonesArray *MClepLV1 = new TClonesArray("TLorentzVector");

  fEvents1->SetBranchAddress("Jets", &jetsLV1);
  fEvents1->SetBranchAddress("Leptons", &lepLV1);
  // fEvents1->SetBranchAddress("MET", &MCmetLV1);
  // fEvents1->SetBranchAddress("TruthJets", &MCjetsLV1);
  // fEvents1->SetBranchAddress("TruthLep", &MClepLV1);

  TH1F *ht1 = new TH1F("ht1", "H_{t}", 100, 100, 3000);
  // TH1F *MCht1 = new TH1F("MCht1", "MC H_{t}", 100, 100, 3000);
  TH1F *M1 = new TH1F("M1", "Invariant Mass", 100, 100, 3000);
  // TH1F *MCM1 = new TH1F("MCM1", "Invariant Mass", 100, 100, 3000);

  //------------------------------
  //-- Reconstruct events
  //------------------------------

  Int_t nentries = (Int_t) fEvents1->GetEntriesFast();
  std::cout << nentries << std::endl;
  for (Int_t ev = 0; ev < nentries; ++ev)
    {
      jetsLV1->Clear();
      lepLV1->Clear();
      // MCmetLV1->Clear();
      // MCjetsLV1->Clear();
      // MClepLV1->Clear();

      Int_t njets = jetsLV1->GetEntriesFast();
      Int_t nlep = lepLV1->GetEntriesFast();
      // Int_t nMCmet = MCmetLV1->GetEntriesFast();
      // Int_t nMCjets = MCjetsLV1->GetEntriesFast();
      // Int_t nMClep = MClepLV1->GetEntriesFast();

      Float_t HT = 0.0;
      // Float_t MCHT = 0.0;
      TLorentzVector tt(0,0,0,0);
      // TLorentzVector MCtt(0,0,0,0);
      TLorentzVector *lv = NULL;

      fEvents1->GetEntry(ev);

      // adding the 4-vectors
      for (Int_t i = 0; i < njets; ++i)
	{
	  lv = (TLorentzVector*) jetsLV1->At(i);
	  tt += *lv;
	  HT += lv->Pt();
	}
      for (Int_t i = 0; i < nlep; ++i)
	{
	  lv = (TLorentzVector*) lepLV1->At(i);
	  tt += *lv;
	  HT += lv->Pt();
	}
      // for (Int_t i = 0; i < nMCmet; ++i)
      // 	{
      // 	  lv = (TLorentzVector*) MCmetLV1->At(i);
      // 	  tt += *lv;
      // 	  HT += lv->Pt();
      // 	  MCtt += *lv;
      // 	  MCHT += lv->Pt();
      // 	}
      // for (Int_t i = 0; i < nMCjets; ++i)
      // 	{
      // 	  lv = (TLorentzVector*) MCjetsLV1->At(i);
      // 	  MCtt += *lv;
      // 	  MCHT += lv->Pt();
      // 	}
      // for (Int_t i = 0; i < nMClep; ++i)
      // 	{
      // 	  lv = (TLorentzVector*) MClepLV1->At(i);
      // 	  MCtt += *lv;
      // 	  MCHT += lv->Pt();
      // 	}

      ht1->Fill(HT/1000.);
      // MCht1->Fill(MCHT/1000.);
      M1->Fill(tt.M()/1000.);
      // MCM1->Fill(MCtt.M()/1000.);
    }

  TCanvas *distrib = new TCanvas("distrib","Invariant mass & HT distributions", 1000, 400);
  TPad *ht_1 = new TPad("ht_1","SM",0.02,0.02,0.48,0.98);
  TPad *M_1 = new TPad("M_1","SM",0.52,0.02,0.98,0.98);

  ht_1->Draw();
  M_1->Draw();

  ht_1->cd();
  ht1->Draw();
  // MCht1->Draw();

  M_1->cd();
  M1->Draw();
  // MCM1->Draw();

  distrib->cd();
  distrib->Update();
}