#include "TTree.h"
#include "TLorentzVector.h"
#include "TH1.h"
#include "TH2.h"
#include "TFile.h"
#include <vector>
#include <iostream>

//HELPFUL VARIABLES:
//******************************************
Double_t m_proton = 0.938;
Double_t m_pion = 0.139;
Double_t m_pi0 = 0.135;

TLorentzVector target;
TLorentzVector beam; //proton beam, defined a few lines below...
//******************************************

//CALULATE BEAM / TARGET 4-MOMENTA:
//******************************************
void set_beam_target_4vector(Double_t beam_mom){
    Double_t beam_energy = TMath::Sqrt( beam_mom*beam_mom + m_proton*m_proton );
    
    beam.SetPxPyPzE(0.0,0.0,beam_mom,beam_energy);
    target.SetPxPyPzE(0.0,0.0,0.0,m_proton); //the proton target is at rest....
}
//******************************************

//ACCESS TREE VALUES:
//************************************************
Double_t beam_momentum = 0.0;

//Reconstructed:
TLorentzVector *rec_proton1 = 0;
TLorentzVector *rec_proton2 = 0;
TLorentzVector *rec_pip = 0;
TLorentzVector *rec_pim = 0;
TLorentzVector *rec_photon1 = 0;
TLorentzVector *rec_photon2 = 0;

//Fitted:
//P4 3pi:
TLorentzVector *fit_2P2Pi2G_proton1 = 0;
TLorentzVector *fit_2P2Pi2G_proton2 = 0;
TLorentzVector *fit_2P2Pi2G_pip = 0;
TLorentzVector *fit_2P2Pi2G_pim = 0;
TLorentzVector *fit_2P2Pi2G_photon1 = 0;
TLorentzVector *fit_2P2Pi2G_photon2 = 0;

Double_t pull_2P2Pi2G_proton1[3];
Double_t pull_2P2Pi2G_proton2[3];
Double_t pull_2P2Pi2G_pip[3];
Double_t pull_2P2Pi2G_pim[3];
Double_t pull_2P2Pi2G_photon1[3];
Double_t pull_2P2Pi2G_photon2[3];

Double_t chi2_2P2Pi2G = 0.0;
Double_t prob_2P2Pi2G = 0.0;

//P4 2pi:
TLorentzVector *fit_2P2Pi_proton1 = 0;
TLorentzVector *fit_2P2Pi_proton2 = 0;
TLorentzVector *fit_2P2Pi_pip = 0;
TLorentzVector *fit_2P2Pi_pim = 0;

Double_t pull_2P2Pi_proton1[3];
Double_t pull_2P2Pi_proton2[3];
Double_t pull_2P2Pi_pip[3];
Double_t pull_2P2Pi_pim[3];

Double_t chi2_2P2Pi = 0.0;
Double_t prob_2P2Pi = 0.0;

//P4 3pi + m(pi0):
TLorentzVector *fit_2P2PiPi0_proton1 = 0;
TLorentzVector *fit_2P2PiPi0_proton2 = 0;
TLorentzVector *fit_2P2PiPi0_pip = 0;
TLorentzVector *fit_2P2PiPi0_pim = 0;
TLorentzVector *fit_2P2PiPi0_photon1 = 0;
TLorentzVector *fit_2P2PiPi0_photon2 = 0;

Double_t pull_2P2PiPi0_proton1[3];
Double_t pull_2P2PiPi0_proton2[3];
Double_t pull_2P2PiPi0_pip[3];
Double_t pull_2P2PiPi0_pim[3];
Double_t pull_2P2PiPi0_photon1[3];
Double_t pull_2P2PiPi0_photon2[3];

Double_t chi2_2P2PiPi0 = 0.0;
Double_t prob_2P2PiPi0 = 0.0;

//P4 3pi = m(eta):
TLorentzVector *fit_2PEta_proton1 = 0;
TLorentzVector *fit_2PEta_proton2 = 0;
TLorentzVector *fit_2PEta_pip = 0;
TLorentzVector *fit_2PEta_pim = 0;
TLorentzVector *fit_2PEta_photon1 = 0;
TLorentzVector *fit_2PEta_photon2 = 0;

Double_t pull_2PEta_proton1[3];
Double_t pull_2PEta_proton2[3];
Double_t pull_2PEta_pip[3];
Double_t pull_2PEta_pim[3];
Double_t pull_2PEta_photon1[3];
Double_t pull_2PEta_photon2[3];

Double_t chi2_2PEta = 0.0;
Double_t prob_2PEta = 0.0;
//************************************************

