#include <iostream>
#include <string>
#include <stdio.h>

%jsroot on

TChain *dataset = new TChain("mini");

dataset->Add("mc_105986.ZZ.root");
//dataset->Add("http://opendata.atlas.cern/release/samples/MC/mc_147770.Zee.root");

const int vs = 5;

Int_t   lepton_n = -1,
        lepton_charge[vs], //electrical charge of the lepton
        lepton_type[vs];   //flavour of the lepton

Float_t lepton_pt[vs],
        lepton_eta[vs],
        lepton_phi[vs],
        lepton_E[vs];

dataset->SetBranchAddress("lep_n",      &lepton_n);
dataset->SetBranchAddress("lep_charge", &lepton_charge);
dataset->SetBranchAddress("lep_type",   &lepton_type);
dataset->SetBranchAddress("lep_pt",     &lepton_pt);
dataset->SetBranchAddress("lep_eta",    &lepton_eta);
dataset->SetBranchAddress("lep_phi",    &lepton_phi);
dataset->SetBranchAddress("lep_E",      &lepton_E);

TH1F *hist = new TH1F("variable","Mass of the Z boson; mass [GeV]; events",30,40,140);
hist->SetFillColor(kBlue);

int nentries, nbytes, i;
nentries = (Int_t)dataset->GetEntries();

// IMPORTANT: fraction events we want to run
fraction_events = 1;
events_to_run = nentries*fraction_events;

std::cout << "Total # events = "  << nentries
          << ". Events to run = " << events_to_run
          << " corresponding to " << fraction_events*100
          << "% of total events!" << std::endl;

TLorentzVector l1;         
                TLorentzVector l2;  
                               TLorentzVector invmass;  

for (i = 0; i < nentries; i++)
{
    nbytes = dataset->GetEntry(i);
    
    // Cut #1: At least 2 leptons
    if(lepton_n == 2)
    {
        // Cut #2: Leptons with opposite charge
        if(lepton_charge[0] != lepton_charge[1])
        {
            // Cut #3: Leptons of the same family (2 electrons or 2 muons)
            if(lepton_type[0] == lepton_type[1])
            {
                l1.SetPtEtaPhiE(lepton_pt[0]/1000., lepton_eta[0], lepton_phi[0], lepton_E[0]/1000.);
                l2.SetPtEtaPhiE(lepton_pt[1]/1000., lepton_eta[1], lepton_phi[1], lepton_E[1]/1000.);
                // The default value in the root file is in MeV, so, we divide by 1000 to get it in GeV
                invmass = l1 + l2;
                hist->Fill(invmass.M());
            }
        }
    }
}

TCanvas *canvas = new TCanvas("c","c",10,10,700,700);
hist->Draw();
canvas->Draw();