#include "Riostream.h"
#include "TString.h"
#include "TH1D.h"
#include "TCanvas.h"
#include "TTree.h"
#include "TFile.h"
#include "TLorentzVector.h"

using namespace std;

void makeExercise(){

    string file_name ="data_file.txt";
    ifstream file_in(file_name);
    
    if(!file_in.is_open()) {
        cout << "Cannot open data file!" << endl;
        return;
    }
    
    double k_px,  k_py,  k_pz;
    double pi_px, pi_py, pi_pz;
    double B_m, B_de; // the variables that I want to calculate
    
    TTree* dataTree = new TTree("dataTree","B0toKpi data");
    dataTree->Branch("k_px",&k_px,"k_px/D");
    dataTree->Branch("k_py",&k_py,"k_py/D");
    dataTree->Branch("k_pz",&k_pz,"k_pz/D");
    dataTree->Branch("pi_px",&pi_px,"pi_px/D");
    dataTree->Branch("pi_py",&pi_py,"pi_py/D");
    dataTree->Branch("pi_pz",&pi_pz,"pi_pz/D");
    //add the two new variables to the tree
    dataTree->Branch("B_m",&B_m,"B_m/D");
    dataTree->Branch("B_de",&B_de,"B_de/D");
    
    //Let's define the histograms to look at the distributions
    TH1D* h_m = new TH1D("h_m"," ",40,5.25,5.30);
    TH1D* h_de = new TH1D("h_de"," ",40,-0.15,0.15);
    
    //usefull constants
    const double pi_m = 0.13957018; //pion mass in GeV/c2
    const double k_m = 0.493667;    //kaon mass in GeV/c2
    const double sqs = 10.5794;     //cms energy in GeV (Y(4S) mass...)
        
    while(file_in.is_open()){

        file_in >> k_px  >> k_py  >> k_pz
                >> pi_px >> pi_py >> pi_pz;
               
        if(file_in.eof()) break;
        
        //define the 4-momentum of the pion and the kaon
        TLorentzVector pi_p, k_p;
        pi_p.SetXYZM(pi_px,pi_py,pi_pz,pi_m);//set the components for the pion
        k_p.SetXYZM(k_px,k_py,k_pz,k_m); //and for the kaon

        TLorentzVector B_p = pi_p+k_p;//the B is the sum of the pion and kaon

        B_de = B_p.E() - sqs/2; //easy to get the energy
        B_m = sqrt( sqs*sqs/4 - B_p.Vect().Mag2() ); //and the mass
        
        //fill my histograms
        h_m->Fill(B_m);
        h_de->Fill(B_de);
        
        //fill the tree
        dataTree->Fill();
        
    }
    
    file_in.close();
    

    cout << "Candidates in the tree: " << dataTree->GetEntries() << endl;
   
    //save everything in a file
    TFile* dataFile = new TFile("data_B0toKpi.root","RECREATE");
    dataTree->Write();
    h_m->Write();
    h_de->Write();
    dataFile->Close();
    
    //let's make some plot
    gStyle->SetOptStat(1110);//this is a global style set
    TCanvas* c1 = new TCanvas("c1","c1",1600,600);
    c1->Divide(2,1);//I split my canvas into two part (called pad)
    c1->cd(1);//and go into the first pad
    h_m->GetXaxis()->SetTitle("m(K#pi) [GeV/c^{2}]"); //set title x
    h_m->GetYaxis()->SetTitle(Form("Candidates per %.1f [Mev/c^{2}]",
                                   1.e3*h_m->GetXaxis()->GetBinWidth(1)));//title y
    h_m->GetYaxis()->SetRangeUser(0,1400); //set the interval to draw in y
    h_m->Draw(); // and draw

    c1->cd(2);//go to the second pad, and draw the other histogram
    h_de->GetXaxis()->SetTitle("#DeltaE [GeV]");
    h_de->GetYaxis()->SetTitle(Form("Candidates per %.1f [Mev]",
                                   1.e3*h_de->GetXaxis()->GetBinWidth(1)));
    h_de->GetYaxis()->SetRangeUser(0,1400);
    h_de->Draw();
        
    return;

}