//-----------------------------------------------------------------------------------------
//-----------------------------------------------------------------------------------------
//  
//  173.308.02 - Advanced Physics Lab
//  Professor: Tobias Marriage
//  
//  Franck_Hertz_Analysis.C (bla version 2)
//  Written by Jake Mokris, 2/20/2011
//  E-mail: jmokris1@jhu.edu
//
//  Adapted from Franck_Hertz.C, a macro by Prof. Petar Maksimovic. His macro can be found
//  at http://www.pha.jhu.edu/~c173_608/franck-hertz/root/Franck_Hertz.C. 
//  
//  This program performs data analysis for the Franck-Hertz experiment: it reads in a
//  data file (formatted in two columns - the first column must be the voltage, and the
//  second must be the current corresponding to that voltage), plots the data,  fits it
//  to a sum of gaussians with a quadratic background, and saves the resulting graph to 
//  the file Franck_Hertz_Fit.png. The number of gaussians is variable, and the program 
//  allows for a linear background as well. 
//  
//  To run this program, open ROOT and type
//  
//  root[0] .L Franck_Hertz_Analysis.C++
//  root[1] Franck_Hertz_Analysis()
//  
//  You will be prompted for the name of a data file, which should be in the same
//  directory as this macro. After that, there will be many more prompts; just do what
//  they say! 
//
//  When inputting values for the fit function parameters, you will probably find that 
//  your initial guesses aren't too great, resulting in a poor fit to the data. You may
//  want to fit each peak one at a time, to get more accurate parameter values (the
//  program outputs the final parameter values after it fits). 
//  
//  I am sure that this program will work IF you set up ROOT properly. If you're having
//  trouble with ROOT, CHECK THE MANUAL FIRST to see if it's set up properly (you can find
//  the manual at the ROOT website, root.cern.ch). ROOT should be set up on the PUC lab
//  computers. If you're still having trouble, e-mail me and I'll probably answer before 
//  the day is over. Again, if you have any suggestions as to how the program can be 
//  improved (I'm a novice programmer), please tell me. 
//  
//-----------------------------------------------------------------------------------------
//-----------------------------------------------------------------------------------------


#include<fstream>
#include<sstream>
#include<iostream>
#include<TCanvas.h>
#include<TF1.h>
#include<TMath.h>
#include<TGraph.h>
#include<TAxis.h>

using namespace std;


//-----------------------------------------------------------------------------------------
// The ifstream constructor (and various other constructors) accept only const char *, so
// convert takes a string s and returns a char *. This is where my knowledge of C++ fails;
// I'm not sure if such a function violates proper programming style, transgresses some 
// property of C++, or is just plain silly. I have this function because I want the 
// program to prompt for file names and such, and I don't know how to read in a char * 
// object directly from standard input. I do know how to read in strings. 
//-----------------------------------------------------------------------------------------

char *convert(string s)
{
  const int length = s.length();
  char *name = new char[length+1]; // Add one to the length to hold the terminal character
  for (Int_t i = 0; i<length; i++)
  {
    name[i] = s[i];
  }
  name[length]='\0'; // Tack the terminal character onto the end
  return(name);
}

//-----------------------------------------------------------------------------------------
// A gaussian function
//-----------------------------------------------------------------------------------------

Double_t gaussian(Double_t x, Double_t norm, Double_t mean, Double_t sigma)
{
  if (sigma != 0) {
    Double_t arg = (x-mean)/sigma;
    return(norm*TMath::Exp(-0.5*arg*arg));
  }
  else {
    return(0);
  }
}

