// The purpose of this tutorial is to write a ROOT macro that receives a ROOT file with
// all the particles of an event and returns another ROOT file with only the Z boson that
// will be reconstructed.

void ROOT_Intermediate_part1_simpler(){

    // read input file
    TFile *fInput = new TFile("zjet_unrec.root", "read");
    
    // load input tree
    TTree *t = (TTree*) fInput->Get("Tdata");

    // declare TH1 object for Z boson mass
    TH1D *h = new TH1D("Zmass", "Z boson mass", 100, 0., 150);

    // declare the variables to assess each event
    vector<int> *id = 0;
	vector<double> *m = 0;
	vector<double> *px = 0;
	vector<double> *py = 0;
	vector<double> *pz = 0;
	vector<double> *E = 0;

    // link these variables to the tree
    t->SetBranchAddress("id", &id);
    t->SetBranchAddress("px", &px);
    t->SetBranchAddress("py", &py);
    t->SetBranchAddress("pz", &pz);
    t->SetBranchAddress("m", &m);
    t->SetBranchAddress("En", &E);



    // create output file
    TFile *fOutput = new TFile("zjet_filtered.root", "recreate");

    // create output tree
    TTree *tOutput = new TTree("Tdata_output", "Tdata_output");

    // declare the variables for output tree
	double m_output = 0;
	double px_output = 0;
	double py_output = 0;
	double pz_output = 0;
	double E_output = 0;

    // link these variables to the tree
    tOutput->Branch("px", &px_output);
    tOutput->Branch("py", &py_output);
    tOutput->Branch("pz", &pz_output);
    tOutput->Branch("m", &m_output);
    tOutput->Branch("En", &E_output);

    // loop through tree
    int NEvents = t->GetEntries();
    for(int i = 0; i <= NEvents; i++){

        t->GetEntry(i);

        // loop through every particle and find muon and antimuon
        int i_mu_p = 0, i_mu_n = 0;
        int flag_mu_p = 0, flag_mu_n = 0;

        for(int iObj = 0; iObj <= (*id).size(); iObj++){
            //muon
            if((*id)[iObj]==13){
                i_mu_p = iObj;
                flag_mu_p++;
            }
            //antimuon
            if((*id)[iObj]==-13){
                i_mu_n = iObj;
                flag_mu_n++;
            }
            
        }

        // check if we found both and calculate invariant mass
        if(flag_mu_p && flag_mu_n){

            // calculate invariant mass
            double Et = (*E)[i_mu_p] + (*E)[i_mu_n];
            double pxt = (*px)[i_mu_p] + (*px)[i_mu_n];
            double pyt = (*py)[i_mu_p] + (*py)[i_mu_n];
            double pzt = (*pz)[i_mu_p] + (*pz)[i_mu_n];

            double invar_mass = sqrt(Et*Et - (pxt*pxt + pyt*pyt + pzt*pzt));
            h->Fill(invar_mass);

            // save event to new tree
            E_output = Et;
            px_output = pxt;
            py_output = pyt;
            pz_output = pzt;
            m_output = invar_mass;
            tOutput->Fill();


        }
        
    }

    // verify spectrum
    new TCanvas();
    h->Draw("hist");

    // save new tree to file
    tOutput->Write();
    fOutput->Close();

    return;
}