// 24 luglio 2007
//----------------------------
//	anna.parravicini@cnao.it
//----------------------------

//routine per plottare distribuzione particelle
//nello spazio fasi ed ellisse di Twiss

//>>>> Difference with previous versions:

//*the reference ellipse is plotted together with the RMS ellipse
//*geometrical emittance plot added
//*TProfile instead of TH2F, in order to have the average z-entry
// rather than the sum of z-entries.
//*x-bins [mm] are centred on the slit median axis width. namely,
// the bin is drawn from xo-(slit step)/2 to xo+(slit step)/2
//*main parameters written on the graphs

//>>>> To be done:

// * aggiungere in modo facilmente toglibile la possibilita' di 
//  plottare tutto decentrato di x_center, xp_center, cosi' che sia
//  piu' realistico (il centramento viene fatto via software solo per 
//  calcolare l'ellisse RMS)

//----------------------------------------------------------

#include "TGraph.h"
#include "TStyle.h"
#include "TMath.h"
#include "TCanvas.h"
#include "TTree.h"
#include "TProfile2D.h" //user's guide, page 44. See prova_TH2F_TProfile2D.C per esempio.

int Main(){

	// === VARIABLES DECLARATION ===================================================
	Int_t Nentries; //events number per branch

	Float_t pos[50000];//[mm] - array da riempire con Branches
	Float_t div[50000];//[mrad] - array da riempire con Branches
	Float_t charge[50000];//[arbitrary units] - array da riempire con Branches

	Float_t pos_t=0;//variabili d'appoggio per leggere Tree
    Float_t div_t=0;
	Float_t charge_t=0;

	Double_t Xmax;//extreme values for x and xp arrays
	Double_t Xmin;
	Double_t XPmax;
	Double_t XPmin;
	Double_t Imax;
	Double_t Imin;

	//RMS ellipse parameters
    Double_t alpha;
    Double_t beta;
    Double_t gamma;
    Double_t emi;
    Double_t limit;

	Int_t n=44;
	Double_t ell[88];
    Double_t x[88];
    Double_t step;

	//reference ellipse parameters
	Double_t alpha_ref;
    Double_t beta_ref;
    Double_t gamma_ref;
    Double_t emi_ref;
    Double_t limit_ref;

	Double_t ell_ref[88];
    Double_t x_ref[88];
    Double_t step_ref;

	//geometrical ellipse threshold
	Double_t Ithreshold;
	Bool_t overloop; //in case it doesn't reach the required accuracy
	Float_t P; //we want compute 90% emittance
	Float_t Peff; //effective percentage of the emittance computation
	Double_t emiP;//[mm*mrad] geometrical emittance output value
	Double_t emiErr;//[mm*mrad] geometrical emittance error

	//scanning parameters
	Int_t sl_Nsteps;//number of slit steps
	Float_t sl_step;//slit step amplitude
	Int_t ws_points;// ws samples per profile, at constant velocity

	//general settings
	gStyle->SetOptStat(kFALSE);//not to plot the statistic box
	gStyle->SetPalette(1);

	// === CREO UN CANVAS in cui fare i plot e lo divido in 3 =======================

	TCanvas *c=new TCanvas ("c", "Phase Space Plots", 800, 800);
	c->Divide(2,2);

	// ==== CARICO PUNTI SP FASI DA FILE (Il file deve stare in C:\root) =============

	// --- from file to the tree-------------------------------

	// per fare histo 2D da file con 3 colonne, di cui la terza e' il "peso"
	// da dare alla cella
	TTree *t=new TTree ("t", "albero per caricare file di 3 colonne");
	t->ReadFile("meas_distribution.txt", "x:xp:I");

	Nentries=(Int_t)t->GetEntries(); //numero di eventi per ramo
	Xmax=t->GetMaximum("x");
	Xmin=t->GetMinimum("x");
	XPmax=t->GetMaximum ("xp");
	XPmin=t->GetMinimum ("xp");
	Imax=t->GetMaximum ("I");
	Imin=t->GetMinimum ("I");
	//printf ("Imax=%d \n", Imax);

	// --- from tree to the arrays------------------------------

	//se mi bastano histo e non voglio grafico,
	//non e' necessario leggere l'albero e riempire gli array
	
	//before reading branches, we must give them an address
	t->SetBranchAddress("x",&pos_t);
	t->SetBranchAddress("xp",&div_t);
	t->SetBranchAddress("I",&charge_t);
	
	//fill the array from the branches
	for (Int_t j=0;j<Nentries;j++){
		t->GetEntry(j);
		pos[j]=pos_t;
		div[j]=div_t;
		charge[j]=charge_t;
	}

	// === CARICO PARAMETRI DI TWISS e preparo i punti per disegnare l'ellisse ==========

	ifstream infile("twisspara.txt");
    Double_t param[17];//alpha, beta, emittance,
			//alpha_ref, beta_ref, emitt_ref,
			//x_c, xp_c (distribution barycenter): I don't use it at now
			//Ithreshold
	//sl_Nsteps=number of slit steps
	//sl_step=slit step amplitude
	//ws_points=ws samples per profile, at constant velocity
	//overloop=boolean variable saying if the geomtrical emittance computation was concluded before
	//			reaching the required accuracy

    for (Int_t i=0;i<17;i++)
	   infile>>param[i];
	
	//RMS ellipse parameters
    alpha=param[0];
    beta=param[1];
    gamma=(1+alpha*alpha)/beta;
    emi=param[2];
    limit=sqrt(emi*beta);

	//reference ellipse parameters
	alpha_ref=param[3];
    beta_ref=param[4];
    gamma_ref=(1+alpha_ref*alpha_ref)/beta_ref;
    emi_ref=param[5];
    limit_ref=sqrt(emi_ref*beta_ref);

	//scanning settings
	sl_Nsteps=param[9];
	sl_step=param[10];
	ws_points=param[11];

	//geometrical ellipse
	Ithreshold=param[8];
	
	overloop=param[12];
	if (overloop)
		cout<<"Geom emitt is far from user's required accuracy."<<endl; //da aggiungere anche sul plot
	
	P=param[13]*100;// %
	Peff=param[14]*100;// %
	emiP=param[15];//[mm*mrad]
	emiErr=param[16];//[mm*mrad]

	//writing the RMS ellipse 
	step=limit/(n/2);
   
    for (Int_t i=0;i<n;i++){ //riempie gli array con i punti dell' ellisse
	   x[i]=-limit+step*i;
	   ell[i]=(-2*alpha*x[i]+sqrt((2*alpha*x[i])*(2*alpha*x[i])-4*beta*(gamma*x[i]*x[i]-emi)))/(2*beta);
 	   x[i+n]=limit-step*i;
	   ell[i+n]=(-2*alpha*x[i+n]-sqrt((2*alpha*x[i+n])*(2*alpha*x[i+n])-4*beta*(gamma*x[i+n]*x[i+n]-emi)))/(2*beta);
    }
   	
	//writing the reference ellipse
	step_ref=limit_ref/(n/2);
   
    for (Int_t i=0;i<n;i++){ //riempie gli array con i punti dell' ellisse
	   x_ref[i]=-limit_ref+step_ref*i;
	   ell_ref[i]=(-2*alpha_ref*x_ref[i]+sqrt((2*alpha_ref*x_ref[i])*(2*alpha_ref*x_ref[i])-4*beta_ref*(gamma_ref*x_ref[i]*x_ref[i]-emi_ref)))/(2*beta_ref);
 	   x_ref[i+n]=limit_ref-step_ref*i;
	   ell_ref[i+n]=(-2*alpha_ref*x_ref[i+n]-sqrt((2*alpha_ref*x_ref[i+n])*(2*alpha_ref*x_ref[i+n])-4*beta_ref*(gamma_ref*x_ref[i+n]*x_ref[i+n]-emi_ref)))/(2*beta_ref);
    }
   	
	// === costruiamo i GRAFICI che RIPRODUCANO l'ELLISSE ===============================
	
	TGraph *gEll=new TGraph(2*n,x,ell);
	gEll->SetLineColor(1);
	gEll->SetLineWidth(2);
	
	TGraph *gEllRef=new TGraph(2*n,x_ref,ell_ref);
	gEllRef->SetLineColor(1);
	gEllRef->SetLineStyle(2);//2=dash, 4=dash-dot
	gEllRef->SetLineWidth(2);

	// === PHASE SPACE DISTRIBUTION (x, xp, I) PLOT ===============================================================
	c->cd(1);

	TProfile2D *h3D = new TProfile2D("h3d","#font[42]{Phase space distribution}", sl_Nsteps,Xmin-sl_step/2, Xmax+sl_step/2,ws_points, XPmin, XPmax+0.0001*XPmax);
	//TH2F *h3d = new TH2F("h3d","#font[42]{Phase space distribution}", sl_steps,Xmin, Xmax,ws_points, XPmin, XPmax);
	
	//TH2F plots the sum of all the z-entries of each bin;
	//TProfile2D plots the average of all the z-entries of each bin.

	for (int j=0;j<Nentries;j++)
		h3d->Fill(pos[j],div[j],charge[j]);
	h3d->GetYaxis()->SetRange(ws_points/3,ws_points*2/3); //range di bin da plottare
	h3d->GetXaxis()->SetTitle("#font[42]{x[mm]}");
	h3d->GetYaxis()->SetTitle("#font[42]{xp[mrad]}");
	h3d->Draw("lego2z fb bb");
	//fb=remove front surface; bb=remove back surface of the cartesian axis

	// === PHASE SPACE TOP-VIEW WITH RMS AND REFERENCE ELLIPSES  ==================================================
	
	//cout <<"hello1"<<endl;

	c->cd(3);

	// histo 2D, with weigth I for each point.
	// If more points have the same x and xp, their I-values are sum up; it 
	// follows that a I value larger than Imax is displayed: it's due to the binning!
	TProfile2D *h = new TProfile2D("h","#font[42]{Statistical emittance and ref ellipse}", sl_Nsteps,Xmin-sl_step/2, Xmax+sl_step/2,ws_points, XPmin, XPmax+0.001*XPmax);
	for (int j=0;j<Nentries;j++)
		h->Fill(pos[j],div[j],charge[j]);
	h->GetYaxis()->SetRange(ws_points/3,ws_points*2/3); //range di bin da plottare
	h->GetXaxis()->SetTitle("#font[42]{x[mm]}");
	h->GetYaxis()->SetTitle("#font[42]{xp[mrad]}");
	gStyle->SetOptStat(kFALSE);//not to plot the statistic box
	h->Draw("colz");

	// add the legend
	leg = new TLegend(0.47,0.11,0.9,0.25);
	leg->AddEntry(gEll,"RMS ellipse","l");//"l" use the line that markes the gEll graph
	leg->AddEntry(gEllRef,"Reference ellipse","l");
	leg->Draw();

	//Add statistical emitt value on the plot
	TPaveText *ptrms=new TPaveText (0.20,0.81,0.9,0.91,"NDC");
	//NDC: normalized coordinates (x1,y1,x2,y2)
	char text[100];
	sprintf (text,"#font[42]{#splitline{alpha RMS= %.3g - beta RMS= %.3g m}{Emitt RMS= %.3g mm#upointmrad}}", alpha, beta,emi);
	//mancano unita' di misura e errore, if any
	TText *trms=ptrms->AddText(text);
	trms->SetTextColor(1);
	trms->SetTextSize(0.04);
	trms->SetTextAlign(12);//=H and V centred

	// t->Draw("xp:x","I","colz"); 
	// if it is drawn directly from the tree, a x-xp binning is performed
	// out of the user's control. Passing by an histo, more options are available

	//cout <<"hello2"<<endl;

	gEll->Draw("C");//non riportiamo gli assi perche' gia' presenti. 
	gEllRef->Draw("C");
	ptrms->Draw();

	// === GEOMETRICAL EMITTANCE PLOT ============================================================================
	c->cd(4);
	
	// param[8]=Ithreshold, namely, all the bins with current >Ithreshold
	// contribute to the geometrical emittance, at the required percentage.
	// I would like a top-view of the 3D phase-space plot, with a different fill color
	// for bins over and under threshold
	
	// Only pints with I >Ithreshold are taken into account

	char hgeom_name[40];
	sprintf(hgeom_name,"#font[42]{Geometrical emittance surface @ %.0f%%}",P);

	TProfile2D *hgeom = new TProfile2D("hgeom",hgeom_name,sl_Nsteps,Xmin-sl_step/2, Xmax+sl_step/2,ws_points, XPmin, XPmax+0.001*XPmax);
	for (int j=0;j<Nentries;j++){
		if (charge[j]>Ithreshold)
			hgeom->Fill(pos[j],div[j],charge[j]);
	}
	hgeom->GetYaxis()->SetRange(ws_points/3,ws_points*2/3); //range di bin da plottare
	hgeom->GetXaxis()->SetTitle("#font[42]{x[mm]}");
	hgeom->GetYaxis()->SetTitle("#font[42]{xp[mrad]}");

	//Add geom emitt value on the plot
	TPaveText *ptgeom=new TPaveText (0.25,0.81,0.9,0.91,"NDC");
	//NDC: normalized coordinates (x1,y1,x2,y2)
	char text[100];
	sprintf (text,"#font[42]{#splitline{Geom Emitt=(%.2f#pm%.2f)mm#upointmrad}{(effective percentage: %.2f%%)}}",emiP,emiErr,Peff);
	TText *tgeom=ptgeom->AddText(text);
	tgeom->SetTextColor(1);
	tgeom->SetTextSize(0.04);
	tgeom->SetTextAlign(12);//=H and V centred

	//Add warning in case of bad geom emitt computation
	TPaveText *ptwarn=new TPaveText (0.15,0.15,0.7,0.3,"NDC");
	//NDC: normalized coordinates (x1,y1,x2,y2)
	ptwarn->SetBorderSize(1);
	char text[100];
	sprintf (text,"#font[102]{#splitline{Geom emitt accuracy far}{from user request}}");
	TText *twarn=ptwarn->AddText(text);
	twarn->SetTextColor(2);
	twarn->SetTextSize(0.04);
	twarn->SetTextAlign(12);//=H and V centred

	hgeom->Draw("colz");
	ptgeom->Draw();
	if (overloop)
		ptwarn->Draw();


	// === SCATTER PLOT (not very useful at now) ==================================================
	c->cd(2);
	
	//temporary text
	TPaveText *ptmess=new TPaveText (0.1,0.1,0.8,0.4,"NDC");
	//NDC: normalized coordinates (x1,y1,x2,y2)
	ptmess->SetBorderSize(1);
	
	TText *t1=ptmess->AddText("#splitline{we could add a summary of computed}{and expected paramaters here,}");
	t1->SetTextColor(2);
	t1->SetTextSize(0.04);
	t1->SetTextFont(42);
	//t1->SetTextAlign(12);//=H and V centred
	TText *t2=ptmess->AddText("if no other plots are needed");
	t2->SetTextColor(2);
	t2->SetTextSize(0.04);
	t2->SetTextFont(42);
	//t2->SetTextAlign(12);//=H and V centred

	ptmess->Draw();


	/*
	// SCATTER PLOT 3D (non usa i colori, quindi non si capisce molto)
	TH3F *h3 = new TH3F("h3","#font[42]{Phase space scatter plot}", sl_Nsteps,Xmin, Xmax,ws_points, XPmin, XPmax, 100, Imin, Imax);
	for (int j=0;j<Nentries;j++)
		h3->Fill(pos[j],div[j],charge[j]);
	h3->GetYaxis()->SetRange(ws_points/3,ws_points*2/3); //range di bin da plottare
	h3->GetXaxis()->SetTitle("#font[42]{x[mm]}");
	h3->GetYaxis()->SetTitle("#font[42]{xp[mrad]}");
	h3->GetZaxis()->SetTitle("#font[42]{I[a.u.]}");
	h3->SetMarkerStyle(20);//pag 111 user's guide
    h3->SetMarkerSize(0.5);//pag 111 user's guide
	h3->Draw("");
	*/

	// === SAVE THE OUTPUT ===========================================================================
	c->cd();
	c->SaveAs("spfasi_plot_v4.gif");

	return 0;
}