#include<TTree.h>
#include<TFile.h>
#include<TMath.h>
#include<TGraph.h>
#include<TCanvas.h>
#include<TF1.h>

/*
Double_t fit_func(Double_t* x, Double_t* par){
 return TMath::Sqrt((*x)*(*x)+par[0]*par[0]);
}*/

int main(){
 TFile* inf = new TFile("momentum_dist.root");
 //as we remember key was "momentum_dist"
 TTree* tr = (TTree*)inf->Get("momentum_dist");//void pointer cast
 // "tr" is the same tree which was saved in file
  /* in case you're not sure what's inside
 tr->Scan();//just to make sure
 tr->Scan("p_x:p_y:p_z");//only momentum dists
  */
 //now we want to have those dists and do some work e.q. finding out the correlation between p_y & p_z
 //this is done similarly to writing
 Double_t p_z=0, p_y=0;//note that variables are initialized
 Int_t count=0;
 tr->SetBranchAddress("p_z",&p_z);//notates that p_z has to store entries from "p_z" branch
 tr->SetBranchAddress("p_y",&p_y);
 tr->SetBranchAddress("index",&count);//just as an example
 Int_t num_of_entries = (Int_t)tr->GetEntries();//returns number of entries
 //Let's make a graph
 Double_t z[num_of_entries], y[num_of_entries];
 for(Int_t index = 0; index<num_of_entries; ++index){
  tr->GetEntry(index);//fills mentioned variables with corresponding entry
  z[index]=p_z;
  y[index]=p_y;
 }
 //Tree is done here, we are just about to build graph
 TGraph* gr = new TGraph(num_of_entries,y,z);
 TCanvas* C = new TCanvas("Canvas");
 gr->Draw("AC*");//direct correlation? Hmm...
 auto fit_lamb = [](Double_t* x, Double_t* par){return TMath::Sqrt((*x)*(*x)+par[0]*par[0]);};//does not work
 TF1* fit = new TF1("fit",fit_lamb,0,25,1);
 fit->SetParameter(0,0.1);//for only one parameter it seems we have to use SetParameter
 gr->Fit("fit");
 C->SaveAs("temp.pdf");
}