//efficiency check of double misID filter in B+->Kst+mumu
//david gerick


#include "GlobalFunctions.hh"

using namespace std;
using namespace RooFit ;

//////////////////////////////////////////////////////////////////////
/// DoubleMisIDcheck()
///  fake the double-misID of pi+ and mu+ and see how many would pass the double-misID filter from the preselection
/// 
///
///

int checkDoubleMisIDvetoEff(std::string year = "2011", std::string magnet = "down", bool PreSelAndTM = true) {

    TChain* tree = 0;

    //Kst2Kpluspi0 case
    if(Kst2Kpluspi0Resolved || Kst2Kpluspi0Merged){
      std::cout << "[ERROR]\t\tNo double misID possible in pi0 channel!" << std::endl;
      return 0;
    }

    //Kst2Kspiplus case
    if(Kst2Kspiplus){
        if(PreSelAndTM){
            tree=new TChain("DecayTreeTruthMatched");
            tree->Add(GetInputFile(year,magnet,true,true,true,false,false).c_str());
        }
        else{
            tree=new TChain("b2KstKs0pimumu_Tuple/DecayTree");
            tree->Add(GetInputFile(year,magnet,false,true,true,false,false).c_str());
        }
    }

    int N = tree->GetEntries();
    cout << "Old files contain " << N << " events" <<  endl;

    if(N == 0){
        std::cout << "[ERROR]\t\tTree has no entries. Exit!" << std::endl;
        return 0;
    }

    //Disable all branches
    tree->SetBranchStatus("*",0);

    tree->SetBranchStatus( "*PX", 1 );
    tree->SetBranchStatus( "*PY", 1);
    tree->SetBranchStatus( "*PZ", 1);
    tree->SetBranchStatus( "*PE", 1);

    //define variables
    TLorentzVector LorVec_B_plus;
    TLorentzVector LorVec_K_star_plus;
    TLorentzVector LorVec_K_short;
    TLorentzVector LorVec_pi_plus;
    TLorentzVector LorVec_KS_pi_plus;
    TLorentzVector LorVec_KS_pi_minus;
    TLorentzVector LorVec_mu_minus;
    TLorentzVector LorVec_mu_plus;

    //potentially misIDed particles:
    TLorentzVector LorVec_pi_plus_MisIDed;
    TLorentzVector LorVec_mu_plus_MisIDed;
    TLorentzVector LorVec_B_plus_From_DoubleMisIDed;
    TLorentzVector LorVec_Jpsi_From_DoubleMisIDed;
    TLorentzVector LorVec_K_star_plus_From_DoubleMisIDed;

    //recombined particles:
    TLorentzVector LorVec_DiMuon;
    TLorentzVector LorVec_KaonPion;
    TLorentzVector LorVec_ReconB;
    TLorentzVector LorVec_KshortMuMu;

        ///B+
    Double_t B_plus_PX    = 0.;
    Double_t B_plus_PY    = 0.;
    Double_t B_plus_PZ    = 0.;
    Double_t B_plus_PE    = 0.;

        ///K*+
    Double_t K_star_plus_PX     = 0.;
    Double_t K_star_plus_PY     = 0.;
    Double_t K_star_plus_PZ     = 0.;
    Double_t K_star_plus_PE     = 0.;

        ///Kshort
    Double_t K_short_PX     = 0.;
    Double_t K_short_PY     = 0.;
    Double_t K_short_PZ     = 0.;
    Double_t K_short_PE     = 0.;

        ///pi+ (KS)
    Double_t Ks_pi_plus_PX     = 0.;
    Double_t Ks_pi_plus_PY     = 0.;
    Double_t Ks_pi_plus_PZ     = 0.;
    Double_t Ks_pi_plus_PE     = 0.;

    Float_t f_Ks_pi_plus_PX_DTF[100];
    Float_t f_Ks_pi_plus_PY_DTF[100];
    Float_t f_Ks_pi_plus_PZ_DTF[100];
    Float_t f_Ks_pi_plus_PE_DTF[100];

        ///pi- (KS)
    Double_t Ks_pi_minus_PX     = 0.;
    Double_t Ks_pi_minus_PY     = 0.;
    Double_t Ks_pi_minus_PZ     = 0.;
    Double_t Ks_pi_minus_PE     = 0.;

    Float_t f_Ks_pi_minus_PX_DTF[100];
    Float_t f_Ks_pi_minus_PY_DTF[100];
    Float_t f_Ks_pi_minus_PZ_DTF[100];
    Float_t f_Ks_pi_minus_PE_DTF[100];

        ///pi+
    Double_t pi_plus_PX     = 0.;
    Double_t pi_plus_PY     = 0.;
    Double_t pi_plus_PZ     = 0.;
    Double_t pi_plus_PE     = 0.;

    Float_t f_pi_plus_PX_DTF[100];
    Float_t f_pi_plus_PY_DTF[100];
    Float_t f_pi_plus_PZ_DTF[100];
    Float_t f_pi_plus_PE_DTF[100];

        ///mu-
    Double_t mu_minus_PX     = 0.;
    Double_t mu_minus_PY     = 0.;
    Double_t mu_minus_PZ     = 0.;
    Double_t mu_minus_PE     = 0.;

    Float_t f_mu_minus_PX_DTF[100];
    Float_t f_mu_minus_PY_DTF[100];
    Float_t f_mu_minus_PZ_DTF[100];
    Float_t f_mu_minus_PE_DTF[100];

        ///mu+
    Double_t mu_plus_PX     = 0.;
    Double_t mu_plus_PY     = 0.;
    Double_t mu_plus_PZ     = 0.;
    Double_t mu_plus_PE     = 0.;

    Float_t f_mu_plus_PX_DTF[100];
    Float_t f_mu_plus_PY_DTF[100];
    Float_t f_mu_plus_PZ_DTF[100];
    Float_t f_mu_plus_PE_DTF[100];

    ///////////////////////////////////////
    /// LINK VARIABLES TO BRANCHES
    ///////////////////////////////////////


    //particle masses and momenta
    tree -> SetBranchAddress( "B_plus_PX" , &B_plus_PX );
    tree -> SetBranchAddress( "B_plus_PY" , &B_plus_PY );
    tree -> SetBranchAddress( "B_plus_PZ" , &B_plus_PZ );
    tree -> SetBranchAddress( "B_plus_PE" , &B_plus_PE );

    tree -> SetBranchAddress( "K_star_plus_PX" , &K_star_plus_PX );
    tree -> SetBranchAddress( "K_star_plus_PY" , &K_star_plus_PY );
    tree -> SetBranchAddress( "K_star_plus_PZ" , &K_star_plus_PZ );
    tree -> SetBranchAddress( "K_star_plus_PE" , &K_star_plus_PE );

    tree -> SetBranchAddress( "K_short_PX" , &K_short_PX );
    tree -> SetBranchAddress( "K_short_PY" , &K_short_PY );
    tree -> SetBranchAddress( "K_short_PZ" , &K_short_PZ );
    tree -> SetBranchAddress( "K_short_PE" , &K_short_PE );

    tree -> SetBranchAddress( "pi_plus_PX" , &pi_plus_PX );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_piplus_PX" , &f_pi_plus_PX_DTF );
    tree -> SetBranchAddress( "pi_plus_PY" , &pi_plus_PY );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_piplus_PY" , &f_pi_plus_PY_DTF );
    tree -> SetBranchAddress( "pi_plus_PZ" , &pi_plus_PZ );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_piplus_PZ" , &f_pi_plus_PZ_DTF );
    tree -> SetBranchAddress( "pi_plus_PE" , &pi_plus_PE );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_piplus_PE" , &f_pi_plus_PE_DTF );

    tree -> SetBranchAddress( "Ks_pi_minus_PX" , &Ks_pi_minus_PX );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_PX" , &f_Ks_pi_minus_PX_DTF );
    tree -> SetBranchAddress( "Ks_pi_minus_PY" , &Ks_pi_minus_PY );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_PY" , &f_Ks_pi_minus_PY_DTF );
    tree -> SetBranchAddress( "Ks_pi_minus_PZ" , &Ks_pi_minus_PZ );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_PZ" , &f_Ks_pi_minus_PZ_DTF );
    tree -> SetBranchAddress( "Ks_pi_minus_PE" , &Ks_pi_minus_PE );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_PE" , &f_Ks_pi_minus_PE_DTF );

    tree -> SetBranchAddress( "Ks_pi_plus_PX" , &Ks_pi_plus_PX );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_0_PX" , &f_Ks_pi_plus_PX_DTF );
    tree -> SetBranchAddress( "Ks_pi_plus_PY" , &Ks_pi_plus_PY );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_0_PY" , &f_Ks_pi_plus_PY_DTF );
    tree -> SetBranchAddress( "Ks_pi_plus_PZ" , &Ks_pi_plus_PZ );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_0_PZ" , &f_Ks_pi_plus_PZ_DTF );
    tree -> SetBranchAddress( "Ks_pi_plus_PE" , &Ks_pi_plus_PE );
    tree -> SetBranchAddress( "B_plus_DTF_Kst_892_plus_KS0_piplus_0_PE" , &f_Ks_pi_plus_PE_DTF );

    tree -> SetBranchAddress( "mu_minus_PX" , &mu_minus_PX );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_0_PX" , &f_mu_minus_PX_DTF );
    tree -> SetBranchAddress( "mu_minus_PY" , &mu_minus_PY );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_0_PY" , &f_mu_minus_PY_DTF );
    tree -> SetBranchAddress( "mu_minus_PZ" , &mu_minus_PZ );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_0_PZ" , &f_mu_minus_PZ_DTF );
    tree -> SetBranchAddress( "mu_minus_PE" , &mu_minus_PE );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_0_PE" , &f_mu_minus_PE_DTF );

    tree -> SetBranchAddress( "mu_plus_PX" , &mu_plus_PX );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_PX" , &f_mu_plus_PX_DTF );
    tree -> SetBranchAddress( "mu_plus_PY" , &mu_plus_PY );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_PY" , &f_mu_plus_PY_DTF );
    tree -> SetBranchAddress( "mu_plus_PZ" , &mu_plus_PZ );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_PZ" , &f_mu_plus_PZ_DTF );
    tree -> SetBranchAddress( "mu_plus_PE" , &mu_plus_PE );
    tree -> SetBranchAddress( "B_plus_DTF_J_psi_1S_muminus_PE" , &f_mu_plus_PE_DTF );

    int n_DoubleMisID = 0;

    bool IsBplus = true;

    ///////////////////////////////////////////////////////////////////////////
    ///
    /// loop over events
    ///
    ///////////////////////////////////////////////////////////////////////////

    for(int i=0; i< N; i++){

        //output the different counters!
        if ((0ul == (i % 10000ul) || i + 1 == N) && i != 0)
            cout << "Read event " << i << "/" << N << std::endl;
        tree->GetEntry(i);

        //define the Lorentz vectors
        if(UseDTF){ //start useDTF for LorentzVectors
            LorVec_mu_minus.SetPxPyPzE        (f_mu_minus_PX_DTF[0],           f_mu_minus_PY_DTF[0],            f_mu_minus_PZ_DTF[0],            f_mu_minus_PE_DTF[0]);
            LorVec_mu_plus.SetPxPyPzE         (f_mu_plus_PX_DTF[0],            f_mu_plus_PY_DTF[0],             f_mu_plus_PZ_DTF[0],             f_mu_plus_PE_DTF[0]);
            LorVec_pi_plus.SetPxPyPzE         (f_pi_plus_PX_DTF[0],            f_pi_plus_PY_DTF[0],             f_pi_plus_PZ_DTF[0],             f_pi_plus_PE_DTF[0]);
            LorVec_KS_pi_plus.SetPxPyPzE      (f_Ks_pi_plus_PX_DTF[0],         f_Ks_pi_plus_PY_DTF[0],          f_Ks_pi_plus_PZ_DTF[0],          f_Ks_pi_plus_PE_DTF[0]);
            LorVec_KS_pi_minus.SetPxPyPzE     (f_Ks_pi_minus_PX_DTF[0],        f_Ks_pi_minus_PY_DTF[0],         f_Ks_pi_minus_PZ_DTF[0],         f_Ks_pi_minus_PE_DTF[0]);
            LorVec_pi_plus_MisIDed.SetXYZM    (f_pi_plus_PX_DTF[0],            f_pi_plus_PY_DTF[0],             f_pi_plus_PZ_DTF[0],             PDGMASS.MU);
            LorVec_mu_plus_MisIDed.SetXYZM    (f_mu_plus_PX_DTF[0],            f_mu_plus_PY_DTF[0],             f_mu_plus_PZ_DTF[0],             PDGMASS.PI_PLUS);
            LorVec_K_short = LorVec_KS_pi_plus + LorVec_KS_pi_minus;
        }
        else{
          std::cout << "[ERROR]\tTest not implemented for non-DTF variables" << std::endl;
          return 0;
        }

        LorVec_DiMuon = LorVec_mu_minus + LorVec_mu_plus;
        LorVec_KaonPion = LorVec_K_short + LorVec_pi_plus;

        LorVec_K_star_plus = LorVec_KaonPion;
        LorVec_ReconB = LorVec_DiMuon + LorVec_K_star_plus;

        //double misIDed
        if(Kst2Kspiplus){
            LorVec_K_star_plus_From_DoubleMisIDed = LorVec_K_short + LorVec_mu_plus_MisIDed;
            LorVec_Jpsi_From_DoubleMisIDed = LorVec_pi_plus_MisIDed + LorVec_mu_minus;
            LorVec_B_plus_From_DoubleMisIDed = LorVec_Jpsi_From_DoubleMisIDed + LorVec_K_star_plus_From_DoubleMisIDed;
        }

        //remove double mis-identified events:
        if(Kst2Kspiplus){
            if(TMath::Abs(LorVec_Jpsi_From_DoubleMisIDed.M() - PDGMASS.J_PSI) > 50 && TMath::Abs(LorVec_K_star_plus_From_DoubleMisIDed.M() - PDGMASS.K_STAR_PLUS) > 30)
              if(TMath::Abs(LorVec_K_star_plus_From_DoubleMisIDed.M() - PDGMASS.K_STAR_PLUS) < 100)
                if( LorVec_Jpsi_From_DoubleMisIDed.M2() > 0.1e6 
                && (LorVec_Jpsi_From_DoubleMisIDed.M2() < 0.98e6 || LorVec_Jpsi_From_DoubleMisIDed.M2() > 1.1e6) 
                && (LorVec_Jpsi_From_DoubleMisIDed.M2() < 8.68e6 || LorVec_Jpsi_From_DoubleMisIDed.M2() > 10.09e6) 
                && (LorVec_Jpsi_From_DoubleMisIDed.M2() < 12.9e6 || LorVec_Jpsi_From_DoubleMisIDed.M2() > 14.4e6) 
                &&  LorVec_Jpsi_From_DoubleMisIDed.M2() < 19.0e6)
                n_DoubleMisID++;
                continue;
        }
    }

    std::cout << "[INFO]\t\tThe double-misID veto was passed by " << n_DoubleMisID << " events. Out of " << N << ". Making the vetoes efficiency = " << 1. * n_DoubleMisID / N << std::endl;

    return 1;
}

