#include <cstdlib>
#include <vector>
#include <iostream>
#include <map>
#include <string>

#include "TChain.h"
#include "TFile.h"
#include "TMath.h"
#include "TTree.h"
#include "TString.h"
#include "TSystem.h"
#include "TROOT.h"
#include "TStopwatch.h"

#include "../GlobalFunctions.hh"
#include "../HeidelbergFitter/LHCbStyle.h"

using namespace std;

//////////////////////////////////////////////////////
///  GetDoubleMissIDs()
/// creates J_psi mass spectra from the non-resonant selected events after the BDT cut,
/// with the mu+ and pi+ double  mis-identified

Int_t GetDoubleMissIDs(bool AfterBDTcut = true, Int_t Run = 1){

    if(!Kst2Kspiplus){
        std::cout << "[ERROR]\t\tDouble mis-identification only possible in Kshort Piplus subdecay!" << std::endl;
        return 0;
    }

    TChain* tree = NULL;
    if(AfterBDTcut){
        tree = new TChain("SelectionOutput");
        tree->Add(Form("%s/data/%s_BDToutput.root",path_to_data.c_str(),TheDecay.c_str()));
    }
    else {
        tree = new TChain("DecayTree");
	if(Run == 1 || Run == 12){
            tree->Add(Form("%s/data/2011_%s_DD_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
            tree->Add(Form("%s/data/2011_%s_LL_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
            tree->Add(Form("%s/data/2012_%s_DD_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
            tree->Add(Form("%s/data/2012_%s_LL_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
	}
	if(Run == 2 || Run == 12){
            tree->Add(Form("%s/data/2015_%s_DD_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
            tree->Add(Form("%s/data/2015_%s_LL_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
            tree->Add(Form("%s/data/2016_%s_DD_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
            tree->Add(Form("%s/data/2016_%s_LL_BDTinput.root",path_to_data.c_str(),TheDecay.c_str()));
        }
    }
    Int_t N = tree->GetEntries();
    if(N == 0){
        std::cout << "[ERROR]\t\tNo events found for tree in given files!" << std::endl;
        return 0;
    }
    std::cout << "[INFO]\t\tTree loaded: Seeing " << N << " events!" << std::endl;

    TLorentzVector LorVec_mu_minus;
    TLorentzVector LorVec_mu_plus;
    TLorentzVector LorVec_pi_plus;
    TLorentzVector LorVec_K_short;
    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;

    Double_t pi_plus_PX = 0.;
    Double_t pi_plus_PY = 0.;
    Double_t pi_plus_PZ = 0.;

    Double_t mu_minus_PX = 0.;
    Double_t mu_minus_PY = 0.;
    Double_t mu_minus_PZ = 0.;

    Double_t mu_plus_PX = 0.;
    Double_t mu_plus_PY = 0.;
    Double_t mu_plus_PZ = 0.;

    Double_t K_short_PX = 0.;
    Double_t K_short_PY = 0.;
    Double_t K_short_PZ = 0.;

    tree -> SetBranchAddress( "pi_plus_PX" , &pi_plus_PX );
    tree -> SetBranchAddress( "pi_plus_PY" , &pi_plus_PY );
    tree -> SetBranchAddress( "pi_plus_PZ" , &pi_plus_PZ );

    tree -> SetBranchAddress( "mu_minus_PX" , &mu_minus_PX );
    tree -> SetBranchAddress( "mu_minus_PY" , &mu_minus_PY );
    tree -> SetBranchAddress( "mu_minus_PZ" , &mu_minus_PZ );

    tree -> SetBranchAddress( "mu_plus_PX" , &mu_plus_PX );
    tree -> SetBranchAddress( "mu_plus_PY" , &mu_plus_PY );
    tree -> SetBranchAddress( "mu_plus_PZ" , &mu_plus_PZ );

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


    Double_t dReconstructedBM;
    tree->Branch("BplusInvMass", &dReconstructedBM, "dReconstructedBM/D");

    std::cout << "[INFO]\t\tVariables created and linked to branches!" << std::endl;

    TH2 * h_BplusM_vs_BplusM_DoubleMisIDed = new TH2D("h_BplusM_vs_BplusM_DoubleMisIDed" , "comparison of B+ mass versus B+ mass (mu+ and pi+ double misIDed", 100, 5000, 6000, 100, 5000, 6000);
    h_BplusM_vs_BplusM_DoubleMisIDed->GetXaxis()->SetTitle("B_plus_M [MeV]");
    h_BplusM_vs_BplusM_DoubleMisIDed->GetYaxis()->SetTitle("B_plus_M (double misIDed) [MeV]");

    TH2 * h_BplusM_vs_JpsiM_DoubleMisIDed = new TH2D("h_BplusM_vs_JpsiM_DoubleMisIDed" , "overlay of J_psi mass and B+ mass (both mu+ and pi+ double misIDed", 50, 5000, 6000, 200, 0, 4000);
    h_BplusM_vs_JpsiM_DoubleMisIDed->GetXaxis()->SetTitle("B_plus_M (double misIDed) [MeV]");
    h_BplusM_vs_JpsiM_DoubleMisIDed->GetYaxis()->SetTitle("J_psi_M (double misIDed) [MeV]");

    TH1 * h_Jpsi_Spectrum_MisIDed_Muon = new TH1D("h_Jpsi_Spectrum_MisIDed_Muon" , "Invariant mass of Jpsi with a pion misIDed as a muon", 200, 0, 4000);
    h_Jpsi_Spectrum_MisIDed_Muon->GetXaxis()->SetTitle("J_psi_M (double misIDed) [MeV]");
    h_Jpsi_Spectrum_MisIDed_Muon->GetYaxis()->SetTitle("Candidates/20MeV");

    TH2 * h_JpsiM_vs_KstplusM_MisIDed_Muon = new TH2D("h_Jpsi_vs_Kstarplus_InvMass" , "m(Jpsi) with pi+ misIDed as a mu+ and K_star_plus with mu+ misIDed as a pi+", 100, 0, 4000, 100, 600, 2600);
    h_JpsiM_vs_KstplusM_MisIDed_Muon->GetXaxis()->SetTitle("J_psi_M (double misIDed) [40MeV]");
    h_JpsiM_vs_KstplusM_MisIDed_Muon->GetYaxis()->SetTitle("K_star_plus_M (double misIDed) [20MeV]");

    TH1 * h_Q2_Spectrum_MisIDed_Muon = new TH1D("h_Q2_Spectrum_MisIDed_Muon" , "Q2 of non-resonant DiMuons with a pion misIDed as a muon", 100, 0, 20);
    h_Q2_Spectrum_MisIDed_Muon->GetXaxis()->SetTitle("Q2 (double misIDed) [GeV^{2}]");
    h_Q2_Spectrum_MisIDed_Muon->GetYaxis()->SetTitle("Candidates / 0.2 GeV^{2}");

    std::cout << "[INFO]\t\tHistograms initialized!" << std::endl;

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

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

        tree->GetEntry(i);

        //define the Lorentz vectors
        LorVec_mu_minus.SetXYZM         (mu_minus_PX,       mu_minus_PY,        mu_minus_PZ,        PDGMASS.MU);
        LorVec_mu_plus.SetXYZM          (mu_plus_PX,        mu_plus_PY,         mu_plus_PZ,         PDGMASS.MU);
        LorVec_pi_plus.SetXYZM          (pi_plus_PX,        pi_plus_PY,         pi_plus_PZ,         PDGMASS.PI_PLUS);
        LorVec_K_short.SetXYZM          (K_short_PX,        K_short_PY,         K_short_PZ,         PDGMASS.K_SHORT);

        LorVec_pi_plus_MisIDed.SetXYZM          (pi_plus_PX,        pi_plus_PY,         pi_plus_PZ,         PDGMASS.MU);
        LorVec_mu_plus_MisIDed.SetXYZM          (mu_plus_PX,        mu_plus_PY,         mu_plus_PZ,         PDGMASS.PI_PLUS);


        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;


        Double_t dJpsiMisIDed_M = LorVec_Jpsi_From_DoubleMisIDed.M();

        h_Jpsi_Spectrum_MisIDed_Muon->Fill(dJpsiMisIDed_M);
        h_BplusM_vs_BplusM_DoubleMisIDed->Fill(LorVec_B_plus_From_DoubleMisIDed.M(), dReconstructedBM);
        h_BplusM_vs_JpsiM_DoubleMisIDed->Fill(LorVec_B_plus_From_DoubleMisIDed.M(), dJpsiMisIDed_M);

        h_Q2_Spectrum_MisIDed_Muon->Fill(LorVec_Jpsi_From_DoubleMisIDed.M2() * 1e-6);
        h_JpsiM_vs_KstplusM_MisIDed_Muon->Fill(dJpsiMisIDed_M, LorVec_K_star_plus_From_DoubleMisIDed.M());

    }

    std::cout << "[INFO]\t\tHistograms filled! Save to file:" << std::endl;

    TFile * f = new TFile(Form("%s/data/%s_MisIDs_%s.root",path_to_data.c_str(),TheDecay.c_str(), AfterBDTcut ? "AfterBDTcut" : "BeforeBDTcut"), "RECREATE");
    f->cd();
    h_Jpsi_Spectrum_MisIDed_Muon->Write("",TObject::kWriteDelete);
    h_JpsiM_vs_KstplusM_MisIDed_Muon->Write("",TObject::kWriteDelete);
    h_Q2_Spectrum_MisIDed_Muon->Write("",TObject::kWriteDelete);
    h_BplusM_vs_BplusM_DoubleMisIDed->Write("",TObject::kWriteDelete);
    h_BplusM_vs_JpsiM_DoubleMisIDed->Write("",TObject::kWriteDelete);
    f->Close();

    std::cout << "[INFO]\t\tFinished!" << std::endl;

    return 1;

}