#include "TChain.h"
#include "TFile.h"
#include "TH1D.h"
#include "TLorentzVector.h"
#include <vector>
#include <iostream>

using namespace std;

bool FindGrandMother(vector<int> *pdgId, vector<vector<int> > *parent_index, int kk, int particle, int grandmother) {
    int pdgid_mother_index = kk; //= ( *parent_index )[kk][0];
    int pdgid_mother = (*pdgId)[pdgid_mother_index];
    while (abs(pdgid_mother) == particle) {
        pdgid_mother_index = (*parent_index)[pdgid_mother_index][0];
        pdgid_mother = (*pdgId)[pdgid_mother_index];
    }

    if (abs(pdgid_mother) == grandmother) {
        return true;
    }

    return false;
}

double WHelicity(TLorentzVector lep, TLorentzVector b, TLorentzVector w) {
    lep.Boost(-w.BoostVector());
    b.Boost(-w.BoostVector());
    return -TMath::Cos(b.Angle(lep.BoostVector()));
}

int main() {
    auto t = new TChain("truth");
    t->Add("/afs/cern.ch/work/a/alhroob/private/HadronicZ_Letonictop/test.ntuple.root");

    auto h_Top_m = new TH1D("Top_m", "", 100, 140, 210);
    h_Top_m->Sumw2();
    h_Top_m->GetXaxis()->SetTitle("M_{top} [GeV]");
    h_Top_m->SetDefaultSumw2(kTRUE);

    auto h_Top_pt = new TH1D("Top_pt", "", 50, 0, 300);
    h_Top_pt->GetXaxis()->SetTitle("P_{T}(top) [GeV]");
    auto h_Top_phi = new TH1D("Top_phi", "", 50, -4, 4);
    h_Top_phi->GetXaxis()->SetTitle("#Phi(top)");
    auto h_Top_eta = new TH1D("Top_eta", "", 50, -10, 10);
    h_Top_eta->GetXaxis()->SetTitle("#eta(top)");

    auto h_W_m = new TH1D("W_m", "", 50, 60, 100);
    h_W_m->GetXaxis()->SetTitle("M_{W} [GeV]");
    auto h_W_pt = new TH1D("W_pt", "", 50, 0, 250);
    h_W_pt->GetXaxis()->SetTitle("P_{T}(W) [GeV]");
    auto h_W_phi = new TH1D("W_phi", "", 50, -4, 4);
    h_W_phi->GetXaxis()->SetTitle("#Phi(W)");
    auto h_W_eta = new TH1D("W_eta", "", 50, -7, 7);
    h_W_eta->GetXaxis()->SetTitle("#eta(W)");
    auto h_W_Helicity = new TH1D("W_Helicity", "", 50, -1, 1);
    h_W_Helicity->GetXaxis()->SetTitle("W Helicity");

    auto h_b_pt = new TH1D("b_pt", "", 50, 0, 250);
    h_b_pt->GetXaxis()->SetTitle("P_{T}(b) [GeV]");
    auto h_b_phi = new TH1D("b_phi", "", 50, -4, 4);
    h_b_phi->GetXaxis()->SetTitle("#Phi(b)");
    auto h_b_eta = new TH1D("b_eta", "", 50, -5, 5);
    h_b_eta->GetXaxis()->SetTitle("#eta(b)");

    auto h_DelR_W_b = new TH1D("DelR_W_b", "", 50, 0, 5);
    h_DelR_W_b->GetXaxis()->SetTitle("#DeltaR(W,b)");
    auto h_DelPhi_W_b = new TH1D("DelPhi_W_b", "", 50, -4, 4);
    h_DelPhi_W_b->GetXaxis()->SetTitle("#Delta#Phi(W,b)");

    auto h_Lepton_pt = new TH1D("Lepton_pt", "", 50, 0, 200);
    h_Lepton_pt->GetXaxis()->SetTitle("P_{T}(l) [GeV]");
    auto h_Lepton_phi = new TH1D("Lepton_phi", "", 50, -4, 4);
    h_Lepton_phi->GetXaxis()->SetTitle("#Phi(l)");
    auto h_Lepton_eta = new TH1D("Lepton_eta", "", 50, -5, 5);
    h_Lepton_eta->GetXaxis()->SetTitle("#eta(l)");

    auto h_Neutrino_pt = new TH1D("Neutrino_pt", "", 50, 0, 200);
    h_Neutrino_pt->GetXaxis()->SetTitle("P_{T}(#nu) [GeV]");
    auto h_Neutrino_phi = new TH1D("Neutrino_phi", "", 50, -4, 4);
    h_Neutrino_phi->GetXaxis()->SetTitle("#Phi(#nu)");
    auto h_Neutrino_eta = new TH1D("Neutrino_eta", "", 50, -5, 5);
    h_Neutrino_eta->GetXaxis()->SetTitle("#eta(#nu)");

    auto h_DelR_Lepton_Neutrino = new TH1D("DelR_Lepton_Neutrino", "", 50, 0, 5);
    h_DelR_Lepton_Neutrino->GetXaxis()->SetTitle("#DeltaR(l,#nu)");
    auto h_DelPhi_Lepton_Neutrino = new TH1D("DelPhi_Lepton_Neutrino", "", 50, -4, 4);
    h_DelPhi_Lepton_Neutrino->GetXaxis()->SetTitle("#Delta#Phi(l,#nu)");
    auto h_DelR_Lepton_b = new TH1D("DelR_Lepton_b", "", 50, 0, 5);
    h_DelR_Lepton_b->GetXaxis()->SetTitle("#DeltaR(l,b)");
    auto h_DelPhi_Lepton_b = new TH1D("DelPhi_Lepton_b", "", 50, -4, 4);
    h_DelPhi_Lepton_b->GetXaxis()->SetTitle("#Delta#Phi(l,b)");
    auto h_DelR_b_Neutrino = new TH1D("DelR_b_Neutrino", "", 50, 0, 5);
    h_DelR_b_Neutrino->GetXaxis()->SetTitle("#DeltaR(b,#nu)");
    auto h_DelPhi_b_Neutrino = new TH1D("DelPhi_b_Neutrino", "", 50, -4, 4);
    h_DelPhi_b_Neutrino->GetXaxis()->SetTitle("#Delta#Phi(b,#nu)");

    auto h_Quark1_pt = new TH1D("Quark1_pt", "", 50, 0, 250);
    h_Quark1_pt->GetXaxis()->SetTitle("P_{T}(q_{1}(Z)) [GeV]");
    auto h_Quark1_phi = new TH1D("Quark1_phi", "", 50, -4, 4);
    h_Quark1_phi->GetXaxis()->SetTitle("#Phi(q_{1}(Z))");
    auto h_Quark1_eta = new TH1D("Quark1_eta", "", 50, -5, 5);
    h_Quark1_eta->GetXaxis()->SetTitle("#eta(q_{1}(Z))");

    auto h_Quark2_pt = new TH1D("Quark2_pt", "", 100, 0, 250);
    h_Quark2_pt->GetXaxis()->SetTitle("P_{T}(q_{2}(Z)) [GeV]");
    auto h_Quark2_phi = new TH1D("Quark2_phi", "", 50, -4, 4);
    h_Quark2_phi->GetXaxis()->SetTitle("#Phi(q_{2}(Z))");
    auto h_Quark2_eta = new TH1D("Quark2_eta", "", 50, -5, 5);
    h_Quark2_eta->GetXaxis()->SetTitle("#eta(q_{2}(Z))");

    auto h_DelR_Quark1_Quark2 = new TH1D("DelR_Quark1_Quark2", "", 50, 0, 5);
    h_DelR_Quark1_Quark2->GetXaxis()->SetTitle("#DeltaR(q1,q2)");
    auto h_DelPhi_Quark1_Quark2 = new TH1D("DelPhi_Quark1_Quark2", "", 50, -4, 4);
    h_DelPhi_Quark1_Quark2->GetXaxis()->SetTitle("#Delta#Phi(q1,q2)");

    auto h_Quark_forward_pt = new TH1D("Quark_forward_pt", "", 50, 0, 300);
    h_Quark_forward_pt->GetXaxis()->SetTitle("P_{T}(q^{'})");
    auto h_Quark_forward_phi = new TH1D("Quark_forward_phi", "", 50, -4, 4);
    h_Quark_forward_phi->GetXaxis()->SetTitle("#Phi_{q^{'}}");
    auto h_Quark_forward_eta = new TH1D("Quark_forward_eta", "", 50, -5, 5);
    h_Quark_forward_eta->GetXaxis()->SetTitle("#eta_{q^{'}}  ");

    auto h_Z_m = new TH1D("Z_m", "", 50, 70, 110);
    h_Z_m->GetXaxis()->SetTitle("M_{Z} [GeV]");
    auto h_Z_pt = new TH1D("Z_pt", "", 50, 0, 250);
    h_Z_pt->GetXaxis()->SetTitle("P_{T}(Z) [GeV]");
    auto h_Z_phi = new TH1D("Z_phi", "", 50, -4, 4);
    h_Z_phi->GetXaxis()->SetTitle("#Phi(Z)");
    auto h_Z_eta = new TH1D("Z_eta", "", 50, -7, 7);
    h_Z_eta->GetXaxis()->SetTitle("#eta(Z)");

    vector<float> *mc_pt = nullptr;
    vector<float> *mc_m = nullptr;
    vector<float> *mc_eta = nullptr;
    vector<float> *mc_phi = nullptr;
    vector<int> *mc_pdgId = nullptr;
    vector<vector<double> > *mcevt_weight = nullptr;
    vector<vector<int> > *mc_parent_index = nullptr;

    t->SetBranchStatus("*", 0);
    t->SetBranchStatus("mc_pt", 1);
    t->SetBranchStatus("mc_eta", 1);
    t->SetBranchStatus("mc_phi", 1);
    t->SetBranchStatus("mc_pdgId", 1);
    t->SetBranchStatus("mcevt_weight", 1);
    t->SetBranchStatus("mc_parent_index", 1);

    t->SetBranchAddress("mc_pt", &mc_pt);
    t->SetBranchAddress("mc_m", &mc_m);
    t->SetBranchAddress("mc_eta", &mc_eta);
    t->SetBranchAddress("mc_phi", &mc_phi);
    t->SetBranchAddress("mc_pdgId", &mc_pdgId);
    t->SetBranchAddress("mcevt_weight", &mcevt_weight);
    t->SetBranchAddress("mc_parent_index", &mc_parent_index);

    for (int Event = 0; t->GetEntry(Event) != 0; ++Event) {
        if (Event % 100 == 0)cout << "Event " << Event << "  is finished\n";
        //if (Event == 5)break;
        TLorentzVector TopVec;
        TLorentzVector WVec;
        TLorentzVector bVec;
        TLorentzVector LeptonVec;
        TLorentzVector NeutrinoVec;
        TLorentzVector Quark1;
        TLorentzVector Quark2;
        TLorentzVector QuarkForward;
        TLorentzVector ZVec;
        float mc_weght = (*mcevt_weight).at(0).at(0);

        int pdgid = -999;
        int idxLep = -1, idxNeu = -1, idxb = -1, idxq1 = -1, idxq2 = -1;

        for (unsigned int kk = 0; kk < (*mc_pdgId).size(); kk++) {
            pdgid = (*mc_pdgId)[kk];
            if (abs(pdgid) == 11 || abs(pdgid) == 13 || abs(pdgid) == 15) {//may be I have to check the once
                if (FindGrandMother(mc_pdgId, mc_parent_index, (*mc_parent_index)[kk][0], 24, 6)) {
                    idxLep = kk;
                }
                continue;
            }
            if (abs(pdgid) == 12 || abs(pdgid) == 14 || abs(pdgid) == 16) {//may be I have to check the once
                if (FindGrandMother(mc_pdgId, mc_parent_index, (*mc_parent_index)[kk][0], 24, 6)) {
                    idxNeu = kk;
                }
                continue;
            }
            if (abs(pdgid) == 5 && (idxb == -1)) {//think about this
                if (abs((*mc_pdgId)[(*mc_parent_index)[kk][0]]) == 6) {
                    idxb = kk;
                    continue;
                }
            }

            if ((pdgid != 0 && abs(pdgid) < 6) && (idxq1 == -1))//check the logic
                if (abs((*mc_pdgId)[(*mc_parent_index)[kk][0]]) == 23) {
                    idxq1 = kk;
                    continue;
                }
            if ((pdgid != 0 && abs(pdgid) < 6) && (idxq2 == -1))
                if (abs((*mc_pdgId)[(*mc_parent_index)[kk][0]]) == 23) {
                    idxq2 = kk;
                }

            if (idxLep != -1 && idxNeu != -1 && idxb != -1 && idxq1 != -1 && idxq2 != -1)break;

        }

        bVec.SetPtEtaPhiM((*mc_pt)[idxb]*0.001, (*mc_eta)[idxb], (*mc_phi)[idxb], (*mc_m)[idxb]*0.001);
        LeptonVec.SetPtEtaPhiM((*mc_pt)[idxLep]*0.001, (*mc_eta)[idxLep], (*mc_phi)[idxLep], (*mc_m)[idxLep]*0.001);
        NeutrinoVec.SetPtEtaPhiM((*mc_pt)[idxNeu]*0.001, (*mc_eta)[idxNeu], (*mc_phi)[idxNeu], (*mc_m)[idxNeu]*0.001);

        TLorentzVector NeutrinoVec_new = NeutrinoVec;
        NeutrinoVec_new.SetPz(0);

        WVec = LeptonVec + NeutrinoVec;
        TopVec = WVec + bVec;

        Quark1.SetPtEtaPhiM((*mc_pt)[idxq1]*0.001, (*mc_eta)[idxq1], (*mc_phi)[idxq1], (*mc_m)[idxq1]*0.001);
        Quark2.SetPtEtaPhiM((*mc_pt)[idxq2]*0.001, (*mc_eta)[idxq2], (*mc_phi)[idxq2], (*mc_m)[idxq2]*0.001);

        TLorentzVector temp;
        if (Quark1.Pt() < Quark2.Pt()) {
            temp = Quark1;
            Quark1 = Quark2;
            Quark2 = temp;
        }


        ZVec = Quark1 + Quark2;

        if (!(*mc_parent_index).empty()) {//the following is only valid because I know that the forward quark is the 4th particle in the list
            if ((abs((*mc_pdgId)[3]) > 0) && (abs((*mc_pdgId)[3]) < 6))
                QuarkForward.SetPtEtaPhiM((*mc_pt)[3]*0.001, (*mc_eta)[3], (*mc_phi)[3], (*mc_m)[3]*0.001);
            else cout << (*mc_pdgId)[3] << endl;
        }

        h_Top_m->Fill(TopVec.M(), mc_weght);
        h_Top_pt->Fill(TopVec.Pt(), mc_weght);
        h_Top_phi->Fill(TopVec.Phi(), mc_weght);
        h_Top_eta->Fill(TopVec.Eta(), mc_weght);

        h_W_m->Fill(WVec.M(), mc_weght);
        h_W_pt->Fill(WVec.Pt(), mc_weght);
        h_W_phi->Fill(WVec.Phi(), mc_weght);
        h_W_eta->Fill(WVec.Eta(), mc_weght);

        h_b_pt->Fill(bVec.Pt(), mc_weght);
        h_b_phi->Fill(bVec.Phi(), mc_weght);
        h_b_eta->Fill(bVec.Eta(), mc_weght);

        h_DelR_W_b->Fill(WVec.DeltaR(bVec), mc_weght);
        h_DelPhi_W_b->Fill(WVec.DeltaPhi(bVec), mc_weght);

        h_Lepton_pt->Fill(LeptonVec.Pt(), mc_weght);
        h_Lepton_phi->Fill(LeptonVec.Phi(), mc_weght);
        h_Lepton_eta->Fill(LeptonVec.Eta(), mc_weght);

        h_Neutrino_pt->Fill(NeutrinoVec.Pt(), mc_weght);
        h_Neutrino_phi->Fill(NeutrinoVec.Phi(), mc_weght);
        h_Neutrino_eta->Fill(NeutrinoVec.Eta(), mc_weght);

        h_DelR_Lepton_Neutrino->Fill(LeptonVec.DeltaR(NeutrinoVec), mc_weght);
        h_DelPhi_Lepton_Neutrino->Fill(LeptonVec.DeltaPhi(NeutrinoVec), mc_weght);
        h_DelR_Lepton_b->Fill(LeptonVec.DeltaR(bVec), mc_weght);
        h_DelPhi_Lepton_b->Fill(LeptonVec.DeltaPhi(bVec), mc_weght);
        h_DelR_b_Neutrino->Fill(bVec.DeltaR(NeutrinoVec), mc_weght);
        h_DelPhi_b_Neutrino->Fill(bVec.DeltaPhi(NeutrinoVec), mc_weght);

        double W_heli = WHelicity(LeptonVec, bVec, WVec);
        h_W_Helicity->Fill(W_heli);


        h_Quark1_pt->Fill(Quark1.Pt(), mc_weght);
        h_Quark1_eta->Fill(Quark1.Eta(), mc_weght);
        h_Quark1_phi->Fill(Quark1.Phi(), mc_weght);

        h_Quark2_pt->Fill(Quark2.Pt(), mc_weght);
        h_Quark2_eta->Fill(Quark2.Eta(), mc_weght);
        h_Quark2_phi->Fill(Quark2.Phi(), mc_weght);

        h_DelR_Quark1_Quark2->Fill(Quark1.DeltaR(Quark2), mc_weght);
        h_DelPhi_Quark1_Quark2->Fill(Quark1.DeltaPhi(Quark2), mc_weght);

        h_Z_m->Fill(ZVec.M(), mc_weght);
        h_Z_pt->Fill(ZVec.Pt(), mc_weght);
        h_Z_phi->Fill(ZVec.Phi(), mc_weght);
        h_Z_eta->Fill(ZVec.Eta(), mc_weght);


        h_Quark_forward_pt->Fill(QuarkForward.Pt(), mc_weght);
        h_Quark_forward_eta->Fill(QuarkForward.Eta(), mc_weght);
        h_Quark_forward_phi->Fill(QuarkForward.Phi(), mc_weght);

        /* 
         //needs to be debugged and tuned carefully 
         vector<TLorentzVector> Jets( 3 );
         Jets[0] = Quark2;
         Jets[1] = Quark1;
         Jets[2] = bVec;
         TLorentzVector B1;
         TLorentzVector B2;
         double matric = 1e9;
         int index[3] = { 0 };
         for ( int k = 0; k < Jets.size( ); k++ ) {
         B1 = WVec + Jets[k];
         TLorentzVector Zbucket;               
         int p = 0;
         int l[2];
         for ( int m = 0; m < Jets.size( ); m++ ) {
         if ( k == m )continue;
         Zbucket += Jets[m];
         l[p] = m;
         p++;
         }
         B2 = Zbucket;
         double matric1 = B1.M( )*0.001 - 172.5;
         double matric2 = B2.M( )*0.001 - 91.2;
         double matric_new = matric1 * matric1 + matric2 * matric2;
         if ( matric_new < matric ) {
         matric = matric_new;
         index[0] = k;
         index[1]=l[0];
         index[2]=l[1];
           
         }
         cout << matric_new << "  " << matric << endl;
           
         }
         cout << index[0] << "  " << index[1] << "  " << index[2] << endl;
           
         exit( 0 );*/

    }

    auto f = new TFile("AfterShower.Herwig.Zt13TeV.root", "recreate");

    h_Top_m->Write();
    h_Top_pt->Write();
    h_Top_phi->Write();
    h_Top_eta->Write();

    h_W_m->Write();
    h_W_pt->Write();
    h_W_phi->Write();
    h_W_eta->Write();

    h_DelR_W_b->Write();
    h_DelPhi_W_b->Write();

    h_b_pt->Write();
    h_b_phi->Write();
    h_b_eta->Write();

    h_Lepton_pt->Write();
    h_Lepton_phi->Write();
    h_Lepton_eta->Write();

    h_Neutrino_pt->Write();
    h_Neutrino_phi->Write();
    h_Neutrino_eta->Write();

    h_DelR_Lepton_Neutrino->Write();
    h_DelPhi_Lepton_Neutrino->Write();
    h_DelR_Lepton_b->Write();
    h_DelPhi_Lepton_b->Write();
    h_DelR_b_Neutrino->Write();
    h_DelPhi_b_Neutrino->Write();
    h_W_Helicity->Write();

    h_Quark1_pt->Write();
    h_Quark1_eta->Write();
    h_Quark1_phi->Write();

    h_Quark2_pt->Write();
    h_Quark2_eta->Write();
    h_Quark2_phi->Write();
    h_DelR_Quark1_Quark2->Write();
    h_DelPhi_Quark1_Quark2->Write();

    h_Z_m->Write();
    h_Z_pt->Write();
    h_Z_phi->Write();
    h_Z_eta->Write();

    h_Quark_forward_pt->Write();
    h_Quark_forward_eta->Write();
    h_Quark_forward_phi->Write();

    f->Close();
    delete f;
    delete t;

    delete h_Top_m;
    delete h_Top_pt;
    delete h_Top_phi;
    delete h_Top_eta;
    delete h_W_m;
    delete h_W_pt;
    delete h_W_phi;
    delete h_W_eta;
    delete h_W_Helicity;
    delete h_b_pt;
    delete h_b_phi;
    delete h_b_eta;
    delete h_DelR_W_b;
    delete h_DelPhi_W_b;
    delete h_Lepton_pt;
    delete h_Lepton_phi;
    delete h_Lepton_eta;
    delete h_Neutrino_pt;
    delete h_Neutrino_phi;
    delete h_Neutrino_eta;
    delete h_DelR_Lepton_b;
    delete h_DelR_Lepton_Neutrino;
    delete h_DelPhi_Lepton_Neutrino;
    delete h_DelPhi_Lepton_b;
    delete h_DelR_b_Neutrino;
    delete h_DelPhi_b_Neutrino;
    delete h_Quark1_pt;
    delete h_Quark1_phi;
    delete h_Quark1_eta;
    delete h_Quark2_pt;
    delete h_Quark2_phi;
    delete h_Quark2_eta;
    delete h_DelR_Quark1_Quark2;
    delete h_DelPhi_Quark1_Quark2;
    delete h_Quark_forward_pt;
    delete h_Quark_forward_phi;
    delete h_Quark_forward_eta;
    delete h_Z_m;
    delete h_Z_pt;
    delete h_Z_phi;
    delete h_Z_eta;

    //return 0;
}