#include "TCanvas.h"
#include "TROOT.h"
#include "TFile.h"
#include "TTree.h"
#include "TBrowser.h"
#include "TH2.h"
#include "TRandom.h"

double fitf(double *v, double *par);

void tree1r()
{
  // Open Root file Selected_All_EEM.root
  TFile *f = new TFile("Selected_All_EEM.root");

  // Signal events    (tree named WZSignal)
  TTree *sig = (TTree*)f->Get("WZSignal");
  Double_t pt1, eta1, phi1, E1;
  Double_t pt2, eta2, phi2, E2;
  Double_t pt3, eta3, phi3, E3;
  Double_t MZ, MET, Trackd0cutWMu, TrackIsoWmu;
  Double_t Weight;

  //get some variables for SIGNAL EVENTS
  sig->SetBranchAddress("pt1",&pt1);
  sig->SetBranchAddress("eta1",&eta1);
  sig->SetBranchAddress("phi1",&phi1);
  sig->SetBranchAddress("E1",&E1);
    
    sig->SetBranchAddress("pt2",&pt2);
    sig->SetBranchAddress("eta2",&eta2);
    sig->SetBranchAddress("phi2",&phi2);
    sig->SetBranchAddress("E2",&E2);
    
    sig->SetBranchAddress("Trackd0cutWMu",&Trackd0cutWMu);
    sig->SetBranchAddress("TrackIsoWmu",&TrackIsoWmu);
    
  sig->SetBranchAddress("MZ",&MZ);
  sig->SetBranchAddress("Weight",&Weight); 
  // add other variables ... if you like


  ////get some variables for BACKGROUND EVENTS (Tree named ttbar)
  TTree *ttbar = (TTree*)f->Get("ttbar");
  Double_t pt1_bkg, eta1_bkg, phi1_bkg, E1_bkg;
    Double_t pt2_bkg, eta2_bkg, phi2_bkg, E2_bkg;
    Double_t MET_bkg, Trackd0cutWMu_bkg, TrackIsoWmu_bkg;

  Double_t MZ_bkg;
  Double_t Weight_bkg;

  //get some variables for ttbar
  ttbar->SetBranchAddress("pt1",&pt1_bkg);
  ttbar->SetBranchAddress("eta1",&eta1_bkg);
  ttbar->SetBranchAddress("phi1",&phi1_bkg);
  ttbar->SetBranchAddress("E1",&E1_bkg);
    
    ttbar->SetBranchAddress("pt2",&pt2_bkg);
    ttbar->SetBranchAddress("eta2",&eta2_bkg);
    ttbar->SetBranchAddress("phi2",&phi2_bkg);
    ttbar->SetBranchAddress("E2",&E2_bkg);
    ttbar->SetBranchAddress("Trackd0cutWMu",&Trackd0cutWMu_bkg);
    ttbar->SetBranchAddress("TrackIsoWmu",&TrackIsoWmu_bkg);
    
  ttbar->SetBranchAddress("MZ",&MZ_bkg);
  ttbar->SetBranchAddress("Weight",&Weight_bkg); 
  // add other variables ...    if you like

  //create requested histograms (for SIGNAL and ttbar)
    
  TH1F *h_pt1_SIG   = new TH1F("h_pt1_SIG","Pt most energetic Z lepton SIG",100,0,150);
  TH1F *h_pt1_BKG   = new TH1F("h_pt1_BKG","Pt most energetic Z lepton BKG",100,0,150);
    
  TH1F *h_d0_SIG   = new TH1F("h_d0_SIG","Impact parameter W lepton SIG",50,0,10);
   TH1F *h_d0_BKG   = new TH1F("h_d0_BKG","Impact parameter W lepton BKG",50,0,10);
    
    TH1F *h_iso_SIG   = new TH1F("h_iso_SIG","Isolation W lepton SIG",50,0,10);
    TH1F *h_iso_BKG   = new TH1F("h_iso_BKG","Isolation W lepton BKG",50,0,10);
  
    
  TH1F *h_MZ      = new TH1F("h_MZ","MZ distribution All events",40,65,115);
  TH1F *h_MZ_SIG  = new TH1F("h_MZ_SIG","MZ distribution SIG",40,65,115);
  TH1F *h_MZ_BKG  = new TH1F("h_MZ_BKG","MZ distribution BKG",40,65,115);

    TH1F *h_MZrec_SIG  = new TH1F("h_MZrec_SIG","MZreco distribution SIG",40,65,115);
    TH1F *h_MZrec_BKG  = new TH1F("h_MZrec_BKG","MZreco distribution BKG",40,65,115);
  
    TH1F *h_MZrec_ALL = new TH1F("h_MZrec_ALL","MZreco distribution ALL",40,65,115);
    TH1F *h_MZrec_CLEAN = new TH1F("h_MZrec_CLEAN","MZreco distribution CLEAN",40,65,115);
    
  //loop to read all SIGNAL entries and fill the histograms
  Int_t nentries = (Int_t)sig->GetEntries();
  
  for (Int_t i=0;i<nentries;i++) {
    sig->GetEntry(i);
      
      
    h_pt1_SIG->Fill(pt1,1);
    h_d0_SIG->Fill(Trackd0cutWMu,1);
    h_iso_SIG->Fill(TrackIsoWmu,1);
    
    // Compute Invariant mass
      double px1 = pt1*cos(phi1);
      double py1 = pt1*sin(phi1);
      double pz1 = E1*tanh(eta1);
      
      double px2 = pt2*cos(phi2);
      double py2 = pt2*sin(phi2);
      double pz2 = E2*tanh(eta2);
      
      double PZrec_SIG2 =  pow((px1+px2), 2)+ pow( (py1+py2), 2) + pow( (pz1+pz2), 2);
      double EZrec_SIG2 =  (E1+E2)*(E1+E2);
      double MZrec_SIG = 0.;
      if(EZrec_SIG2 > PZrec_SIG2) MZrec_SIG=sqrt(EZrec_SIG2-PZrec_SIG2);


      h_MZrec_SIG->Fill(MZrec_SIG,1);
      h_MZrec_ALL->Fill(MZrec_SIG,1);

      if(TrackIsoWmu < 0.2 ) h_MZrec_CLEAN->Fill(MZrec_SIG,1);
      
    h_MZ_SIG->Fill(MZ,1);
    h_MZ->Fill(MZ,1);
  
  }

  //read all ttbar entries and fill the histograms as before
  Int_t nentries_bkg = (Int_t)ttbar->GetEntries();
  
  for (Int_t i=0;i<nentries_bkg;i++) {
    ttbar->GetEntry(i);
      
    h_pt1_BKG->Fill(pt1_bkg,1);
    h_d0_BKG->Fill(Trackd0cutWMu_bkg,1);
    h_iso_BKG->Fill(TrackIsoWmu_bkg,1);
      
      // Compute Invariant mass
      double px1_bkg = pt1_bkg*cos(phi1_bkg);
      double py1_bkg = pt1_bkg*sin(phi1_bkg);
      double pz1_bkg = E1_bkg*tanh(eta1_bkg);
      
      double px2_bkg = pt2_bkg*cos(phi2_bkg);
      double py2_bkg = pt2_bkg*sin(phi2_bkg);
      double pz2_bkg = E2_bkg*tanh(eta2_bkg);
      
      double PZrec_BKG2 =  pow((px1_bkg+px2_bkg), 2)+ pow ((py1_bkg+py2_bkg), 2)+pow ((pz1_bkg+pz2_bkg),2);
      double EZrec_BKG2 =  (E1_bkg+E2_bkg)*(E1_bkg+E2_bkg);
      double MZrec_BKG = 0.;
      if(EZrec_BKG2 > PZrec_BKG2) MZrec_BKG=sqrt(EZrec_BKG2-PZrec_BKG2);
      
      
      h_MZrec_BKG->Fill(MZrec_BKG,1);
      h_MZrec_ALL->Fill(MZrec_BKG,1);
      
      if(TrackIsoWmu_bkg < 0.2 ) h_MZrec_CLEAN->Fill(MZrec_BKG,1);
      
    h_MZ_BKG->Fill(MZ_bkg,1);
    h_MZ->Fill(MZ_bkg,1);
  }
  
  // example how Draw and  save histograms
  //TCanvas *c = new TCanvas();
  //c->cd();
    //========================================
  /*TCanvas *c0= new TCanvas("c0"," C0  ",950,(Int_t)(950/sqrt(2)));
  c0->cd();
  h_MZ_SIG->Draw();
  h_MZ_BKG->SetLineColor(kRed);
  h_MZ_BKG->Draw("same");
  */
 //========================================
    TCanvas *c1= new TCanvas("c1"," C1  ",950,(Int_t)(950/sqrt(2)));
    c1->Divide(2,2);
    c1->cd(1)->SetGrid();
    c1->cd(2)->SetGrid();
    c1->cd(3)->SetGrid();
    c1->cd(4)->SetGrid();

    c1->cd(1);
    h_pt1_SIG->Draw();
    c1->cd(2);
    h_pt1_BKG->Draw();
    
    
    c1->cd(3);
    //plot the invariant mass of the Z for ALL events
    h_MZrec_ALL->Draw();
    
    c1->cd(4);
    //plot the invariant mass of the Z for SIGNAL and BKG events
    h_MZrec_SIG->Draw();
    h_MZrec_BKG->SetLineColor(kRed);
    h_MZrec_BKG->Draw("same");

//==========================================
    TCanvas *c2= new TCanvas("c2"," C2  ",950,(Int_t)(950/sqrt(2)));
    c2->Divide(2,2);
    c2->cd(1)->SetGrid();
    c2->cd(2)->SetGrid();
    c2->cd(3)->SetGrid();
    c2->cd(4)->SetGrid();    
    
    c2->cd(3);
    h_d0_SIG->Draw();
    c2->cd(4);
    h_d0_BKG->Draw();
    c2->cd(1);
    h_iso_SIG->Draw();
    c2->cd(2);
    h_iso_BKG->Draw();
   
    //==========================================
    TCanvas *c3= new TCanvas("c3"," C3  ",950,(Int_t)(950/sqrt(2)));
    c3->Divide(2,2);
    c3->cd(1)->SetGrid();
    c3->cd(2)->SetGrid();
    c3->cd(3)->SetGrid();
    c3->cd(4)->SetGrid();
    
    c3->cd(1);
    // Reconstructed mass all events
    h_MZrec_ALL->Draw();

    c3->cd(2);
    //Reconstructed mass selected events
    h_MZrec_CLEAN->Draw();
    
    c3->cd(3);
    h_MZ_SIG->Draw();
    
   //TCanvas *c4= new TCanvas("c4"," C4  ",950,(Int_t)(950/sqrt(2)));
    c3->cd(4);
    h_MZ_SIG->Draw();

   TF1 *fit = new TF1("fit",fitf,80,100,3);
   fit->SetParameters(666,90,3);
   fit->SetParNames("Norm","Mean","Larg");
   //func->Draw();
   h_MZrec_CLEAN->Fit("fit");
    
   //c3->cd(3);
   //h_MZ_BKG->Draw();
    
  //c->Print("test_MZ.eps");
}
void macro_final_new()
{
  tree1r();
}
double fitf(double *v, double *par)
{
 Double_t arg = 0;
 if (par[2] != 0) arg = (v[0] - par[1]);
 
    Double_t fitval = par[0]/(arg*arg+par[2]*par[2]);
    //+ par[3]*v[0]+par[4];
 return fitval;
}