//-----------------------------------------------------------------------------------------
// The function to be fitted to the data. It consists of n gaussians atop a quadratic
// polynomial background. For the function to be fitted to data, ROOT requires it to 
// accept as arguments two pointers to double arrays, *x and *par. *x is the array of 
// values inputted into the function; since this function is from $\mathbb{R}$ (the LaTeX
// way to write the symbol for the reals) to $\mathbb R$, we only need the first element
// of the array, x[0]. The *par array indicates the parameters to be fitted to the data.
// There are 3*n + 4 parameters to be fitted in the function I wrote; they are as follows:
//
// par[0] = n, the number of gaussians
//
// For int i = 1 and i < n + 1, the parameters refer to the six gaussians:
// 
// par[3*i-2] is the norm - the area - under the (i+1)th gaussian
// par[3*i-1] is the mean - the location on the x-axis - of the (i+1)th gaussian
// par[3*i] is the sigma - the width - of the (i+1)th gaussian
// 
// The last three parameters are for the quadratic background:
// 
// background = par[3*par[0]+1] + par[3*par[0]+2]*x[0] + par[3*par[0]+3]*x[0]*x[0]
//
// Since the fit function itself is the sum of all the gaussians and the background, the 
// function returns the sum of the values of each gaussian and the background at x[0]. 
//-----------------------------------------------------------------------------------------

Double_t fit_function(Double_t *x, Double_t *par)
{
  
  Double_t value = 0; // The value to be returned
  Int_t ng = par[0]; // Number of gaussians 
  // Add together the contribution of each gaussian
  for(Int_t i = 1; i < ng + 1; i++)
  {
    value += gaussian(x[0], par[3*i-2], par[3*i-1], par[3*i]);
  }
  // Add on the contribution of the background
  value += par[3*ng+1] + par[3*ng+2]*x[0] + par[3*ng+3]*x[0]*x[0];
  return(value);
}

//-----------------------------------------------------------------------------------------
// ROOT does not execute the content of main when running C++ code; the proper way to run
// C++ code in ROOT is to put it into a void function and call the function. 
//-----------------------------------------------------------------------------------------

