#define NanoTree_cxx
#include "NanoTree.h"
#include <TH2.h>
#include <TStyle.h>
#include <TCanvas.h>
#include <fstream>
#include <cmath>
#include "TLorentzVector.h"

ClassImp(NanoTree);

NanoTree::NanoTree(TTree *tree) : fChain(0) 
{
   Init(tree);
}

NanoTree::~NanoTree()
{
   if (!fChain) return;
   delete fChain->GetCurrentFile();
}

Int_t NanoTree::GetEntry(Long64_t entry)
{
   if (!fChain) return 0;
   return fChain->GetEntry(entry);
}
Long64_t NanoTree::LoadTree(Long64_t entry)
{
   if (!fChain) return -5; 
   Long64_t centry = fChain->LoadTree(entry);
   if (centry < 0) return centry;
   if (fChain->GetTreeNumber() != fCurrent) {
      fCurrent = fChain->GetTreeNumber();
      Notify();
   }   
   return centry;
}

void NanoTree::Init(TTree *tree)
{
   // Set branch addresses and branch pointers
   if (!tree) return;
   fChain = tree;
   fCurrent = -1; 
   fChain->SetMakeClass(1);

   fChain->SetBranchAddress("run", &run, &b_run);
   fChain->SetBranchAddress("luminosityBlock", &luminosityBlock, &b_luminosityBlock);
   fChain->SetBranchAddress("event", &event, &b_event);
   fChain->SetBranchAddress("nMuon", &nMuon, &b_nMuon);
   fChain->SetBranchAddress("Muon_eta", Muon_eta, &b_Muon_eta);
   fChain->SetBranchAddress("Muon_mass", Muon_mass, &b_Muon_mass);
   fChain->SetBranchAddress("Muon_phi", Muon_phi, &b_Muon_phi);
   fChain->SetBranchAddress("Muon_pt", Muon_pt, &b_Muon_pt);
   fChain->SetBranchAddress("Muon_tightId", Muon_tightId, &b_Muon_tightId);
   fChain->SetBranchAddress("Muon_charge", Muon_charge, &b_Muon_charge);
   Notify();
}

Bool_t NanoTree::Notify()
{
   return kTRUE;
}

void NanoTree::Show(Long64_t entry)
{
   if (!fChain) return;
   fChain->Show(entry);
}
Int_t NanoTree::Cut(Long64_t entry)
{
   return 1;
}

void NanoTree::Loop(const char* output)
{
    if (fChain == 0) return;
    Long64_t nentries = fChain->GetEntriesFast();
    Long64_t nbytes = 0, nb = 0;

    TFile *testFile = new TFile(output,"RECREATE");
    TTree t1("AnalysisTree","a simple Tree with simple variables");

    Float_t leadingMuPt,leadingMuEta,leadingMuPhi,leadingMuMass,subleadingMuPt,subleadingMuEta,subleadingMuPhi,subleadingMuMass,MuMuMass;
 
    t1.Branch("leadingMuPt",&leadingMuPt,"leadingMuPt/F");
    t1.Branch("leadingMuEta",&leadingMuEta,"leadingMuEta/F");
    t1.Branch("leadingMuPhi",&leadingMuPhi,"leadingMuPhi/F");
    t1.Branch("leadingMuMass",&leadingMuMass,"leadingMuMass/F");
    t1.Branch("subleadingMuPt",&subleadingMuPt,"subleadingMuPt/F");
    t1.Branch("subleadingMuEta",&subleadingMuEta,"subleadingMuEta/F");
    t1.Branch("subleadingMuPhi",&subleadingMuPhi,"subleadingMuPhi/F");
    t1.Branch("subleadingMuMass",&subleadingMuMass,"subleadingMuMass/F");
    t1.Branch("MuMuMass",&MuMuMass,"MuMuMass/F");

    for (Long64_t jentry=0; jentry<nentries;jentry++) {
        Long64_t ientry = LoadTree(jentry);
        if (ientry < 0) break;
        nb = fChain->GetEntry(jentry);   nbytes += nb;

        double mp_pt = -1;
        int mp_index = -1;
        double mm_pt = -1;
        int mm_index = -1;
        for (unsigned int i = 0; i < nMuon; i++) {
            if (Muon_tightId[i] && Muon_charge[i] == +1 && Muon_pt[i] > mp_pt) {
                mp_pt = Muon_pt[i];
                mp_index = i;
            }
            else if (Muon_tightId[i] && Muon_charge[i] == -1 && Muon_pt[i] > mm_pt) {
                mm_pt = Muon_pt[i];
                mm_index = i;
            }
        }
        int leading_muon_index = -1;
        int subleading_muon_index = -1;
        if (mp_pt > mm_pt) {
            leading_muon_index = mp_index;
            subleading_muon_index = mm_index;
        }
        else {
            leading_muon_index = mm_index;
            subleading_muon_index = mp_index;
        }

        if (mm_index != -1 && mp_index != -1) {
            TLorentzVector mm;
            TLorentzVector mp;
            mm.SetPtEtaPhiM(Muon_pt[mm_index],Muon_eta[mm_index],Muon_phi[mm_index],Muon_mass[mm_index]);
            mp.SetPtEtaPhiM(Muon_pt[mp_index],Muon_eta[mp_index],Muon_phi[mp_index],Muon_mass[mp_index]);
            TLorentzVector mumu = mm + mp;
            MuMuMass = mumu.M();
            leadingMuPt = Muon_pt[leading_muon_index];
            leadingMuEta = Muon_eta[leading_muon_index];
            leadingMuPhi = Muon_phi[leading_muon_index];
            leadingMuMass = Muon_mass[leading_muon_index];
            subleadingMuPt = Muon_pt[subleading_muon_index];
            subleadingMuEta = Muon_eta[subleading_muon_index];
            subleadingMuPhi = Muon_phi[subleading_muon_index];
            subleadingMuMass = Muon_mass[subleading_muon_index];

        }
        t1.Fill();
    }

    t1.Write();
    testFile->Close();
}