#include "TFile.h"
#include "TH1.h"
#include "TMath.h"
#include "TMatrixD.h"
#include "TMinuit.h"
#include "TCanvas.h"
#include "TPaveText.h"
#include "TGraphErrors.h"
#include "TVectorD.h"
#include "TF1.h"
#include "TStyle.h"
#include "TLatex.h"
#include "TLine.h"
#include "TVirtualFFT.h"
#include "TVirtualFitter.h"
#include "THStack.h"
#include "TLegend.h"

#include <stdio.h>
#include <math.h>

#include <iostream>
#include <fstream>
#include <cmath>

using namespace std;

//****************************************************//
//********** Set user input parameters here **********//
//****************************************************//

const bool realData = true; // Flag for real data of MC
const bool useTruthInput = false;// Flag to use truth input (for MC)
const double B = 1.4513;//1.451269; // Main SR magnetic field [Tesla]
const double quadNvalue = 0.0;//realData ? 0.108 : 0.0; // Average field index, <n>
const double avg_n1mn = 0.0;//realData ? 0.0804 : 0.0; // <n(1-n)>  (N.B. != <n>(1-<n>) )
const int noRadBins = !useTruthInput ? 25 : 50; // The number of radial bins to solve for
const double beamWidth = 90.0; // Width of storage aperture [mm]
const int noInjTimeBins = !useTruthInput ? 25 : 149; // The number of bins the 149 ns injection pulse is divided up into
const double analysisStart = 4; // 0.15 Analysis start time [us]
const double analysisEnd = 100.0; // Analysis end time [us]
const double curvFactor = 0;//realData ? 5.e3 : 1.e5; // Curvature scale factor for chi^2 penalty

//****************************************************//

TH1D* RescaleAxis(TH1* input, Double_t Scale) {
  int bins = input->GetNbinsX();
  TAxis* xaxis = input->GetXaxis();
  double* ba = new double[bins+1];
  xaxis->GetLowEdge(ba);
  ba[bins] = ba[bins-1] + xaxis->GetBinWidth(bins);
  for (int i = 0; i < bins+1; i++) {
    ba[i] *= Scale;
  }
  TH1D* out = new TH1D(Form("%s_Rescaled",input->GetName()), input->GetTitle(), bins, ba);
  for (int i = 0; i <= bins; i++) {
    out->SetBinContent(i, input->GetBinContent(i));
    out->SetBinError(i, input->GetBinError(i));
  }
  return out;
}

//***********************************************************************//
//************************ Useful Functions ***************************//
//**********************************************************************//
const double c_light = 299.792458; // speed of light [mm/ns]

// Magic momentum [MeV/c]                                                   
double pMagic() {
  const double aMuon = 11659208.9e-10; //muon anomaly                        
  const double MASSMU = 105.6583745; // Mass of muon                         
  static double const pMagic_ = MASSMU/std::sqrt( aMuon );
  return pMagic_;
}

// Magic radius [mm]                                                         
double rMagic() {
  return pMagic()/( c_light * B )*1.0e3;
}

// Muon velocity [mev/c^2] (calculate beta from any momentum p)              
double betaP(double p1) {
  const double MASSMU = 105.6583745; //[MeV/c^2]                             
  double gamma_2_ = (p1/MASSMU)*(p1/MASSMU);
  double beta_ = std::sqrt(1.-1./gamma_2_);
  return beta_;
}

// Convert from momentum p [mev/c] to radis r [mm]                           
double p2r(double pp, double nVal) {
  return rMagic() * ( 1 + (pp-pMagic())/(pMagic() * (1 - nVal)));
} // from momentum p to radius r [mm] with quad value n                       
// Convert from radis r [mm] to momentum p [mev/c]                           
double r2p(double rr, double nVal) {
  return pMagic() * ( 1 + (rr-rMagic())*(1 - nVal)/rMagic());
} // from radius to momentum p [MeV/c] with quad value n 

// Convert from radius [mm] to time [µs]                                     
double r2t(double rt) {
  double tt_ = TMath::TwoPi()*rt/(betaP(r2p(rt,quadNvalue))*c_light)*1e-3;
  return tt_;
}  // [us]                                                                    
// Convert from momentum p [mev/c] to time [µs]                              
double p2t(double pt) {
  double tt_ = TMath::TwoPi()*p2r(pt,quadNvalue)/(betaP(pt)*c_light)*1e-3;
  return tt_;
}  // [us]                                                                    
// Convert from time [µs] to frequency [kH]z                                 
double t2f(double tt) {
  double ff_ = 1.0/tt*1e3;
  return ff_; // [kHz]                                                       
}

// Magic revolution time [µs]                                                
double timeMagic(){
  return p2t(pMagic());
}

//***********************************************************************//
//************** Gauss-Jordan Elimination Function **********************//
//************** Taken from Numerical Recipes for C++ ******************//
//**********************************************************************//

template <class T> void SWAP ( T& a, T& b )
{
  T c(a); a=b; b=c;
}

void gaussj(double** a, int n, double* b) {
  
  int i,icol=0,irow=0,j,k,l,ll;
  double big,dum,pivinv;
  int indxr[n], indxc[n], ipiv[n];
  
  for (j=0;j<n;j++) ipiv[j]=0;
  for (i=0;i<n;i++) {
    big=0.0;
    for (j=0;j<n;j++)
      if (ipiv[j] != 1)
	for (k=0;k<n;k++) {
	  if (ipiv[k] == 0) {
	    if (fabs(a[j][k]) >= big) {
	      big=fabs(a[j][k]);
	      irow=j;
	      icol=k;
	    }
	  }
	}
    ++(ipiv[icol]);
    if (irow != icol) {
      for (l=0;l<n;l++) SWAP(a[irow][l],a[icol][l]);
      SWAP(b[irow],b[icol]);
    }
    indxr[i]=irow;
    indxc[i]=icol;
    if (a[icol][icol] == 0.0) cout << "gaussj: Singular Matrix" << endl;
    pivinv=1.0/a[icol][icol];
    a[icol][icol]=1.0;
    for (l=0;l<n;l++) a[icol][l] *= pivinv;
    b[icol] *= pivinv;
    for (ll=0;ll<n;ll++)
      if (ll != icol) {
	dum=a[ll][icol];
	a[ll][icol]=0.0;
	for (l=0;l<n;l++) a[ll][l] -= a[icol][l]*dum;
	b[ll] -= b[icol]*dum;
      }
  }
  for (l=n-1;l>=0;l--) {
    if (indxr[l] != indxc[l])
      for (k=0;k<n;k++)
	SWAP(a[k][indxr[l]],a[k][indxc[l]]);
  }
  
  
}

//***********************************************************************//
//***************** chi^2 Minimisation Subroutine ******************//
//**********************************************************************//

