//Classes and function used in MCtruthmatching.cpp to truthmatch the samples
//Renata Kopečná

#include "GlobalFunctions.hh"

class counters{
    public:
    int B_plus_VETO;
    int K_star_plus_VETO;
    int mu_VETO;
    int J_psi_VETO;
    int K_plus_VETO;
    int pi_zero_VETO;
    int gammas_VETO;
    int MOTHER_ID_VETO;
    int GD_MOTHER_ID_VETO;
    int GD_GD_MOTHER_ID_VETO;
    int K_short_VETO;
    int Ks_pi_VETO;
    int pi_plus_VETO;
    int Kshort_pi_zero_daughters;
    int ReconstructedPionsWhichAreMuons;

    counters(){ //default constructor
        B_plus_VETO = 0;
        K_star_plus_VETO = 0;
        mu_VETO = 0;
        J_psi_VETO = 0;
        K_plus_VETO = 0;
        pi_zero_VETO = 0;
        gammas_VETO = 0;
        MOTHER_ID_VETO = 0;
        GD_MOTHER_ID_VETO = 0;
        GD_GD_MOTHER_ID_VETO = 0;
        K_short_VETO = 0;
        Ks_pi_VETO = 0;
        pi_plus_VETO = 0;
        Kshort_pi_zero_daughters = 0;
        ReconstructedPionsWhichAreMuons = 0;
    }
    void printCounters(bool B0, bool K1);
    ~counters();   //destuctor
};

void counters::printCounters(bool B0, bool K1){
    coutInfo("The incorrect reconstructed particles are the following:");
    coutInfo("(only the first found-incorrect particle is counted)");
    coutInfo("");
    coutInfo((B0 ? "B_zero:\t\t" : "B_plus:\t\t")+to_string(B_plus_VETO));
    coutInfo((B0 ? "K_star_zero:\t" : ( K1 ? "K_one_plus:\t" : "K_star_plus:\t"))+to_string(K_star_plus_VETO));
    coutInfo("mu:\t\t"+to_string(mu_VETO));
    if(Kst2Kpluspi0Resolved){
        coutInfo("K_plus:\t\t"+to_string(K_plus_VETO));
        coutInfo("pi_zero TM:\t"+to_string(pi_zero_VETO));
        coutInfo("gamma:\t\t"+to_string(gammas_VETO));
    }
    if(Kst2Kspiplus){
        coutInfo("K_short:\t"+to_string(K_short_VETO));
        coutInfo("Ks_pi_minus:\t"+to_string(Ks_pi_VETO));
        coutInfo("pi_plus:\t"+to_string(pi_plus_VETO));
    }
    coutInfo("J_psi:\t\t"+to_string(J_psi_VETO));

    if(Kst2Kspiplus){
        coutInfo("Pi0 daughers of the Kshort:\t"+to_string(Kshort_pi_zero_daughters));
        coutInfo("Muons mis-IDed as pions:\t"+to_string(ReconstructedPionsWhichAreMuons)+"/"+to_string(Ks_pi_VETO + pi_plus_VETO));
    }

    coutInfo("Total incorrect family members of all particles: "+to_string(MOTHER_ID_VETO + GD_MOTHER_ID_VETO + GD_GD_MOTHER_ID_VETO));
    coutInfo("MotherID:\t"+to_string(MOTHER_ID_VETO));
    coutInfo("GDmotherID:\t"+to_string(GD_MOTHER_ID_VETO));
    coutInfo("GDGDmotherID:\t"+to_string(GD_GD_MOTHER_ID_VETO));
    return;
}

counters::~counters(){ //destuctor
    B_plus_VETO = 0;
    K_star_plus_VETO = 0;
    mu_VETO = 0;
    J_psi_VETO = 0;
    K_plus_VETO = 0;
    pi_zero_VETO = 0;
    gammas_VETO = 0;
    MOTHER_ID_VETO = 0;
    GD_MOTHER_ID_VETO = 0;
    K_short_VETO = 0;
    Ks_pi_VETO = 0;
    pi_plus_VETO = 0;
    Kshort_pi_zero_daughters = 0;
    ReconstructedPionsWhichAreMuons = 0;
}

bool isBKGCAT(int BKGCAT, bool ReferenceChannel){
    if (Kst2Kspiplus){
        if(ReferenceChannel){
            if(BKGCAT == 0 || BKGCAT == 50) return true;
            else return false;
        }
        else{
            if(BKGCAT == 10 || BKGCAT == 50) return true;
            else return false;
        }
    }
    else{
        if(ReferenceChannel){
            if(BKGCAT == 0 || BKGCAT == 40 || BKGCAT == 50) return true;
            else return false;
        }
        else{
            if(BKGCAT == 10 || BKGCAT == 40 || BKGCAT == 50) return true;
            else return false;
        }
    }
}

