#include <TFile.h>
#include <TROOT.h>
#include <TH1.h>
#include <TH2.h>
#include <TSystem.h>
#include <TGraph.h>
#include <TLorentzVector.h>
#include <TGraphAsymmErrors.h>
#include <vector>
#include <iostream>
#include <algorithm>
#include <iomanip>

#include "TMath.h"

using namespace std;

#include <TTree.h>
#include <TFile.h>
#include <TTree.h>
#include <TROOT.h>
#include <TH1.h>
#include <TH2.h>
#include <TSystem.h>
#include <vector>
#include <TLorentzVector.h>
#include <math.h>
#include "TMath.h"
#include <iostream>
#include <fstream>
#include <algorithm>
#include <map>
#include "TBranch.h"
#include "TBasket.h"

void jet1()
{
   TFile *f = new TFile("MergedMiniEvents_6.root");

   TTree *t = (TTree*)f->Get("analysis/data");
   Int_t nentries(t->GetEntriesFast());
   t->GetEntry(0);

   // Initialize Histograms
   TH1F *Z_Inv_mass = new TH1F("Z_Inv_mass","Z_Inv_mass",300,0,300);
   TH1F *W_Tran_mass = new TH1F("W_Tran_mass","W_Tran_mass",300,0,300);
   TH1F *z_pt = new TH1F("z_pt","z_pt",60,60,120);

   // Muons
   Int_t nMu;
   Bool_t Mu_isPromptFinalState[50], Mu_isDirectPromptTauDecayProductFinalState[50];
   Int_t Mu_id[50],Mu_charge[50],Mu_pid[50],Mu_g[200];
   Float_t met_pt[50], met_phi[50], Mu_pt[50],Mu_eta[50],Mu_phi[50], Mu_mass[50], Mu_scaleUnc[50], Mu_miniIso[50], Mu_chargedHadronIso[50], Mu_relIso[50], Mu_ip3d[50], Mu_ip3dsig[50],Mu_mva[50];

   // Electrons
   Int_t nEl;
   Bool_t El_isPromptFinalState[50], El_isDirectPromptTauDecayProductFinalState[50];
   Int_t El_id[50],El_charge[50],El_pid[50],El_g[200];
   Float_t El_pt[50],El_eta[50],El_phi[50], El_mass[50], El_scaleUnc[50], El_miniIso[50], El_chargedHadronIso[50], El_relIso[50], El_ip3d[50], El_ip3dsig[50],El_mva[50];

   // Jets
   Int_t nj;
   Float_t j_pt[50],j_eta[50],j_phi[50],j_mass[50],j_csv[50];

   // Jet info
   t->SetBranchAddress("nj",   &nj);
   t->SetBranchAddress("j_pt",        j_pt);
   t->SetBranchAddress("j_eta",      j_eta);
   t->SetBranchAddress("j_phi",      j_phi);
   t->SetBranchAddress("j_mass",    j_mass);
   t->SetBranchAddress("j_csv",     j_csv);

   //Electron info
   t->SetBranchAddress("nEl", &nEl);
   t->SetBranchAddress("El_isPromptFinalState",                         El_isPromptFinalState);
   t->SetBranchAddress("El_isDirectPromptTauDecayProductFinalState",    El_isDirectPromptTauDecayProductFinalState);
   t->SetBranchAddress("El_mva",      El_mva);
   t->SetBranchAddress("El_id",       El_id);
   t->SetBranchAddress("El_pid",      El_pid);
   t->SetBranchAddress("El_g",        El_g);
   t->SetBranchAddress("El_charge",   El_charge);
   t->SetBranchAddress("El_pt",       El_pt);
   t->SetBranchAddress("El_eta",      El_eta);
   t->SetBranchAddress("El_phi",      El_phi);
   t->SetBranchAddress("El_mass",     El_mass);
   t->SetBranchAddress("El_scaleUnc",         El_scaleUnc);
   t->SetBranchAddress("El_chargedHadronIso", El_chargedHadronIso);
   t->SetBranchAddress("El_miniIso",          El_miniIso);
   t->SetBranchAddress("El_relIso",           El_relIso);
   t->SetBranchAddress("El_ip3d",             El_ip3d);
   t->SetBranchAddress("El_ip3dsig",          El_ip3dsig);

   //Muon info
   t->SetBranchAddress("nMu", &nMu);
   t->SetBranchAddress("Mu_isPromptFinalState",                         Mu_isPromptFinalState);
   t->SetBranchAddress("Mu_isDirectPromptTauDecayProductFinalState",    Mu_isDirectPromptTauDecayProductFinalState);
   t->SetBranchAddress("Mu_mva",      Mu_mva);
   t->SetBranchAddress("Mu_id",       Mu_id);
   t->SetBranchAddress("Mu_pid",      Mu_pid);
   t->SetBranchAddress("Mu_g",        Mu_g);
   t->SetBranchAddress("Mu_charge",   Mu_charge);
   t->SetBranchAddress("Mu_pt",       Mu_pt);
   t->SetBranchAddress("Mu_eta",      Mu_eta);
   t->SetBranchAddress("Mu_phi",      Mu_phi);
   t->SetBranchAddress("Mu_mass",     Mu_mass);
   t->SetBranchAddress("Mu_scaleUnc",         Mu_scaleUnc);
   t->SetBranchAddress("Mu_chargedHadronIso", Mu_chargedHadronIso);
   t->SetBranchAddress("Mu_miniIso",          Mu_miniIso);
   t->SetBranchAddress("Mu_relIso",           Mu_relIso);
   t->SetBranchAddress("Mu_ip3d",             Mu_ip3d);
   t->SetBranchAddress("Mu_ip3dsig",          Mu_ip3dsig);
   t->SetBranchAddress("met_pt",      met_pt);
   t->SetBranchAddress("met_phi",      met_phi);

   double min_m =1000.0, m_tmp, mmass_tmp;
   double z_mass = 91.1876, dphi_mu=0, dphi_mu1=999, mhard_pt =0., zmass=0. , zpt =0;
   int j=0, jj=0, m_i =999, m_ii=999;


   //LOOP OVER EVENTS
   for (Int_t iev=0;iev<nentries;iev++) {
      t->GetEntry(iev);
      if (iev%1000==0) printf ("\r [%3.0f/100] done",100.*(float)(iev)/(float)(nentries));

      std::vector<TLorentzVector> selectedMuons, selectedElectrons;
      std::vector<int> mu_charge, el_charge, mu_isTight, mu_ismuon;
      std::vector<float> mu_iso;

         // selected muons
         for (int m=0; m<nMu;m++) {
            TLorentzVector tmp(0,0,0,0);
            bool passTightKin(Mu_pt[m]>=10 && fabs(Mu_eta[m])<2.5); // I made selected muons as a loose to make it equal to veto muons by changing cuts
            float relIso(Mu_relIso[m]);
            bool passIso(relIso<0.5);
            tmp.SetPtEtaPhiM(Mu_pt[m], Mu_eta[m], Mu_phi[m], Mu_mass[m]);
            if (passTightKin ) { //&& passSIP3d)
               if (passIso) {
                  selectedMuons.push_back(tmp);
                  mu_charge.push_back(Mu_charge[m]);
               }
            }
         }

         //Selected Electrons
         for (int e=0; e<nEl;e++) {
            TLorentzVector tmp1(0,0,0,0);
            bool  e_passTightKin(El_pt[e]>=10 && fabs(El_eta[e])<2.5);
            float e_relIso(El_relIso[e]);
            bool  e_passIso(e_relIso<0.5);
            tmp1.SetPtEtaPhiM(El_pt[e], El_eta[e], El_phi[e], El_mass[e]);
            if (e_passTightKin) {
               if (e_passIso) {
                  selectedElectrons.push_back(tmp1);
                  el_charge.push_back(El_charge[e]);
               }
            }
         }

         int nSelMuons = selectedMuons.size();
         int nSelElectrons = selectedElectrons.size();
         int nSelLeptons = nSelElectrons + nSelMuons;

         //THREE LEPTON SELECTION CUT
         if (nSelLeptons != 3) continue;

         // Jets selection

         float csv = -999.;
         std::vector<TLorentzVector> bJets,lightJets, allJets;

         for (int k=0; k<nj; k++) {
            TLorentzVector jp4;
            jp4.SetPtEtaPhiM(j_pt[k],j_eta[k],j_phi[k],j_mass[k]);

            // Delta R between leptons and jets

            float dR_jet_mu=999., dR_jet_el=999.;
            for (unsigned int mu=0; mu <selectedMuons.size(); ++mu) {
               dR_jet_mu = jp4.DeltaR(selectedMuons[mu]);
               if (dR_jet_mu <= 0.4) break;
            }
            if (dR_jet_mu <= 0.4) continue;
            for (unsigned int el=0; el <selectedElectrons.size(); ++el){
               dR_jet_el = jp4.DeltaR(selectedElectrons[el]);
               if (dR_jet_el <= 0.4) break;
            }
            if (dR_jet_el <= 0.4) continue;

            if (jp4.Pt()<20) continue;
            if (fabs(jp4.Eta()) > 4.5) continue;

            csv = j_csv[k];
            bool isBTagged(csv> 0.9535); /// Tight

            //save jet
            if (isBTagged) bJets.push_back(jp4);
            else           lightJets.push_back(jp4);
         }

         // Z Mass from Muons
         j = 0;
         for (unsigned int mu1=0; mu1 <selectedMuons.size(); ++mu1) {
            jj=0;
            for (unsigned int mu2=0; mu2 <selectedMuons.size(); ++mu2) {
               if ( (mu2 > mu1) &&  (mu_charge[mu1]*mu_charge[mu2]) < 0. ) {
                  dphi_mu = (selectedMuons)[mu1].Phi() - (selectedMuons)[mu2].Phi();
                  if ( fabs(dphi_mu)< 0.5) continue;
                  dphi_mu1=dphi_mu;
                  mmass_tmp = ((selectedMuons)[mu1] + (selectedMuons)[mu2]).M();

                  if ( selectedMuons[mu1].Pt() > selectedMuons[mu2].Pt() ) {
                     mhard_pt =  selectedMuons[mu1].Pt();
                  } else {
                     mhard_pt =  selectedMuons[mu2].Pt();
                  }
                  m_tmp =  z_mass - mmass_tmp;
                  if (fabs(m_tmp) <min_m  && mhard_pt > 20.) {
                     min_m = fabs(m_tmp); m_i = j; m_ii = jj;
                  }
               } //end of sfos if loop
               jj++;
            } //end of nested mu 2 loop
            j++;
         } //end of nested mu 1 loop

         bool is2muon = false; j=0; jj=0;
         if (m_i != 999 && m_ii != 999) is2muon =true;
         if (is2muon ) {
            if (mu_isTight[m_i]!=1 && mu_ismuon[m_i] !=1 && mu_iso[m_i] >0.12) continue;
            if (mu_isTight[m_ii]!=1 && mu_ismuon[m_ii] !=1 && mu_iso[m_ii] >0.12) continue;
            zmass = (selectedMuons[m_i] + selectedMuons[m_ii]).M();
            zpt   = (selectedMuons[m_i] + selectedMuons[m_ii]).Pt();
            Z_Inv_mass->Fill(zmass);
            z_pt -> Fill(zpt);
         }

         TLorentzVector met(0,0,0,0),  mup4_w (0,0,0,0);
         bool isWmuon =false;
         if (is2muon ) {
         for ( int mu=0; mu < (int) selectedMuons.size(); ++mu) {
            if ( mu==m_i || mu==m_ii ) continue;
            mup4_w = TLorentzVector((selectedMuons)[mu].Px(), (selectedMuons)[mu].Py(),
                                    (selectedMuons)[mu].Pz(), (selectedMuons)[mu].Energy() );

         }
      }

      //MET and transverse mass
      double  mt_mu=0.; //TLorentzVector mu_neutrino(0,0,0,0);
      met.SetPtEtaPhiM(met_pt[0],0,met_phi[0],0.);
      met.SetPz(0.); met.SetE(met.Pt());

      if (isWmuon) {
         double   dphi_mumet = (mup4_w - met).Phi();
         if ( fabs(dphi_mumet)< 0.5) continue;
         mt_mu = TMath::Sqrt(2* mup4_w.Pt()*met.Pt()*(1-TMath::Cos( dphi_mumet )));
         W_Tran_mass->Fill(mt_mu);
      }

      TFile *ff = new TFile("dg6.root","RECREATE");
      Z_Inv_mass->Write();
      z_pt->Write();
      W_Tran_mass->Write();

      ff->Close();
   }
   f->Close();
}