double Minimisation(TH1D* hSignal, double injt0, double injBinWidth, int FRstartBin, double* radBinWidth, double* radius_i, double* Tc, double* f, double* I) {

  // Declare canvases that we'll put all our plots in for each iteration
  TCanvas* fitComparison = new TCanvas("fitComparison","fitComparison",800,600);
  TCanvas* fitComparisonStack = new TCanvas("fitComparisonStack","fitComparisonStack",800,600);
  TCanvas* timeResiduals = new TCanvas("timeResiduals","timeResiduals",800,600);
  TCanvas* timePulls = new TCanvas("timePulls","timePulls",800,600);
  TCanvas* pullProj = new TCanvas("pullProj","pullProj",800,600);
  TCanvas* residualFFT = new TCanvas("residualFFT","residualFFT",800,600);
  TCanvas* pullFFT = new TCanvas("pullFFT","pullFFT",800,600);
  TCanvas* pulse = new TCanvas("pulse","pulse",800,600);
  TCanvas* pulseStack = new TCanvas("pulseStack","pulseStack",800,600);
  TCanvas* RadialDistribution = new TCanvas("RadialDistribution","RadialDistribution",800,600);

  // Arrays of histograms for the injection pulse colour plots
  THStack* injBinStack = new THStack("injBinStack",""); //THStack containing each injection pulse bin's contribution to the final histogram
  THStack* injPulseStack = new THStack("injPulseStack","Injection Pulse"); //THStack containing color coding of injection pulse bins in above stack

  // Set start time (for plots)
  double FRstart = hSignal->GetBinLowEdge(FRstartBin) + hSignal->GetBinWidth(FRstartBin);

  // Determine and display the true analysis end time                                                                                       
  int analysisEndBin = -1;
  for (int i = 0; i < hSignal->GetNbinsX(); i++){
    if (hSignal->GetBinCenter(i) > analysisEnd){
      analysisEndBin = i-1;
      break;
    }
  }
  cout << "\n\tThe analysis end time is determined to be " << hSignal->GetBinCenter(analysisEndBin) + injt0 << " µs at bin number " << analysisEndBin << endl;

  // Set width of an injection pulse bin 
  double Delta = injBinWidth;

  // Define total number of time bins for plot (start after first pulse has passed)
  int plotStartBin = hSignal->FindBin(timeMagic()/2 - injt0) + 1;
  int noPlotTimeBins = analysisEndBin - plotStartBin - 1;
  int plotFitOffset = FRstartBin - plotStartBin;

  // Redefine the number of fit time bins to exclude bins before the fit start time, FRstart                                                    
  int noFitTimeBins = analysisEndBin - FRstartBin - 1;
  
  double N[noPlotTimeBins]; // The observed number of counts in time bin j
  double C[noPlotTimeBins]; // The expected number of counts in time bin j
  double Z[noPlotTimeBins]; // The weighting factor (i.e., expected number of counts in bin j

  double*** beta_ijk;  // The allocated geomtetry factors, ß(i,j,k) with indices i for radial bin, j for time bin, k for injection time bin
  beta_ijk = new double**[noRadBins];
  for (int i = 0; i < noRadBins; i++) {
    beta_ijk[i] = new double*[noPlotTimeBins];
    for (int j = 0; j < noPlotTimeBins; j++){
      beta_ijk[i][j] = new double[noInjTimeBins];
    }  
  }

  // Claim memory for radial and injection pulse matrices
  double** fMat = new double* [noRadBins];
  for(int i = 0 ; i < noRadBins; i++) fMat[i] = new double [noRadBins];
  double** IMat = new double* [noInjTimeBins];
  for(int i = 0 ; i < noInjTimeBins; i++) IMat[i] = new double [noInjTimeBins];

  // Allocate observed number of counts N in each bin from first selected bin
  for (int j=0; j<noPlotTimeBins; j++) {
    N[j] = hSignal->GetBinContent(plotStartBin+j);
    // Allocate weighting factor, Z[j], ready for chi^2 minimisation
    Z[j] = pow(hSignal->GetBinError(plotStartBin+j),2);
    if(Z[j] == 0) Z[j] = 1; // Ensure no zero variances (or we can't minimise)
  }
  
  // Begin determination of geometry factors ß(i,j,k)
  // Loop over radial bins
  for (int i = 0; i < noRadBins; i++) {

    // Loop over time bins
    for (int j = 0; j < noPlotTimeBins; j++) {

      // Loop over injection pulse time bins
      for (int k = 0; k < noInjTimeBins; k++) {

	// Find the true time of each bin center
	double time_j = hSignal->GetBinCenter(plotStartBin + j) + injt0 - (k - double(noInjTimeBins)/2.0)*(injBinWidth);
	
	// Find the degree of skewness, delta, of the full pulse at the given time
	double delta = radBinWidth[i]/radius_i[i]*time_j;

	// Find the distance in time, x, between the center of the injection time bin of the radial bin and the selected time bin
	double x = fabs(fmod(time_j + delta + Delta/2.0, Tc[i]) - (delta + Delta/2.0)); 	

        // Determine the 6 variants on the geomtery factors, ß(i,j)
	double beta_1 = 0.0;
	double beta_2 = 0.5 - x/delta + Delta/(2.0*delta) + 1.0/(2.0*delta*delta)*(x-Delta/2.0)*(x-Delta/2.0);
	double beta_3 = 0.5 - x/delta + Delta/(2.0*delta) - 1.0/(2.0*delta*delta)*(x-Delta/2.0)*(x-Delta/2.0);
	double beta_4 = Delta/delta - x*Delta/(delta*delta);
	double beta_5 = 1.0;
	double beta_6 = Delta/delta - (x/delta)*(x/delta) - (Delta/(2.0*delta))*(Delta/(2.0*delta));

	// Now, determine geomtery factors ß(i,j) from conditions as prescribed in Muon g-2 Note 271.
	if ( 0.0 < delta && delta < Delta/2.0 ) {
	  if ( (delta + Delta/2.0) < x ) {
	    beta_ijk[i][j][k] = beta_1;
	  }
	  else if ( Delta/2.0 < x && x < (delta + Delta/2.0) ) {
	    beta_ijk[i][j][k] = beta_2;		
	  }
	  else if ( (Delta/2.0 - delta) < x && x < Delta/2.0 ) {
	    beta_ijk[i][j][k] = beta_3;
	  }
	  else if ( 0 < x && x < (Delta/2.0 - delta) ) {	
	    beta_ijk[i][j][k] = beta_5;
	  }
	}
 	else if ( Delta/2.0 < delta && delta < Delta ) {
	  if ( (delta + Delta/2.0) < x ) {
	    beta_ijk[i][j][k] = beta_1;
	  }
	  else if ( Delta/2.0 < x && x < (delta + Delta/2.0) ) {
	    beta_ijk[i][j][k] = beta_2;	
	  }
	  else if ( (delta - Delta/2.0) < x && x < Delta/2.0 ) {
	    beta_ijk[i][j][k] = beta_3;
	  }
	  else if ( 0 < x && x < (delta - Delta/2.0) ) {
	    beta_ijk[i][j][k] = beta_6;
	  }
	}
	else if ( Delta < delta ) {
	  if ( (delta + Delta/2.0) < x ) {
	    beta_ijk[i][j][k] = beta_1;
	  }
	  else if ( (delta - Delta/2.0) < x && x < (delta + Delta/2.0) ) {
	    beta_ijk[i][j][k] = beta_2;
	  }
	  else if ( Delta/2.0 < x && x < (delta - Delta/2.0) ) {
	    beta_ijk[i][j][k] = beta_4;
	  }
	  else if ( 0 < x && x < Delta/2.0) {
	    beta_ijk[i][j][k] = beta_6;
	  }	
	}
	
      } // End k-loop over number of injection pulse time bins
      
    } // End j-loop over number of spectrum time bins
    
  } // End n-loop over number of radial bins
  
  // Initialise chi^2 variables
  double chisq_old = 0; // The chi^2 from the previous iteration
  double chisq = 0.0; // The total chi^2
  int nDoF = 0; // The chi^2 from the previous iteratin number
  double chisq_dof = 0.0; // The chi^2 per degree of freedom
  
  // Define maximum number of iterations
  int itMax = 100;

  // Return value for chisq/dof (this returns -1 unless we were successful in minimisation)
  double chisq_dof_return = -1;
  
  // Loop over chi^2 minimisation iterations
  for (int itNum = 0; itNum < itMax; itNum++ ) {

    // Display iteration number of fit     
    cout <<  "\n --> Iteration " << itNum << "..." << endl;
    
    //////////////////////////////////////////////////////////////////////////
    // Radial bin minimisation
    //////////////////////////////////////////////////////////////////////////

    cout <<  "\n --> Radial Bin Minimisation..." << endl;

    // Extract radial bins (either calculating or using truth)
    if(!useTruthInput){

      // Set ß(i,j) values - for solution values of I[k]
      TMatrixD betaIJ(noRadBins,noFitTimeBins); // The allocated geomtetry factors, ß(i,j) where i is for injection bin
      for (int i = 0; i < noRadBins; i++) {
	for (int j = 0; j < noFitTimeBins; j++) {
	  betaIJ(i,j) = 0.0;
	  for (int k = 0; k < noInjTimeBins; k++) {
	    betaIJ(i,j) += beta_ijk[i][j+plotFitOffset][k]*I[k];
	  }
	}
      }

      // Initialise all variables
      for (int n = 0; n < noRadBins; n++) {
	f[n] = 0.0;
	for (int i = 0; i < noRadBins; i++) {
	  fMat[n][i] = 0.0;
	}
      }

      // Prepare system of equations - populate half of matrix and then use fact it's symmetric
      for (int j = 0; j < noFitTimeBins; j++) {
	for (int n = 0; n < noRadBins; n++) {
	  f[n] += betaIJ(n,j)*N[j+plotFitOffset]/Z[j+plotFitOffset];
	  for (int i = 0; i <= n; i++) {
	    fMat[n][i] += betaIJ(i,j)*betaIJ(n,j)/Z[j+plotFitOffset];
	  }
	}
      }
      for (int n = 0; n < noRadBins; n++) {
	for (int i = n; i < noRadBins; i++) {
	  fMat[n][i] = fMat[i][n];
	}
      }

      // Call gaussj function --> the gauss-jordan elimination to solve for the chi^2 minimum
      gaussj( fMat, noRadBins, f);
      for (int n = 0; n < noRadBins; n++) {
	f[n] = f[n];
      }

    } else {

      for (int n = 0; n < noRadBins; n++) {
	if(fabs(radius_i[n] - rMagic()) > 0.003*rMagic()) f[n] = 0;
	else {
	  f[n] = 1.69e6/noRadBins*TMath::Gaus(radius_i[n],  rMagic(), 0.001*rMagic()); // Parameters from g2gen
	}
      }
    }

    // Calculate expected number of counts, C_j
    for (int j = 0; j < noPlotTimeBins; j++) {
      C[j] = 0.0;
      for (int i = 0; i < noRadBins; i++) {
	for(int k = 0; k < noInjTimeBins; k++){
	  C[j] += beta_ijk[i][j][k]*f[i]*I[k];
	}
      }
    }

    // Calculate norm factor that will minimise chi2 analytically - this is just from differentiating chi2 formula and setting to 0
    if(useTruthInput){
      double expectedSum = 0;
      double weightedSum = 0;
      for (int j = plotFitOffset; j < plotFitOffset + noFitTimeBins; j++) {
	if(N[j] > 0){
	  expectedSum += C[j];
	  weightedSum += (C[j]*C[j])/N[j];
	}
      }

      double normFactor = expectedSum/weightedSum;
      cout <<  "\n --> Radial bins normalisation factor = " << expectedSum/weightedSum << endl;
      for (int n = 0; n < noRadBins; n++) f[n] *= normFactor;
      
      // Re-calculate expected number of counts, C_j with the normalisation factor
      for (int j = 0; j < noPlotTimeBins; j++) {
	C[j] = 0.0;
	for (int i = 0; i < noRadBins; i++) {
	  for(int k = 0; k < noInjTimeBins; k++){
	    C[j] += beta_ijk[i][j][k]*f[i]*I[k];
	  }
	}
      }
    } // useTruthInput

    // Calculate and display chi^2_min/d.o.f.
    chisq = 0.0;
    for (int j = plotFitOffset; j < plotFitOffset + noFitTimeBins; j++) {
      chisq += (C[j]-N[j])*(C[j]-N[j])/Z[j]; //chi^2 function
    }
    nDoF = (noFitTimeBins - noRadBins); // Determine the number of degrees of freedom
    chisq_dof = chisq/nDoF; // Determine the chi^_minimum per degree of freedom
    cout << "\tTotal chi^2 = " << chisq << endl; 
    cout << "\tchi^2_min/d.o.f. = " << chisq_dof << endl;
    cout << "\tp-value = " << TMath::Prob(chisq,nDoF) << endl;

    //////////////////////////////////////////////////////////////////////////
    // Injection pulse minimisation
    //////////////////////////////////////////////////////////////////////////

    cout <<  "\n --> Injection Pulse Minimisation..." << endl;

    // Extract injection pulse bins (either calculating or using truth)
    if(!useTruthInput){

      // Set ß(k,j) values - for solution values of f[i]
      TMatrixD betaKJ(noInjTimeBins,noFitTimeBins); // The allocated geomtetry factors, ß(k,j) where k is for injection bin
      for (int k = 0; k < noInjTimeBins; k++) {
	for (int j = 0; j < noFitTimeBins; j++) {
	  betaKJ(k,j) = 0.0;
	  for (int i = 0; i < noRadBins; i++) {
	    betaKJ(k,j) += beta_ijk[i][j+plotFitOffset][k]*f[i];
	  }
	}
      }

      // Initialise all variables
      for (int n = 0; n < noInjTimeBins; n++) {
	I[n] = 0.0;
	for (int k = 0; k < noInjTimeBins; k++) {
	  IMat[n][k] = 0.0;
	}
      }

      // Prepare system of equations - populate half of matrix and then use fact it's symmetric
      for (int j = 0; j < noFitTimeBins; j++) {
	for (int n = 0; n < noInjTimeBins; n++) {
	  I[n] += betaKJ(n,j)*N[j+plotFitOffset]/Z[j+plotFitOffset];
	  for (int k = 0; k <= n; k++) {
	    IMat[n][k] += betaKJ(k,j)*betaKJ(n,j)/Z[j+plotFitOffset];
	  }
	}
      }
      for (int n = 0; n < noInjTimeBins; n++) {
	for (int k = n; k < noInjTimeBins; k++) {
	  IMat[n][k] = IMat[k][n];
	}
      }

      // Add in curvature term with scale factor (need to scale with entries as it's added into the chi2 as such)
      double dataSum = 0;
      for(int j = 0; j < noFitTimeBins; j++) dataSum += N[j+plotFitOffset];
      double curvScale = curvFactor / dataSum;
      for (int k = 0; k < noInjTimeBins; k++) {
	if(k >= 2) IMat[k][k-2] += curvScale*1;
	else       IMat[k][noInjTimeBins-k-2] += curvScale*1;

	if(k >= 1) IMat[k][k-1] -= curvScale*4;
	else       IMat[k][noInjTimeBins-k-1] -= curvScale*4;

	IMat[k][k]   += curvScale*6;

	if(k < noInjTimeBins-1) IMat[k][k+1] -= curvScale*4;
	else                         IMat[k][noInjTimeBins-1-k] -= curvScale*4;

	if(k < noInjTimeBins-2) IMat[k][k+2] += curvScale*1;
	else                         IMat[k][noInjTimeBins-2-k] += curvScale*1;
      }

      // Call gaussj function --> the gauss-jordan elimination to solve for the chi^2 minimum
      gaussj(IMat, noInjTimeBins, I);
      double injTotal = 0;

    } else {
      for (int n = 0; n < noInjTimeBins; n++) {
	I[n] = TMath::Gaus((n - double(noInjTimeBins)/2) * injBinWidth, 0, 15e-3);
      }
    }

    // Calculate expected number of counts, C_j
    for (int j = 0; j < noPlotTimeBins; j++) {
      C[j] = 0.0;
      for (int i = 0; i < noRadBins; i++) {
	for(int k = 0; k < noInjTimeBins; k++){
	  C[j] += beta_ijk[i][j][k]*f[i]*I[k];
	}
      }
    }

    // Calculate and display chi^2_min/d.o.f.
    chisq = 0.0;
    for (int j = plotFitOffset; j < plotFitOffset + noFitTimeBins; j++) {
      chisq += (C[j]-N[j])*(C[j]-N[j])/Z[j]; //chi^2 function
    }
    nDoF = (noFitTimeBins - noRadBins); // Determine the number of degrees of freedom
    chisq_dof = chisq/nDoF; // Determine the chi^_minimum per degree of freedom
    cout << "\tTotal chi^2 = " << chisq << endl; 
    cout << "\tchi^2_min/d.o.f. = " << chisq_dof << endl;
    cout << "\tp-value = " << TMath::Prob(chisq,nDoF) << endl;

    // Display chi^2_min/d.o.f from previous iteration
    if (itNum > 0) {
      cout << "\tchi^2 for previous iteration = " << chisq_old << endl;
    }

    //////////////////////////////////////////////////////////////////////////
    // Plot result and overlay with data
    //////////////////////////////////////////////////////////////////////////

    // Populate histogram with calculated time spectrum and draw
    TH1D *calcSpectrum = new TH1D(Form("calcSpectrum_%d",itNum),"Calculated bin content", analysisEndBin,  hSignal->GetXaxis()->GetXmin(), hSignal->GetXaxis()->GetXmin() + analysisEndBin*hSignal->GetBinWidth(1));
    calcSpectrum->GetXaxis()->SetTitle("Time [#mus]");
    for(int j = 0; j < noPlotTimeBins; j++){
      calcSpectrum->SetBinContent(j+plotStartBin,C[j]);
    }

    fitComparison->cd();
    hSignal->SetStats(0);
    hSignal->GetXaxis()->SetRangeUser(analysisStart,analysisStart+0.7);
    calcSpectrum->SetLineColor(51+2*itNum);
    if(itNum == 0) hSignal->Draw("HIST");
    calcSpectrum->Draw("SAME");
    fitComparison->Update();

    //////////////////////////////////////////////////////////////////////////
    // Make stacked histograms of each injection bin's contribution
    //////////////////////////////////////////////////////////////////////////

    // Delete everything in the inj bin stack and then start again - easiest way to avoid memory leak!
    if(injBinStack->GetNhists() > 0){
      for(int i = injBinStack->GetNhists() - 1; i >= 0; i--){
     	TObject* thisHist = injBinStack->GetHists()->At(i);
	injBinStack->RecursiveRemove(thisHist);
	delete thisHist;
      }
    }

    // If first time, add new histograms to stack
    for(int k = 0; k < noInjTimeBins; k++){
      TH1F* injBinContrib = new TH1F(Form("injBinContrib_%03d",k), Form("Contribution from injection time bin %d", k), analysisEndBin,  hSignal->GetXaxis()->GetXmin(), hSignal->GetXaxis()->GetXmin() + analysisEndBin*hSignal->GetBinWidth(1));
      injBinContrib->SetFillColor(51 + int(k*50./noInjTimeBins));
      injBinContrib->SetLineColor(51 + int(k*50./noInjTimeBins));
      injBinStack->Add(injBinContrib);
    }

    // Calculate expected number of counts in each bin from this injection pulse, C_k and draw
    for (int j = 0; j < noPlotTimeBins; j++) {
      for(int k = 0; k < noInjTimeBins; k++){
	double C_K = 0;
	for (int i = 0; i < noRadBins; i++) {
	  C_K += beta_ijk[i][j][k]*f[i]*I[k];
	}
	((TH1F*)injBinStack->GetHists()->At(k))->SetBinContent(j+plotStartBin, C_K);
      }
    }
    fitComparisonStack->cd();
    hSignal->Draw("HIST");
    injBinStack->Draw("HIST SAME");
    hSignal->Draw("SAME");
    fitComparisonStack->Update();

    // Make stacked plot with different colors for each injection pulse bin (kind of legend for above stacked plot)
    if(injPulseStack->GetNhists() > 0){
      for(int i = injPulseStack->GetNhists() - 1; i >= 0; i--){
	TObject* thisHist = injPulseStack->GetHists()->At(i);
	injPulseStack->RecursiveRemove(thisHist);
	delete thisHist;
      }
    }
    // Add histograms (one for each injection bin)
    for(int k = 0; k < noInjTimeBins; k++){
      TH1F* injPulseContrib = new TH1F(Form("injPulseContrib_%03d",k), Form("Colour-coding of injection time bin %d", k), noInjTimeBins, -0.5*(noInjTimeBins+1)*injBinWidth-injt0, 0.5*(noInjTimeBins-1)*injBinWidth-injt0);
      injPulseContrib->SetFillColor(51 + int(k*50./noInjTimeBins));
      injPulseContrib->SetLineColor(51 + int(k*50./noInjTimeBins));
      injPulseContrib->SetBinContent(k+1,I[k]/noInjTimeBins);
      injPulseStack->Add(injPulseContrib);
    }
    pulseStack->cd();
    injPulseStack->Draw("HIST");
    pulseStack->Update();

    //////////////////////////////////////////////////////////////////////////
    // Residuals, Pulls, FFTs
    //////////////////////////////////////////////////////////////////////////

    // Calculate residuals, populate histogram with residuals and draw
    TH1D *Residuals = new TH1D(Form("Residuals_%d",itNum),"Fit Residuals;Time [#mus];Measured - Predicted", noFitTimeBins,  FRstart, FRstart + noFitTimeBins*hSignal->GetBinWidth(1));
    double resSum = 0.0;
    for(int i=0; i<noFitTimeBins;i++){
      Residuals->SetBinContent(i,(N[i+plotFitOffset]-C[i+plotFitOffset]));
      resSum += N[i+plotFitOffset]-C[i+plotFitOffset];
    }
    double resMean = resSum/noFitTimeBins;
    //    cout << "\n(Sum of residuals = " << resSum << "; residual mean = " << resMean << ")\n" << endl;
  
    timeResiduals->cd();
    Residuals->SetLineColor(51+2*itNum);
    Residuals->SetStats(0);
    if(itNum == 0) Residuals->Draw();
    else                  Residuals->Draw("SAME");
    timeResiduals->Update();
  
    // Do the same for pull
    TH1D *Pulls = new TH1D(Form("Pulls_%d",itNum),"Fit Pulls;Time [#mus];(Measured - Predicted) / Uncertainty", noFitTimeBins,  FRstart, FRstart + noFitTimeBins*hSignal->GetBinWidth(1));
    TH1D *PullProj = new TH1D(Form("PullProj_%d",itNum),"Fit Pulls;No. Bins;(Measured - Predicted) / Uncertainty", 100, -10, 10);
    for(int i=0; i<noFitTimeBins;i++){
      if(C[i+plotFitOffset] > 0) {
	Pulls->SetBinContent(i,(N[i+plotFitOffset]-C[i+plotFitOffset])/sqrt(Z[i+plotFitOffset]));
	PullProj->Fill((N[i+plotFitOffset]-C[i+plotFitOffset])/sqrt(Z[i+plotFitOffset]));
      }
    }
  
    timePulls->cd();
    Pulls->SetLineColor(51+2*itNum);
    Pulls->SetStats(0);
    if(itNum == 0) Pulls->Draw();
    else                  Pulls->Draw("SAME");
    timePulls->Update();
  
    pullProj->cd();
    PullProj->SetLineColor(51+2*itNum);
    if(itNum == 0) PullProj->Draw();
    else                  PullProj->Draw("SAME");
    pullProj->Update();

    // Do FFT of the residuals
    residualFFT->cd();
    TH1D *hFFTInput = new TH1D(Form("hFFTInput_%d",itNum),";Time [#mus];Counts", Residuals->GetNbinsX(), Residuals->GetXaxis()->GetXmin(), Residuals->GetXaxis()->GetXmax());
    for(int bin = 1; bin <= hFFTInput->GetNbinsX(); bin++){
      hFFTInput->SetBinContent(bin, Residuals->GetBinContent(bin));
    }
    // Do FFT
    TH1 *hm1 = 0;
    TVirtualFFT::SetTransform(0);
    hm1 = hFFTInput->FFT(hm1,"MAG");
    hm1->SetName(Form("Residual_FFTMag_%d",itNum));
    TH1D* fft1 = RescaleAxis(hm1,1./(hFFTInput->GetXaxis()->GetXmax() - hFFTInput->GetXaxis()->GetXmin()));
    fft1->GetXaxis()->SetRangeUser(0.1,35);
    fft1->SetTitle("FFT of Residuals;Frequency (MHz);Magnitude [Arb Units]");
    fft1->SetStats(0);
    fft1->SetLineColor(51+2*itNum);
    if(itNum == 0) fft1->Draw("HIST"); 
    else                  fft1->Draw("HISTSAME");
    residualFFT->Update();
  
    // Do FFT of the pulls
    pullFFT->cd();
    TH1D *hFFTInput2 = new TH1D(Form("hFFTInput2_%d",itNum),";Time [#mus];Counts", Pulls->GetNbinsX(), Pulls->GetXaxis()->GetXmin(), Pulls->GetXaxis()->GetXmax());
    for(int bin = 1; bin <= hFFTInput2->GetNbinsX(); bin++){
      hFFTInput2->SetBinContent(bin, Pulls->GetBinContent(bin));
    }
    // Do FFT
    TH1 *hm2 = 0;
    TVirtualFFT::SetTransform(0);
    hm2 = hFFTInput2->FFT(hm2,"MAG");
    hm2->SetName(Form("Pull_FFTMag_%d",itNum));
    TH1D* fft2 = RescaleAxis(hm2,1./(hFFTInput2->GetXaxis()->GetXmax() - hFFTInput2->GetXaxis()->GetXmin()));
    fft2->GetXaxis()->SetRangeUser(0.1,35);
    fft2->SetTitle("FFT of Pulls;Frequency (MHz);Magnitude [Arb Units]");
    fft2->SetStats(0);
    fft2->SetLineColor(51+2*itNum);
    if(itNum == 0) fft2->Draw("HIST"); 
    else                  fft2->Draw("HISTSAME");
    pullFFT->Update();

    //////////////////////////////////////////////////////////////////////////
    // Injection pulse & radial distribution results
    //////////////////////////////////////////////////////////////////////////
    
    // Populate histogram with injection pulse and draw
    TH1D *injPulse = new TH1D(Form("injPulse_%d",itNum),"Injection Pulse", noInjTimeBins, -0.5*(noInjTimeBins+1)*injBinWidth-injt0, 0.5*(noInjTimeBins-1)*injBinWidth-injt0);
    injPulse->GetXaxis()->SetTitle("Time [ns]");
    for(int i = 1; i <= noInjTimeBins; i++){
      injPulse->SetBinContent(i,I[i-1]);
    }
    pulse->cd();
    injPulse->SetLineColor(51+2*itNum);
    if(itNum == 0) {
      injPulse->Draw();
    }
    else injPulse->Draw("SAMES");
    pulse->Update();
    cout << "\n--> Injection T0 = " << 1000*injPulse->GetMean() << " ns, Injection RMS = " << 1000*injPulse->GetRMS() << "ns" << endl;

    // Fill histogram with determined radial distribution (add in variable bins as two at edges are larger)
    double radBinLowEdges[noRadBins+1];
    radBinLowEdges[0] = rMagic()-45;
    for(int i = 1; i <= noRadBins; i++) radBinLowEdges[i] = radBinLowEdges[i-1] + radBinWidth[i-1];
    TH1D *hRadDistribution = new TH1D(Form("hRadDistribution_%d",itNum),"Radial Distribution;Radius [mm];Arbitrary Units [au]", noRadBins,  radBinLowEdges);
    for(int i = 0; i < noRadBins; i++) hRadDistribution->SetBinContent(i+1,f[i]);

    // Now, determine E-field correction
    double Beta = 0.9994174214209439;
    double xesq = 0.0;
    double C_E = 0.0;
    double count = 0.0;
    for (int i = 0; i < noRadBins; i++){
      count += f[i];
      xesq += (radius_i[i]-rMagic())*(radius_i[i]-rMagic())*f[i];
    }
    xesq = xesq/count;
    // E-field correction, C_E                                        
    C_E = -2.0*avg_n1mn*pow(Beta,2)*xesq/pow(rMagic(),2);
    //Display calculated mean, rms and C_E                                  
    cout << "--> Mean = " << hRadDistribution->GetMean() << " mm, RMS = " << hRadDistribution->GetRMS() << "mm, C_E = " << C_E*1e9  << "ppb\n" << endl; 
  
    // Draw radial distribution with information from fit
    RadialDistribution->cd();
    hRadDistribution->SetLineColor(51+2*itNum);
    if(itNum == 0){
      hRadDistribution->Draw("HIST");
    } else {
      hRadDistribution->Draw("HISTSAMES");
    }
    RadialDistribution->Update();

    
    //////////////////////////////////////////////////////////////////////////
    // Check for chi2 convergence
    //////////////////////////////////////////////////////////////////////////

    // Always stop if useTruthInput is set
    if ( useTruthInput ){
      cout << "\n --> Using Truth Input for radial and injection information.  Stopping iterations..." << endl;
      chisq_dof_return = chisq_dof;
      break;
    }

    // Determine the difference between the current and previous chi^2_min/d.o.f. and look for convergence
    if (  (chisq <= chisq_old ) && ( abs(chisq-chisq_old) / (chisq+chisq_old) < 1e-4) ) {		
      cout << "\n --> Converged with rel. change in chi^2 of " << abs(chisq-chisq_old)/(chisq+chisq_old) << endl;
      cout << "\nFinishing the chi**2-minimisation..." << endl;

      // We've converged, so stop these iterations
      chisq_dof_return = chisq_dof;
      break;
    }

    if ( itNum > 0 and chisq > chisq_old ) {
      cout << "\n!!! chi^2-minimisation is diverging! Stopping..." << endl;
      break;
    }

    // Replace old chi^2 value with chi^2 value from this iteration.
    chisq_old = chisq;
    
    if ( itNum == itMax - 1 ) { //If the minimisation does not converge...
      cout << "\n\t !!! chi^2-minimisation does not converge! Stopping..." << endl;
    }

  }

  // // Plot pretty pictures for animated gif
  // TCanvas* cGifPics = new TCanvas();
  // injPulseStack->SetTitle("Injection Pulse;Time [#mus];Intensity [a.u.]");
  // injPulseStack->GetYaxis()->SetTitleOffset(1.2);
  // injPulseStack->Draw("HIST");
  // cGifPics->SaveAs("Images/InjectionPulse.png");
  
  // injBinStack->GetYaxis()->SetTitleOffset(1.2);
  // injBinStack->SetMinimum(0);
  // injBinStack->Draw("HIST");
  // hSignal->GetXaxis()->UnZoom();
  // hSignal->Draw("SAME");
  // for(int i = 1; i < 671; i ++){
  //   double xMin = i*timeMagic() - injt0 - 0.6*timeMagic();
  //   double xMax = i*timeMagic() - injt0 + 0.6*timeMagic();
  //   injBinStack->GetXaxis()->SetRangeUser(xMin, xMax);
  //   if(xMin > 4){
  //     hSignal->GetXaxis()->SetRangeUser(xMin-0.1, xMax+0.1);
  //     if(hSignal->GetMaximum() > 0) injBinStack->SetMaximum(hSignal->GetMaximum());
  //   }
  //   injBinStack->SetTitle(Form("Fast Rotation Signal %03.2f - %03.2f #mus;Time [us];Entries",xMin, xMax));
  //   cGifPics->SaveAs(Form("Images/FitComparison_%05d.png",i));
  // }

  // De-allocate memory from beta_ijk
  for (int i = 0; i < noRadBins; i++) {
    for (int j = 0; j < noPlotTimeBins; j++){
      delete [] beta_ijk[i][j];
    }
    delete [] beta_ijk[i];
  }
  delete [] beta_ijk;

  return chisq_dof_return;

} // End of minimisation subroutine