//Set tree variables:
//************************************************
//Incoming tree:
void set_tree_variables(TTree *inTree){
    //Reconstructed properties:
    inTree->SetBranchAddress("beam_momentum",&beam_momentum);
    inTree->SetBranchAddress("rec_proton1",&rec_proton1);
    inTree->SetBranchAddress("rec_proton2",&rec_proton2);
    inTree->SetBranchAddress("rec_pip",&rec_pip);
    inTree->SetBranchAddress("rec_pim",&rec_pim);
    inTree->SetBranchAddress("rec_photon1",&rec_photon1);
    inTree->SetBranchAddress("rec_photon2",&rec_photon2);
       
    //Fitted variables:
    //P4 3pi:
    inTree->SetBranchAddress("fit_2P2Pi2G_proton1",&fit_2P2Pi2G_proton1);
    inTree->SetBranchAddress("fit_2P2Pi2G_proton2",&fit_2P2Pi2G_proton2);
    inTree->SetBranchAddress("fit_2P2Pi2G_pip",&fit_2P2Pi2G_pip);
    inTree->SetBranchAddress("fit_2P2Pi2G_pim",&fit_2P2Pi2G_pim);
    inTree->SetBranchAddress("fit_2P2Pi2G_photon1",&fit_2P2Pi2G_photon1);
    inTree->SetBranchAddress("fit_2P2Pi2G_photon2",&fit_2P2Pi2G_photon2);
       
    inTree->SetBranchAddress("pull_2P2Pi2G_proton1",pull_2P2Pi2G_proton1);
    inTree->SetBranchAddress("pull_2P2Pi2G_proton2",pull_2P2Pi2G_proton2);
    inTree->SetBranchAddress("pull_2P2Pi2G_pip",pull_2P2Pi2G_pip);
    inTree->SetBranchAddress("pull_2P2Pi2G_pim",pull_2P2Pi2G_pim);
    inTree->SetBranchAddress("pull_2P2Pi2G_photon1",pull_2P2Pi2G_photon1);
    inTree->SetBranchAddress("pull_2P2Pi2G_photon2",pull_2P2Pi2G_photon2);
       
    inTree->SetBranchAddress("chi2_2P2Pi2G",&chi2_2P2Pi2G);
    inTree->SetBranchAddress("prob_2P2Pi2G",&prob_2P2Pi2G);
    
    //P4 2pi:
    inTree->SetBranchAddress("fit_2P2Pi_proton1",&fit_2P2Pi_proton1);
    inTree->SetBranchAddress("fit_2P2Pi_proton2",&fit_2P2Pi_proton2);
    inTree->SetBranchAddress("fit_2P2Pi_pip",&fit_2P2Pi_pip);
    inTree->SetBranchAddress("fit_2P2Pi_pim",&fit_2P2Pi_pim);
         
    inTree->SetBranchAddress("pull_2P2Pi_proton1",pull_2P2Pi_proton1);
    inTree->SetBranchAddress("pull_2P2Pi_proton2",pull_2P2Pi_proton2);
    inTree->SetBranchAddress("pull_2P2Pi_pip",pull_2P2Pi_pip);
    inTree->SetBranchAddress("pull_2P2Pi_pim",pull_2P2Pi_pim);
        
    inTree->SetBranchAddress("chi2_2P2Pi",&chi2_2P2Pi);
    inTree->SetBranchAddress("prob_2P2Pi",&prob_2P2Pi);
    
    //P4 3pi + m(pi0):
    inTree->SetBranchAddress("fit_2P2PiPi0_proton1",&fit_2P2PiPi0_proton1);
    inTree->SetBranchAddress("fit_2P2PiPi0_proton2",&fit_2P2PiPi0_proton2);
    inTree->SetBranchAddress("fit_2P2PiPi0_pip",&fit_2P2PiPi0_pip);
    inTree->SetBranchAddress("fit_2P2PiPi0_pim",&fit_2P2PiPi0_pim);
    inTree->SetBranchAddress("fit_2P2PiPi0_photon1",&fit_2P2PiPi0_photon1);
    inTree->SetBranchAddress("fit_2P2PiPi0_photon2",&fit_2P2PiPi0_photon2);
       
    inTree->SetBranchAddress("pull_2P2PiPi0_proton1",pull_2P2PiPi0_proton1);
    inTree->SetBranchAddress("pull_2P2PiPi0_proton2",pull_2P2PiPi0_proton2);
    inTree->SetBranchAddress("pull_2P2PiPi0_pip",pull_2P2PiPi0_pip);
    inTree->SetBranchAddress("pull_2P2PiPi0_pim",pull_2P2PiPi0_pim);
    inTree->SetBranchAddress("pull_2P2PiPi0_photon1",pull_2P2PiPi0_photon1);
    inTree->SetBranchAddress("pull_2P2PiPi0_photon2",pull_2P2PiPi0_photon2);
       
    inTree->SetBranchAddress("chi2_2P2PiPi0",&chi2_2P2PiPi0);
    inTree->SetBranchAddress("prob_2P2PiPi0",&prob_2P2PiPi0);
    
    //P4 3pi = m(eta):
    inTree->SetBranchAddress("fit_2PEta_proton1",&fit_2PEta_proton1);
    inTree->SetBranchAddress("fit_2PEta_proton2",&fit_2PEta_proton2);
    inTree->SetBranchAddress("fit_2PEta_pip",&fit_2PEta_pip);
    inTree->SetBranchAddress("fit_2PEta_pim",&fit_2PEta_pim);
    inTree->SetBranchAddress("fit_2PEta_photon1",&fit_2PEta_photon1);
    inTree->SetBranchAddress("fit_2PEta_photon2",&fit_2PEta_photon2);
       
    inTree->SetBranchAddress("pull_2PEta_proton1",pull_2PEta_proton1);
    inTree->SetBranchAddress("pull_2PEta_proton2",pull_2PEta_proton2);
    inTree->SetBranchAddress("pull_2PEta_pip",pull_2PEta_pip);
    inTree->SetBranchAddress("pull_2PEta_pim",pull_2PEta_pim);
    inTree->SetBranchAddress("pull_2PEta_photon1",pull_2PEta_photon1);
    inTree->SetBranchAddress("pull_2PEta_photon2",pull_2PEta_photon2);
       
    inTree->SetBranchAddress("chi2_2PEta",&chi2_2PEta);
    inTree->SetBranchAddress("prob_2PEta",&prob_2PEta);
}
//************************************************