//TRUED B+
class B_plus_IDs{
    public:
    int B_plus_TRUEID;
    B_plus_IDs(){
         B_plus_TRUEID = 0;
    }
    ~B_plus_IDs(){
        B_plus_TRUEID = 0;
    }
    bool isBplus(counters *counter, bool B_plus_TM){
        if (!B_plus_TM) return true;
        if(TMath::Abs(B_plus_TRUEID) != TRUEID.B_PLUS){
            counter->B_plus_VETO++;
            return false;
        }
        else return true;
    }
};

//TRUEID K*+
class K_star_IDs{
    public:
    int K_star_plus_TRUEID;
    int K_star_plus_MOTHER_ID;
    K_star_IDs(){
        K_star_plus_TRUEID = 0;
        K_star_plus_MOTHER_ID = 0;
    }
    ~K_star_IDs(){
        K_star_plus_TRUEID = 0;
        K_star_plus_MOTHER_ID = 0;
    }

    bool isKstar(counters *counter, bool K_star_plus_TM, bool B_plus_TM){
        if (K_star_plus_TM && TMath::Abs(K_star_plus_TRUEID) != TRUEID.K_STAR_PLUS){
            counter->K_star_plus_VETO++;
            return false;
        }
        if (B_plus_TM && TMath::Abs(K_star_plus_MOTHER_ID) != TRUEID.B_PLUS){
            counter->MOTHER_ID_VETO++;
            return false;
        }
        return true;
    }
};

//TRUEID K+
class K_plus_IDs{
    public:
    int K_plus_TRUEID;
    int K_plus_MOTHER_ID;
    int K_plus_GD_MOTHER_ID;

    K_plus_IDs(){
        K_plus_TRUEID = 0;
        K_plus_MOTHER_ID = 0;
        K_plus_GD_MOTHER_ID = 0;
    }
    ~K_plus_IDs(){
        K_plus_TRUEID = 0;
        K_plus_MOTHER_ID = 0;
        K_plus_GD_MOTHER_ID = 0;
    }

    bool isKplus(counters *counter, bool K_plus_TM, bool K_star_plus_TM, bool B_plus_TM){// flags *allFlags){
        if (K_plus_TM && TMath::Abs(K_plus_TRUEID) != TRUEID.K_PLUS){
            counter->K_plus_VETO++;
            return false;
        }
        if (K_star_plus_TM && TMath::Abs(K_plus_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            counter->MOTHER_ID_VETO++;
            return false;
        }
        if (B_plus_TM && TMath::Abs(K_plus_GD_MOTHER_ID) != TRUEID.B_PLUS){
            counter->GD_MOTHER_ID_VETO++;
            return false;
        }
        return true;
    }
};

//TRUEID pi0
class pi_zero_IDs{
public:
    int pi_zero_TRUEID;
    int pi_zero_MOTHER_ID;
    int pi_zero_GD_MOTHER_ID;

    pi_zero_IDs(){
        pi_zero_TRUEID = 0;
        pi_zero_MOTHER_ID = 0;
        pi_zero_GD_MOTHER_ID = 0;
    }
    ~pi_zero_IDs(){
        pi_zero_TRUEID = 0;
        pi_zero_MOTHER_ID = 0;
        pi_zero_GD_MOTHER_ID = 0;
    }

    bool isPi0(counters *counter, bool pi_zero_TM, bool K_star_plus_TM, bool B_plus_TM, double pi_zero_resolved_M, double B_plus_M, TH1D *h_pi0_mismatched_mass, TH1D *h_Bplus_mass_pi0mismatched, TH1D *h_pi0_resolved_mass){

        if (pi_zero_TM){
            if(TMath::Abs(pi_zero_TRUEID) == TRUEID.PI_ZERO){
                h_pi0_resolved_mass->Fill(pi_zero_resolved_M);
            }
            else{
                counter->pi_zero_VETO++;
                h_pi0_mismatched_mass->Fill(pi_zero_resolved_M);
                h_Bplus_mass_pi0mismatched->Fill(B_plus_M);
                return false;
            }
        }
        if (K_star_plus_TM && TMath::Abs(pi_zero_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            counter->MOTHER_ID_VETO++;
            return false;
        }
        if (B_plus_TM && TMath::Abs(pi_zero_GD_MOTHER_ID) != TRUEID.B_PLUS){
            counter->GD_MOTHER_ID_VETO++;
            return false;
        }
        return true;
    }
};

//TRUEID gamma
class gamma_IDs{
public:
    int gamma_TRUEID;
    int gamma_MOTHER_ID;
    int gamma_GD_MOTHER_ID;
    int gamma_GD_GD_MOTHER_ID;