void Franck_Hertz_Analysis()
{
  // Get the name of the file
  string file_string;
  cout << "Enter the name of the data file you wish to read: ";
  cin >> file_string;
  
  // Create ifstream object to input data from file
  const char *file_name = convert(file_string);
  ifstream data_file(file_name);
  
  // Check that the file opened properly
  if (!data_file) 
  {
    cerr << "Failed to open the file.\n";
    exit(EXIT_FAILURE);
  }

  //---------------------------------------------------------------------------------------
  // Count the number of lines of data in the file. Right now this just ignores a large 
  // number of characters until it hits the end-of-line character. Clearly it needs to 
  // ignore an arbitrary number of characters until reaching the end of the line, but I 
  // haven't worked on implementing that yet.
  //---------------------------------------------------------------------------------------

  Int_t number_of_lines = 0;
  while(data_file.ignore(9001, '\n')) // It's over 9000!
  {
    if (!data_file.eof())
    {
    number_of_lines++;
    }
  }
  cout << "The data file " << file_name << " contains " << number_of_lines;
  cout << " data points." << endl;
  
  // Return to beginning of file; input data
  data_file.clear();
  data_file.seekg(0);
  Double_t x_coordinate, y_coordinate;
  Double_t x[number_of_lines], y[number_of_lines];
  
  cout << "Reading data file " << file_name << "." << endl;
  
  for(Int_t i=0; i<number_of_lines && !data_file.eof(); i++)
  {
    data_file >> x_coordinate;
    data_file >> y_coordinate;
    
    if (!data_file.eof())
    {
      x[i]=x_coordinate*10; // The x-coordinate data is in units of 10V
      y[i]=y_coordinate;
    }
  }
  
  cout << "Data input complete." << endl;
  
  // Close the file
  data_file.close();
    
  // Create the TGraph object, and a TCanvas to put it on
  TCanvas *canvas = new TCanvas("Franck-Hertz Data","Franck-Hertz Data",600,600);
  TGraph *graph = new TGraph(number_of_lines, x, y);
  
  // To modify the graph settings (axis titles and such), the graph must be drawn first
  graph->Draw();
  
  // Make the graph fancy
  graph->SetTitle(""); // This gets rid of the pesky graph title box
  // graph->SetMarkerStyle(20); // Sets the point marker to a filled black circle
  
  // x-axis 
  Double_t xmin, xmax;
  cout << "Enter minimum value for x-axis: ";
  cin >> xmin;
  cout << "Enter maximum value for x-axis: ";
  cin >> xmax;
  graph->GetXaxis()->SetTitle("Electron Acceleration Voltage (V)"); // Set x-axis title
  graph->GetXaxis()->SetLimits(xmin,xmax); // Set x-axis range
  graph->GetXaxis()->CenterTitle(); // Axis titles are initially off-center
  
  // y-axis
  graph->GetYaxis()->SetTitle("Current (nA)"); // Set y-axis title
  graph->GetYaxis()->CenterTitle(); // Axis titles are initially off-center
  
  //---------------------------------------------------------------------------------------
  // Draw options:
  // 
  // A - Include axes
  // P - Draw a point marker at each data point
  //---------------------------------------------------------------------------------------
  
  // Redraw the graph
  graph->Draw("AP"); 
  canvas->Update();

  //---------------------------------------------------------------------------------------
  // Create an instance of fit_function, defined above. The TF1 parameters are name of
  // the object, type of function you're making, xmin, xmax, and number of parameters
  // the function type has (which is 3*number_of_gaussians + 4).  
  //---------------------------------------------------------------------------------------
  
  Double_t norm, mean, sigma, p_0, p_1, p_2;
  cout << "To perform the fit, input approximate values for the following ";
  cout << "parameters. Read them off of the graph you just made." << endl;

  Int_t number_of_gaussians;
  cout << "Enter number of gaussians you wish to fit to the data: ";
  cin >> number_of_gaussians;

  TF1 *f = new TF1("Fit function", fit_function, xmin, xmax, 3*number_of_gaussians + 4);
  f->FixParameter(0, number_of_gaussians); // Fixes par[0] to number_of_gaussians

  //---------------------------------------------------------------------------------------
  // Before you use the function, you have to set the parameters to some initial values. 
  // This prompts the user for approximate values for the parameters.
  //---------------------------------------------------------------------------------------

  // Get parameter values for the six gaussians
  for (Int_t i = 1; i < number_of_gaussians + 1; i++)
  {
    // Create stringstreams
    stringstream normname, meanname, sigmaname;
    normname << "Norm";
    meanname << "Mean";
    sigmaname << "Sigma";
    
    // Read in value for the norm of the ith gaussian
    cout << "Enter area under peak " << i << ": ";
    cin >> norm;
    normname << i; // This gives the parameter the name "Norm(i)"
    f->SetParName(3*i - 2, convert(normname.str())); // Set the name of the parameter
    f->SetParameter(3*i - 2, norm); // Set the value of the parameter

    // Read in value for the center of the ith gaussian
    cout << "Enter x-coordinate of peak " << i << ": ";
    cin >> mean;
    meanname << i;
    f->SetParName(3*i - 1, convert(meanname.str()));
    f->SetParameter(3*i - 1, mean);
    
    // Read in value for the sigma of the ith gaussian
    cout << "Enter width of peak " << i << ": ";
    cin >> sigma;
    sigmaname << i;
    f->SetParName(3*i, convert(sigmaname.str()));
    f->SetParameter(3*i, sigma);
  }
  
  // Read in coefficients of the background polynomial  
  cout << "Enter value for coefficient of x^0 in background polynomial: ";
  cin >> p_0;
  f->SetParName(3*number_of_gaussians + 1, "p_0");
  f->SetParameter(3*number_of_gaussians + 1, p_0);

  cout << "Enter value for coefficient of x^1 in background polynomial: ";
  cin >> p_1;
  f->SetParName(3*number_of_gaussians + 2, "p_1");
  f->SetParameter(3*number_of_gaussians + 2, p_1);
  
  // Toggle linear/quadratic background
  int line;
  cout << "Enter 0 to fit background to a line: ";
  cin >> line;
  if (line == 0)
  {
    f->FixParameter(3*number_of_gaussians + 3, 0);
  }
  else 
  {
  cout << "Enter value for coefficient of x^2 in background polynomial: ";
  cin >> p_2;
  f->SetParName(3*number_of_gaussians + 3, "p_2");
  f->SetParameter(3*number_of_gaussians + 3, p_2);
  }
  
  // Fit the function to the data and plot it on the graph
  graph->Fit(f, "R");

  // Save the image
  canvas->SaveAs("Franck_Hertz_Fit.png");
}