//***********************************************************************//
//****************** Main Fast Rotation  Subroutine ****************//
//**********************************************************************//

void FRRadDist_Merge1_Updated(std::string inputFile,std::string inputHist){
 
  // Check configuration options
  if(realData && useTruthInput){
    cout << "ERROR: Conflicting flags of realData and useTruthInput" << endl;
    return;
  }

  // Open input data file and read in signal
  TFile* file = TFile::Open(inputFile.c_str());
  TH1D* hSignal = (TH1D*)file->Get(inputHist.c_str());

  // Welcome message
  cout <<  "\n* FRRadDist v2.0 (A.Keshavarzi & J.Mott), The muon g-2 experiment, Fermilab, 2018 *" << endl;
  cout << "* CERN III style chi^2 minimisation routine to solve for average radial distribution of stored beam *" << endl;
  cout << "* [Notes on the CERN III fast rotation analysis, H. Jostlein and P. Hattersley, 1967] *" << endl;
  cout <<  "---------------------------------------------------------------------------------------" << endl;

  // Output all input parameters and force user to confirm
  cout << "\n**************************************" << endl;
  cout << "********** Input parameters **********" << endl;
  cout << "**************************************" << endl;
  if (realData){
    cout << "• Use real data = Yes" << endl;
    cout << "  [Data file: " << inputFile << "]"<< endl;
  }
  else{
    cout << "• Use real data =  No [Simulation file: " << inputFile << "]"<< endl;
    if (useTruthInput){
      cout << "   -->Use truth input = Yes" << endl;
    }
    else{
      cout << "   -->Use truth input = No" << endl;
    }
  }
  cout << "• Magnetic field, B = " << B << " T" << endl;
  cout << "   --> Magic momentum, p_magic = " << pMagic() << " MeV/c" << endl;
  cout << "   --> Magic radius, r_magic = " << rMagic() << " mm" << endl;
  cout << "• Average field index, <n> = " << quadNvalue << endl;
  cout << "• Average n(1-n) over ring azimuth, <n(1-n)> = " << avg_n1mn << endl;
  cout << "• Number of radial bins to solve for, i = " << noRadBins << endl;
  cout << "   --> Width of beam aperture = " << beamWidth << " mm" << endl;
  cout << "• Number of injection pulse time bins to solve for, k = " << noInjTimeBins << endl;
  cout << "• Analysis start time = " << analysisStart << " µs" << endl;  
  cout << "• Analysis end time  = " << analysisEnd << " µs" << endl;
  cout << "• Curvature scale (chi^2 penalty) factor for injection pulse minimisation = " << curvFactor << endl; 
  cout << "**************************************" << endl;
  int num = 0;
  while (num >= 0){
    num += 1;
    cout << "\nAre all the input parameters listed above correct? (Please type Y or N):" << endl;
    string inputConfirm;
    getline (cin, inputConfirm);
    if (inputConfirm == "Y" or inputConfirm == "y"){
      cout << "\n--> All input parameters confirmed to be correct. Proceeding..." << endl;
      break;
    }
    else if (inputConfirm == "N" or inputConfirm == "n"){
      cout << "\n-->Input parameters are incorrect!! Please adjust necessary input parameters to correct values and re-run program." << endl;
      cout << "Exiting..." << endl;
      return;
    }
    else{
      cout << "Cannot recognise user input! Only type either 'Y' (yes) or 'N.' (no)" << endl;
    }
  }
  
  // Declare arrays for radial information
  double radBinWidth[noRadBins]; // The width of the defined number of radial bins
  double radius_i[noRadBins]; // The radius value of each radial bin 
  double Tc[noRadBins]; // The central revoluation time of each radial bin 

  // Find the radial bin width for all the defined radial bins (first and last bin are 1.5 times wider than the rest)                                  
  for (int i= 0; i < noRadBins; i++) {
    if ( i == 0 or i == (noRadBins-1) ) {
      radBinWidth[i] = beamWidth/(noRadBins+1) * 1.5;  // We have noRadBin-2 at binwidth and 2 and binWidth*1.5 (so we have noRadBins+1 * binWidth)       
    } else {
      radBinWidth[i] = beamWidth/(noRadBins+1);
    }
  }

  // Find the radius and central rotation time values for all the defined radial bins                                                                     
  for (int i= 0; i < noRadBins; i++) {
    if ( i == 0 ) {
      radius_i[i] = rMagic() - beamWidth/2.0 + 0.5*radBinWidth[i]; // [mm]
    } else if ( i == 1 or i == noRadBins-1 ) {
      radius_i[i] =  radius_i[i-1] + 0.5*radBinWidth[i-1] + 0.5*radBinWidth[i]; // [mm]
    } else {
      radius_i[i] =  radius_i[i-1] + radBinWidth[i]; // [mm]
    }
    Tc[i] = r2t(radius_i[i]); // [µs]                                         
  }

  // Determine approximate start and end time of the chosen injection pulse.
  hSignal->GetXaxis()->SetRangeUser(analysisStart,analysisStart+0.2); 
  int injPulseStartBin = hSignal->GetMinimumBin();
  double injPulseStart = hSignal->GetBinCenter(injPulseStartBin);
  hSignal->GetXaxis()->SetRangeUser(injPulseStart+0.5*timeMagic(),injPulseStart+1.5*timeMagic()); 
  int injPulseEndBin = hSignal->GetMinimumBin();

  // Find start of injection pulse and define as the fast rotation fit start time
  int FRstartBin = injPulseStartBin;
  double FRstart = hSignal->GetBinLowEdge(FRstartBin) + hSignal->GetBinWidth(FRstartBin);
  cout << "\n\tFast rotation fit start time = " << FRstart << " µs, at bin number " << FRstartBin << endl;
  
  // Have a good guess at the t0 offset (injection pulse shape should take care of the rest)
  hSignal->GetXaxis()->SetRange(injPulseStartBin,injPulseEndBin); 
  double injPulseMeanPoint = hSignal->GetMean();
  double injt0 = useTruthInput ? 0 : timeMagic()*round(injPulseMeanPoint/timeMagic()) - injPulseMeanPoint;
  hSignal->GetXaxis()->UnZoom(); // reset

  // Overwrite number of injection time bins with our new number of bins
  double injBinWidth = timeMagic()/noInjTimeBins;
  cout << "\n\tThe injection pulse width is set to be " << noInjTimeBins*injBinWidth << " µs"  << endl;
  cout << "\tThe width of the " << noInjTimeBins << " bins of the injection pulse are " << injBinWidth << " µs" << endl;

  // Declare radial distribution & injection pulse arrays
  double* f = new double [noRadBins];
  double* I = new double [noInjTimeBins];

  // Canvas for comparing injection pulses
  TCanvas *pulseGuess = new TCanvas("pulseGuess","pulseGuess",800,600);

  // Put some checks on our injection pulse binning & check that input histogram bins are equal to 1 us
  if(hSignal->GetBinWidth(1) != 0.001){
    cout << "ERROR: Input histogram \"" << inputHist << "\" has bin width of " << 1e3*hSignal->GetBinWidth(1) << "\t ns, not 1 ns as expected. Injection pulse rebinning will fail. Stopping..." << endl;
    return;
  }
  if(noInjTimeBins > 1e3*timeMagic()){
    cout << "ERROR: We're trying to rebin our 1 ns input bins to " << 1e3*injBinWidth << " ns bins.  Rebinning will fail. Stopping..." << endl;
    return;
  }

  // Determine the content of each injection pulse time bin
  for(int i = 0 ; i < noInjTimeBins; i++){

    if(useTruthInput) I[i] = TMath::Gaus((i - double(noInjTimeBins)/2) * injBinWidth, 0, 40e-3);  // Parameters from g2gen

    else {

      // Set low and high edges that we're aiming for in new 
      double binLowEdge = injPulseMeanPoint + (i - double(noInjTimeBins/2))*injBinWidth;
      double binHighEdge = binLowEdge+injBinWidth;

      // Loop over input histogram and fill injection pulse bins
      I[i] = 0;
      for(int j = 1; j <= hSignal->GetNbinsX(); j++){

	// Set input bin low/high edges
	double inLowEdge = hSignal->GetBinLowEdge(j);
	double inHighEdge = hSignal->GetBinLowEdge(j) + hSignal->GetBinWidth(j);

	// Case 1 : First bin of interest (binLowEdge falls in this bin)
	if(inLowEdge < binLowEdge and inHighEdge > binLowEdge){
	  double addFrac = (inHighEdge - binLowEdge)/hSignal->GetBinWidth(j);
	  I[i] += addFrac*hSignal->GetBinContent(j);
	}
	
	// Case 2 : Bin fully in our target range
	if(inLowEdge > binLowEdge and inHighEdge < binHighEdge){
	  I[i] += hSignal->GetBinContent(j);
	}

	// Case 3: Last bin of interest (binHighEdge falls in this bin)
	if(inLowEdge < binHighEdge and inHighEdge > binHighEdge){
	  double addFrac = (binHighEdge - inLowEdge)/hSignal->GetBinWidth(j);
	  I[i] += addFrac*hSignal->GetBinContent(j);
	}
      }
    }
  }
  
  // Populate histogram with injection pulse and draw
  TH1D *injPulseGuess = new TH1D("injPulseGuess","Injection Pulse Guess", noInjTimeBins, -0.5*(noInjTimeBins+1)*injBinWidth-injt0, 0.5*(noInjTimeBins-1)*injBinWidth-injt0);
  injPulseGuess->GetXaxis()->SetTitle("Time [ns]");
  for(int i = 0; i < noInjTimeBins; i++){
    injPulseGuess->SetBinContent(i+1,I[i]);
  }
  pulseGuess->cd();
  injPulseGuess->SetLineColor(51);
  injPulseGuess->Draw();
  pulseGuess->Update();

  // Display messages of chosen scan results 
  cout << "\n--> Starting minimisation with fixed t0 of " << 1e3*injt0 << " ns ..." << endl;
  
  //*****************************************//
  //******** Main minimisation call *********//    
  //*****************************************//

  // Call subroutine to minimisation chi^2 function for final time - return value is chisq_dof
  double chisq_dof = Minimisation(hSignal, injt0, injBinWidth, FRstartBin, radBinWidth, radius_i, Tc, f, I);

  if(chisq_dof < 0) {
    cout << "!!! Minimisation failed!" << endl;
    return;
  }

  //*****************************************//
  //******** Injection pulse results ********//    
  //*****************************************//

  // Display calculated injection pulse
  cout << "\n Calculation of injection pulse complete!" << endl;
  cout << "\n * "  << "Injection Bin" << "		"  << " * "  << "Intensity [AU]" << endl;
  cout << "---------------------------------------------" << endl;
  for (int k = 0; k < noInjTimeBins; k++) {
    cout << "   "  << k << "		"  << "   "  << I[k] << endl;
  }

  TCanvas *InjPulseFinal = new TCanvas("InjPulseFinal","InjPulse Final",800,600);
  TH1D *injPulseFinal = new TH1D("injPulseFinal", "Injection Pulse;Time [ns]", noInjTimeBins, -0.5*(noInjTimeBins+1)*injBinWidth-injt0, 0.5*(noInjTimeBins-1)*injBinWidth-injt0);
  for(int i = 1; i <= noInjTimeBins; i++) injPulseFinal->SetBinContent(i,I[i-1]);
  injPulseFinal->SetLineColor(1);
  injPulseFinal->SetFillColor(5);
  injPulseFinal->Draw("p*CF");

  if(!realData){
    TH1D* inputInjPulse = RescaleAxis((TH1D*)file->Get("hT0Pulse")->Clone("hInputInjPulse"), 1e-3); // Convert input to us
    inputInjPulse->Scale(injPulseFinal->Integral("WIDTH")/inputInjPulse->Integral("WIDTH"));
    inputInjPulse->SetLineColor(2);
    inputInjPulse->Draw("HISTSAME");
  }

  //*****************************************//                              
  //****** Radial distribution results ******//        
  //*****************************************// 

  // Display calculated radial distribution
  cout << "\n Calculation of radial distribution complete!" << endl;
  cout << "\n * "  << "Radius[mm]" << "		"  << " * "  << "Intensity [AU]" << endl;
  cout << "---------------------------------------------" << endl;
  for (int n = 0; n < noRadBins; n++) {
    cout << "   "  << radius_i[n] << "		"  << "   "  << f[n] << endl;
  }
  // Plot radial distribtion
  TCanvas *RadialDistributionFinal = new TCanvas("RadialDistributionFinal","Radial Distribution Final",800,600);
  double radBinLowEdges[noRadBins+1];
  radBinLowEdges[0] = rMagic()-45;
  for(int i = 1; i <= noRadBins; i++) radBinLowEdges[i] = radBinLowEdges[i-1] + radBinWidth[i-1];
  TH1D *hRadDistributionFinal = new TH1D("hRadDistributionFinal","Radial Distribution;Radius [mm];Arbitrary Units [au]", noRadBins,  radBinLowEdges);
  for(int i = 0; i < noRadBins; i++) hRadDistributionFinal->SetBinContent(i+1,f[i]);
  hRadDistributionFinal->SetLineColor(1);
  hRadDistributionFinal->SetFillColor(5);
  hRadDistributionFinal->SetStats(0);
  hRadDistributionFinal->Draw("p*CF");

  if(!realData){
    TH1D* inputRadialDistribution = RescaleAxis((TH1D*)(file->Get("hRadius")->Clone("hInputRadius")), 10);  // Convert input to mm
    inputRadialDistribution->Scale(hRadDistributionFinal->Integral("WIDTH")/inputRadialDistribution->Integral("WIDTH"));
    inputRadialDistribution->SetLineColor(2);
    inputRadialDistribution->Draw("HISTSAME");
  }

  // Determine Radial radial mean and rms
  cout << "\n Determining the mean, RMS and E-field correction from the calculated radial distribution:" << endl;
  double sum = 0.0;
  double squares = 0.0;
  double count = 0.0;
  for (int i = 0; i < noRadBins; i++) {
    count += f[i];
    sum += radius_i[i]*f[i];
    squares += radius_i[i]*radius_i[i]*f[i];
  }
  double mean = sum/count;
  double var = (squares/count)-(mean*mean);
  double rms = sqrt(var);
  // Determine E-field correction
  double Beta = 0.9994174214209439;
  double xesq = 0.0;
  double C_E = 0.0;
  for (int i = 0; i < noRadBins; i++){
    xesq += (radius_i[i]-rMagic())*(radius_i[i]-rMagic())*f[i];
  }
  xesq = xesq/count;
  // E-field correction, C_E                                                
  C_E = -2.0*avg_n1mn*pow(Beta,2)*xesq/pow(rMagic(),2);
  //Display calculated mean, rms and C_E                                  
  cout << "\n\t --> Mean = " << mean << " mm, RMS = " << rms << "mm, C_E = " << C_E*1e9  << "ppb\n" << endl;
  // Update radial distribution plot with fit info and C_E
  TPaveText* fitInfo = new TPaveText(7070., 0.7*hRadDistributionFinal->GetBinContent(hRadDistributionFinal->GetMaximumBin()), 7100., hRadDistributionFinal->GetBinContent(hRadDistributionFinal->GetMaximumBin()));
  fitInfo->AddText(Form("Preferred injection t0 = %.2f ns" ,1e3*injPulseFinal->GetMean()));
  fitInfo->AddText(Form("Analysis start time = %g #mus" ,FRstart));
  fitInfo->AddText(Form("Analysis end time = %g #mus" , analysisEnd));
  fitInfo->AddText(Form("#chi^{2}_{min}/d.o.f. = %g" ,chisq_dof));
  fitInfo->AddText(Form("Mean radius = %g mm" ,mean));
  fitInfo->AddText(Form("RMS = %g mm" ,rms));
  fitInfo->AddText(Form("r_mean - r_magic = %g mm" ,mean - rMagic()));
  fitInfo->AddText(Form("C_E = %g ppb", C_E*1e9));
  fitInfo->SetFillColor(33);
  fitInfo->Draw("SAME");
  TLine *r_magicLine = new TLine(rMagic(),0,rMagic(),hRadDistributionFinal->GetBinContent(hRadDistributionFinal->GetMaximumBin()));
  r_magicLine->SetLineColor(kBlue);
  r_magicLine->SetLineWidth(3);
  r_magicLine->Draw();
  TLine *r_meanLine = new TLine(mean,0,mean,hRadDistributionFinal->GetBinContent(hRadDistributionFinal->GetMaximumBin()));
  r_meanLine->SetLineColor(kRed);
  r_meanLine->SetLineWidth(3);
  r_meanLine->Draw();

  //*****************************************//                              
  //***** Momentum distribution results *****//      
  //*****************************************//

  // Calculate Momentum distribution                                     
  double MomBinLowEdges[noRadBins+1];
  double radius = 0.0;
  radius = radBinLowEdges[0];
  MomBinLowEdges[0] = r2p(radius,quadNvalue)*1e-3;
  for(int i = 1; i <= noRadBins; i++){
    radius = radBinLowEdges[i];
    MomBinLowEdges[i] = r2p(radius,quadNvalue)*1e-3;
  }
  TH1D *hMomDistribution = new TH1D("hMomDistribution","Momentum Distribution", noRadBins, MomBinLowEdges);
  hMomDistribution->GetXaxis()->SetTitle("Momentum [GeV/c]");
  hMomDistribution->GetYaxis()->SetTitle("arbitrary units [au]");
  for(int i = 0; i < noRadBins; i++) {
    hMomDistribution->SetBinContent(i+1,f[i]);
  }
  // Display calculated Momentum distribution                            
  cout << "\n Calculation of Momentum distribution complete" << endl;
  cout << "\n * "  << "Momentum [GeV/c]" << "         "  << " * "  << "Intensity [AU]" << endl;
  cout << "---------------------------------------------" << endl;
  for (int n = 0; n < noRadBins; n++) {
    cout <<  hMomDistribution->GetBinCenter(n+1) << "               "  << "   "  << hMomDistribution->GetBinContent(n+1) << endl;
  }
  // Calculate mean of momentum distribution                              
  double sum_Mom = 0.0;
  double squares_Mom = 0.0;
  double count_Mom = 0.0;
  for (int i = 0; i < noRadBins; i++) {
    count_Mom += hMomDistribution->GetBinContent(i);
    sum_Mom += hMomDistribution->GetBinCenter(i)*hMomDistribution->GetBinContent(i);
    squares_Mom += hMomDistribution->GetBinCenter(i)*hMomDistribution->GetBinCenter(i)*hMomDistribution->GetBinContent(i);
  }
  double mean_Mom = sum_Mom/count_Mom;
  double var_Mom = (squares_Mom/count_Mom)-(mean_Mom*mean_Mom);
  double rms_Mom = sqrt(var_Mom);
  TCanvas *MomentumDistributionFinal = new TCanvas("MomentumDistribution","MomentumDistribution",800,600);
  hMomDistribution->SetLineColor(1);
  hMomDistribution->SetFillColor(5);
  hMomDistribution->Draw("p*CF");
  //double CERN3max = hMomDistribution->GetBinContent(hMomDistribution->GetMaximumBin());
  //  double FTmax = FTmomDist->GetBinContent(FTmomDist->GetMaximumBin());
  //  double FTnorm = CERN3max/FTmax;
  //  for (int i = 0; i < FTmomDist->GetNbinsX(); i++){
  //    FTmomDist->SetBinContent(i,FTmomDist->GetBinContent(i)*FTnorm);
  //  }
  //  FTmomDist->SetLineColor(2);
  //  FTmomDist->Draw("SAME");
  TLegend* momLegend = new TLegend(0.6,0.7,0.95,0.9);
  momLegend->AddEntry(hMomDistribution,"CERNIII result","lp");
  //  momLegend->AddEntry(FTmomDist,"Truth","l");
  momLegend->Draw("SAME");
  TPaveText* MomFitInfo = new TPaveText(hMomDistribution->GetBinCenter(1),0.7*hMomDistribution->GetBinContent(hMomDistribution->GetMaximumBin()), hMomDistribution->GetBinCenter(11), hMomDistribution->GetBinContent(hMomDistribution->GetMaximumBin()));
  MomFitInfo->AddText(Form("Preferred injection t0 = %g #mus" ,injt0));
  MomFitInfo->AddText(Form("Analysis start time = %g #mus" ,FRstart));
  MomFitInfo->AddText(Form("Analysis end time = %g #mus" , analysisEnd));
  MomFitInfo->AddText(Form("#chi^{2}_{min}/d.o.f. = %g" ,chisq_dof));
  MomFitInfo->AddText(Form("Mean Momentum = %g GeV/c" ,mean_Mom));
  MomFitInfo->AddText(Form("RMS = %g GeV/c" ,rms_Mom));
  MomFitInfo->AddText(Form("p_mean - p_magic = %g GeV/c" ,mean_Mom-pMagic()));
  MomFitInfo->AddText(Form("C_E = %g ppb", C_E*1e9));                      
  MomFitInfo->SetFillColor(33);
  MomFitInfo->Draw("SAME");

  //*****************************************//                              
  //**** Frequency distribution results *****//                           
  //*****************************************// 

  // Calculate frequency distribution                                     
  double freqBinLowEdges[noRadBins+1];
  radius = 0.0;
  radius = radBinLowEdges[0];
  freqBinLowEdges[noRadBins] = t2f(r2t(radius));
  for(int i = 1; i <= noRadBins; i++){
    radius = radBinLowEdges[i];
    freqBinLowEdges[noRadBins-i] = t2f(r2t(radius));
  }
  TH1D *hFreqDistribution = new TH1D("hFreqDistribution","Frequency Distribution", noRadBins, freqBinLowEdges);
  hFreqDistribution->GetXaxis()->SetTitle("Frequency [kHz]");
  hFreqDistribution->GetYaxis()->SetTitle("arbitrary units [au]");
  for(int i = 1; i <= noRadBins; i++) {
    hFreqDistribution->SetBinContent(i, f[noRadBins-i]);
  }
  // Display calculated frequency distribution                            
  cout << "\n Calculation of frequency distribution complete" << endl;
  cout << "\n * "  << "Frequency[kHz]" << "         "  << " * "  << "Intensity [AU]" << endl;
  cout << "---------------------------------------------" << endl;
  for (int n = 0; n < noRadBins; n++) {
    cout << "   "  << hFreqDistribution->GetBinCenter(n+1) << "             "  << "   "  << hFreqDistribution->GetBinContent(n+1) << endl;
  }  
  // Calculate mean of frequency distribution                             
  double sum_freq = 0.0;
  double squares_freq = 0.0;
  double count_freq = 0.0;
  for (int i = 0; i < noRadBins; i++) {
    count_freq += hFreqDistribution->GetBinContent(i+1);
    sum_freq += hFreqDistribution->GetBinCenter(i+1)*hFreqDistribution->GetBinContent(i+1);
    squares_freq += hFreqDistribution->GetBinCenter(i+1)*hFreqDistribution->GetBinCenter(i+1)*hFreqDistribution->GetBinContent(i+1);
  }
  double mean_freq = sum_freq/count_freq;
  double var_freq = (squares_freq/count_freq)-(mean_freq*mean_freq);
  double rms_freq = sqrt(var_freq);
  // Plot calculate frequency distribution                                
  TCanvas *FrequencyDistributionFinal = new TCanvas("FrequencyDistribution","FrequencyDistribution",800,600);
  hFreqDistribution->SetLineColor(1);
  hFreqDistribution->SetFillColor(5);
  hFreqDistribution->Draw("p*CF");
  //  CERN3max = hFreqDistribution->GetBinContent(hFreqDistribution->GetMaximumBin());
  //  FTmax = FTfreqDist->GetBinContent(FTfreqDist->GetMaximumBin());
  //  FTnorm = CERN3max/FTmax;
  //for (int i = 0; i < FTfreqDist->GetNbinsX(); i++){
  //    FTfreqDist->SetBinContent(i,FTfreqDist->GetBinContent(i)*FTnorm);
  //  }
  //  FTfreqDist->SetLineColor(2);
  //  FTfreqDist->Draw("SAME");
  TLegend* freqLegend = new TLegend(0.6,0.7,0.95,0.9);
  freqLegend->AddEntry(hFreqDistribution,"CERNIII result","l");
  //  freqLegend->AddEntry(FTfreqDist,"Truth","l");
  freqLegend->Draw("SAME");
  TPaveText* freqFitInfo = new TPaveText(hFreqDistribution->GetBinCenter(1), 0.7*hFreqDistribution->GetBinContent(hFreqDistribution->GetMaximumBin()),hFreqDistribution->GetBinCenter(11), hFreqDistribution->GetBinContent(hFreqDistribution->GetMaximumBin()));
  freqFitInfo->AddText(Form("Preferred injection t0 = %g #mus" ,injt0));
  freqFitInfo->AddText(Form("Analysis start time = %g #mus" ,FRstart));
  freqFitInfo->AddText(Form("Analysis end time = %g #mus" , analysisEnd));
  freqFitInfo->AddText(Form("#chi^{2}_{min}/d.o.f. = %g" ,chisq_dof));
  freqFitInfo->AddText(Form("Mean frequency = %g kHz" ,mean_freq));
  freqFitInfo->AddText(Form("RMS = %g kHz" ,rms_freq));
  freqFitInfo->AddText(Form("f_mean - f_magic = %g kHz" ,mean_freq-t2f(r2t(rMagic()))));
  freqFitInfo->AddText(Form("C_E = %g ppb", C_E*1e9));                      
  freqFitInfo->SetFillColor(33);
  freqFitInfo->Draw("SAME");

}