#include "TFile.h"
#include "TTree.h"
#include "TCanvas.h"
#include "TH1F.h"
#include <iostream>

using namespace std;

void readEvents(){
  
  //TASK 1
  //Load the file
  TFile * f = new TFile("muons.root");

  //Get the tree(called POOLCollectionTree) from the file
  TTree *tree = (TTree *) f->Get("POOLCollectionTree");

  //Get entries of the file
  int nEntries = tree->GetEntries();

  //Print the number of entries the file found
  cout << "There are " << nEntries << " entries in your ntuple" << endl;
  

  //TASK 2
  // create local variables for the tree's branches
  // We need Float_t type for LooseMuonsPhi1, LooseMuonsPt1, LooseMuonsEta2, LooseMuonsPhi2,LooseMuonsPt2
  UInt_t NLooseMuons;
  Float_t LooseMuonsEta1;
  Float_t LooseMuonsPhi1;
  Float_t LooseMuonsPt1;

  Float_t LooseMuonsEta2;
  Float_t LooseMuonsPhi2;
  Float_t LooseMuonsPt2;

  //Task 3  
  // set the tree's braches to the local variables
  tree->SetBranchAddress("NLooseMuon", &NLooseMuons);
  tree->SetBranchAddress("LooseMuonEta1", &LooseMuonsEta1);
  tree->SetBranchAddress("LooseMuonPhi1", &LooseMuonsPhi1);
  tree->SetBranchAddress("LooseMuonPt1", &LooseMuonsPt1);
  
  tree->SetBranchAddress("LooseMuonEta2", &LooseMuonsEta2);
  tree->SetBranchAddress("LooseMuonPhi2", &LooseMuonsPhi2);
  tree->SetBranchAddress("LooseMuonPt2", &LooseMuonsPt2);
  
  //Task 4
  //declare some histograms
  //declare a histogram called muPt, with 50 bins, and range from 0 to 200
  TH1F *muPt = new TH1F("muPt", ";p_{T} [GeV/c];Events", 50, 0, 200);

  //declare a histogram called muEta, with 50 bins and range -3 to 3
  TH1F *muEta = new TH1F("muEta", ";#eta;Events", 50, -3, 3);

  //declare a histogram called muPhi with 50 bins and range from -4 to 4
  TH1F *muPhi = new TH1F("muPhi", ";#phi;Events", 50, -4, 4);
 
  //declare a histogram called muE with 50 bins and range from 0 to 200
  TH1F *muE = new TH1F("muE", ";Energy;Events", 50, 0, 200);
  
  //declare a histogram called Zmass with 50 bins and range from 0 to 200
  TH1F *Zmass = new TH1F("Zmass", ";Z Mass;Events", 50, 0, 200);

  //declare a histogram called ZmassBkg with 50 bins and range from 0 to 200
  TH1F *ZmassBkg = new TH1F("ZmassBkg", ";Z Mass;Events", 50, 0, 200);
  

  // loop over each entry (event) in the tree
  for( int entry=0; entry<nEntries; entry++ ){
    //Task A
    if( entry%10000 == 0 ) cout << "Entry:" << entry << endl;
    
    // check that the event is read properly
    int entryCheck = tree->GetEntry( entry );
    if( entryCheck <= 0 ){  continue; }


    //Task B    
    // only look at events containing at least 2 leptons
    if(NLooseMuons < 2) continue;
    
    // require the leptons to be greater than 20 GeV
    if(abs(LooseMuonsPt1) *0.001 < 20 ) continue;
    if(abs(LooseMuonsPt1) *0.001 < 20 ) continue;
    
    // make a LorentzVector from the muon
    TLorentzVector Muons1;
    Muons1.SetPtEtaPhiM(fabs(LooseMuonsPt1), LooseMuonsEta1, LooseMuonsPhi1, 0);
    
    TLorentzVector Muons2;
    Muons2.SetPtEtaPhiM(fabs(LooseMuonsPt2), LooseMuonsEta2, LooseMuonsPhi2, 0);
    
    // print out the details of an electron every so often
    if( entry%100000 == 0 ){ 
      cout << "Muons pt: " << LooseMuonsPt1 << " eta: " << LooseMuonsEta1 << " phi " << LooseMuonsPhi1 << endl;
    }
    
    //Task C
    // fill our histograms
    muPt->Fill(Muons1.Pt()*0.001);
    muEta->Fill(Muons1.Eta());
    muPhi->Fill(Muons1.Phi());
    muE->Fill(Muons1.E()*0.001);
   
    //Define the Z boson vector as a sum of the 2 muon vectors
    TLorentzVector Z = Muons1 + Muons2;     

    //Fill the Zmass and ZmassBkg histograms 
    if((LooseMuonsPt1 *LooseMuonsPt2) < 0 ){
      Zmass->Fill(Z.M()*0.001);
    }
    else{
      ZmassBkg->Fill(Z.M()*0.001);
    }
    
  }//THIS IS THE END OF OUR ENTRIES LOOP
  

  //Task 5
  // draw the Zmass distribution
  Zmass->Draw();
  
  // make a ROOT output file to store your histograms
  TFile *outFile = new TFile("histograms.root", "recreate");
  muPt->Write();
  muEta->Write();
  muPhi->Write();
  muE->Write();
  Zmass->Write();
  ZmassBkg->Write();

}// END OF readEventsLoop