    gamma_IDs(){
        gamma_TRUEID = 0;
        gamma_MOTHER_ID = 0;
        gamma_GD_MOTHER_ID = 0;
        gamma_GD_GD_MOTHER_ID = 0;
    }
    ~gamma_IDs(){
        gamma_TRUEID = 0;
        gamma_MOTHER_ID = 0;
        gamma_GD_MOTHER_ID = 0;
        gamma_GD_GD_MOTHER_ID = 0;
    }

    //Check if gamma was converted
    bool isGammaConversion(counters *counter, bool gamma_TM, bool pi_zero_TM, bool K_star_plus_TM){
        if(gamma_TM && TMath::Abs(gamma_TRUEID) != TRUEID.ELECTRON){
            ++counter->gammas_VETO;
            return false;
        }
        if(gamma_TM && TMath::Abs(gamma_MOTHER_ID) != TRUEID.GAMMA){
            ++counter->gammas_VETO;
            return false;
        }
        if(pi_zero_TM && TMath::Abs(gamma_GD_MOTHER_ID) != TRUEID.PI_ZERO){
            ++counter->MOTHER_ID_VETO;
            return false;
        }
        if(K_star_plus_TM && TMath::Abs(gamma_GD_GD_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            ++counter->GD_MOTHER_ID_VETO;
            return false;
        }
        return true;
    }    

    //TRUEID gamma
    bool isGammaTrue(counters *counter, bool gamma_TM,bool gamma_TM_full, bool pi_zero_TM, bool K_star_plus_TM, bool B_plus_TM){
        if (!gamma_TM_full) return true;
        if(gamma_TM && TMath::Abs(gamma_TRUEID) != TRUEID.GAMMA){
            ++counter->gammas_VETO;
            return false;
        }

        if (pi_zero_TM && TMath::Abs(gamma_MOTHER_ID) != TRUEID.PI_ZERO){
            ++counter->MOTHER_ID_VETO;
            return false;
        }

        if (K_star_plus_TM && TMath::Abs(gamma_GD_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            ++counter->GD_MOTHER_ID_VETO;
            return false;
        }

        if (B_plus_TM && TMath::Abs(gamma_GD_GD_MOTHER_ID) != TRUEID.B_PLUS){
            ++counter->GD_GD_MOTHER_ID_VETO;
            return false;
        }
        return true;
    }
};

//TRUEID Ks
class K_short_IDs{
public:
    int K_short_TRUEID;
    int K_short_MOTHER_ID;
    int K_short_GD_MOTHER_ID;

    K_short_IDs(){
        K_short_TRUEID = 0;
        K_short_MOTHER_ID = 0;
        K_short_GD_MOTHER_ID = 0;
    }
    ~K_short_IDs(){
        K_short_TRUEID = 0;
        K_short_MOTHER_ID = 0;
        K_short_GD_MOTHER_ID = 0;
    }

    bool isKshort(counters *counter){
        if(TMath::Abs(K_short_TRUEID) != TRUEID.K_SHORT){
            counter->K_short_VETO++;
            return false;
        }

        if(TMath::Abs(K_short_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            counter->MOTHER_ID_VETO++;
            return false;
        }

        if(TMath::Abs(K_short_GD_MOTHER_ID) != TRUEID.B_PLUS){
            counter->GD_MOTHER_ID_VETO++;
            return false;
        }
        return true;
    }
};

//TRUEID pi+
class pi_plus_IDs{
public:
    int pi_plus_TRUEID;
    int pi_plus_MOTHER_ID;
    int pi_plus_GD_MOTHER_ID;

    pi_plus_IDs(){
        pi_plus_TRUEID = 0;
        pi_plus_MOTHER_ID = 0;
        pi_plus_GD_MOTHER_ID = 0;
    }
    ~pi_plus_IDs(){
        pi_plus_TRUEID = 0;
        pi_plus_MOTHER_ID = 0;
        pi_plus_GD_MOTHER_ID = 0;
    }

    bool isPiPlus(counters *counter){
        if(TMath::Abs(pi_plus_TRUEID) != TRUEID.PI_PLUS){
            if(TMath::Abs(pi_plus_TRUEID) == TRUEID.MU_MINUS) counter->ReconstructedPionsWhichAreMuons++;
            counter->pi_plus_VETO++;
            return false;
        }

        if(TMath::Abs(pi_plus_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            counter->MOTHER_ID_VETO++;
            return false;
        }
        if(TMath::Abs(pi_plus_GD_MOTHER_ID) != TRUEID.B_PLUS){
            counter->GD_MOTHER_ID_VETO++;
            return false;
        }
        return true;
    }
};

//TRUEID pions coming from Ks
class Ks_pi_IDs{
public:
    int Ks_pi_TRUEID;
    int Ks_pi_MOTHER_ID;
    int Ks_pi_GD_MOTHER_ID;
    int Ks_pi_GD_GD_MOTHER_ID;

    Ks_pi_IDs(){
        Ks_pi_TRUEID = 0;
        Ks_pi_MOTHER_ID = 0;
        Ks_pi_GD_MOTHER_ID = 0;
        Ks_pi_GD_GD_MOTHER_ID = 0;
    }
    ~Ks_pi_IDs(){
        Ks_pi_TRUEID = 0;
        Ks_pi_MOTHER_ID = 0;
        Ks_pi_GD_MOTHER_ID = 0;
        Ks_pi_GD_GD_MOTHER_ID = 0;
    }

    bool isKsPi(counters *counter){

        if(TMath::Abs(Ks_pi_TRUEID) != TRUEID.PI_PLUS){
            if(TMath::Abs(Ks_pi_TRUEID) == TRUEID.PI_ZERO)counter->Kshort_pi_zero_daughters++;
            if(TMath::Abs(Ks_pi_TRUEID) == TRUEID.MU_MINUS)counter->ReconstructedPionsWhichAreMuons++;
            counter->Ks_pi_VETO++;
            return false;
        }
        if(TMath::Abs(Ks_pi_MOTHER_ID) != TRUEID.K_SHORT){
            counter->MOTHER_ID_VETO++;
            return false;
        }
        if(TMath::Abs(Ks_pi_GD_MOTHER_ID) != TRUEID.K_STAR_PLUS){
            counter->GD_MOTHER_ID_VETO++;
            return false;
        }
        if(TMath::Abs(Ks_pi_GD_GD_MOTHER_ID) != TRUEID.B_PLUS){
            counter->GD_GD_MOTHER_ID_VETO++;
            return false;
        }
        return true;

    }
};

//TRUEID muon
class mu_IDs{
public:
    int mu_TRUEID;
    int mu_MOTHER_ID;
    int mu_GD_MOTHER_ID;
    bool ReferenceChannel = false;

    mu_IDs(bool isRefChan){
        mu_TRUEID = 0;
        mu_MOTHER_ID = 0;
        mu_GD_MOTHER_ID = 0;
        ReferenceChannel = isRefChan;
    }
    ~mu_IDs(){
        mu_TRUEID = 0;
        mu_MOTHER_ID = 0;
        mu_GD_MOTHER_ID = 0;
        ReferenceChannel = 0;
    }

    bool isMu(counters *counter, bool mu_TM, bool J_psi_TM, bool B_plus_TM){//flags *allFlags){
        if (mu_TM && TMath::Abs(mu_TRUEID) != TRUEID.MU_MINUS){
            counter->mu_VETO++;
            return false;
        }
        if(ReferenceChannel){
            if (J_psi_TM && TMath::Abs(mu_MOTHER_ID) != TRUEID.J_PSI){
                counter->MOTHER_ID_VETO++;
                return false;
            }
            if(B_plus_TM && TMath::Abs(mu_GD_MOTHER_ID) != TRUEID.B_PLUS){
                counter->GD_MOTHER_ID_VETO++;
                return false;
            }
        }
        else{
            if(B_plus_TM && TMath::Abs(mu_MOTHER_ID) != TRUEID.B_PLUS){
                counter->MOTHER_ID_VETO++;
                return false;

            }
        }
        return true;
    }
};

//TRUEID J/psi
class J_psi_IDs{
public:
    int J_psi_TRUEID;
    int J_psi_MOTHER_ID;
    int J_psi_GD_MOTHER_ID;
    bool ReferenceChannel = false;

    J_psi_IDs(bool isRefChan){
        J_psi_TRUEID = 0;
        J_psi_MOTHER_ID = 0;
        J_psi_GD_MOTHER_ID = 0;
        ReferenceChannel = isRefChan;
    }
    ~J_psi_IDs(){
        J_psi_TRUEID = 0;
        J_psi_MOTHER_ID = 0;
        J_psi_GD_MOTHER_ID = 0;
        ReferenceChannel = 0;
    }

    bool isJpsi(counters *counter, bool J_psi_TM, bool B_plus_TM){
        if(ReferenceChannel){
            if(J_psi_TM && TMath::Abs(J_psi_TRUEID) != TRUEID.J_PSI){
                counter->J_psi_VETO++;
                return false;
            }
            if(B_plus_TM && TMath::Abs(J_psi_MOTHER_ID) != TRUEID.B_PLUS){
                counter->MOTHER_ID_VETO++;
                return false;
            }
        }
        else{
            if(J_psi_TM && TMath::Abs(J_psi_TRUEID) == TRUEID.J_PSI){
                counter->J_psi_VETO++;
                return false;
            }
        }
        return true;
    }

};