// ----------------------------------------------------------------------------------- //
/*
  ROOT macro for reading text (ASCII) files and fitting data written.

  The data is from a damped harmonic oscillator, in this case a spring hanging with a
  weight but also a round piece of cardboard, which increases the drag.

  The basic fit is to a simple damped harmonic oscillator, which has parameters:
   1: Offset constant (it is not swinging around 0!).
   2: Amplitude
   3: Damping coefficient
   4: Phase
   5: Frequency

  Author: Troels C. Petersen (NBI)
  Email:  petersen@nbi.dk
  Date:   4th of October 2011
*/
// ----------------------------------------------------------------------------------- //

#include <TFile.h>
#include <fstream>
#include <iostream>


// ----------------------------------------------------------------------------------- //
void HarmonicOscillator_Fys2Lab_FitData_Solution() {
// ----------------------------------------------------------------------------------- //

  // Setting of general plotting style:
  gStyle->SetCanvasColor(0);
  gStyle->SetFillColor(0);
  // Setting what to be shown in statistics box:
  gStyle->SetOptStat("e");
  gStyle->SetOptFit(1111);


  //---------------------------------------------------------------------------------- 
  // Reading the data:
  //---------------------------------------------------------------------------------- 
  int NeventsMax = 3617;
  int Nevents = 0;
  double t, x;
  char name[80];

  double t_min = 0.0;
  double t_max = 0.01 * double(NeventsMax-1);
  TH1F* Hist_Osc = new TH1F("Hist_Osc", "Hist_Osc", NeventsMax, t_min-0.005, t_max+0.005);
  TH1F* Resi_Osc = new TH1F("Resi_Osc", "Resi_Osc", NeventsMax, t_min-0.005, t_max+0.005);

  // File to read:
  ifstream infile("Exp_Osc_V2.txt");
  if (infile.is_open()) {
    while ((!infile.eof()) && (Nevents < NeventsMax)) {
      infile >> t >> x;
      Hist_Osc->SetBinContent(Nevents+1, x);       // Value measured at time t
      Hist_Osc->SetBinError(Nevents+1, 0.0025);    // Error ESTIMATED for this value!

      // Write out the first 10 measurements:
      if (Nevents < 10) printf("  %2d. events read: %6.3f  %6.3f \n", Nevents, t, x);
      Nevents++;
    }
  }
  infile.close();
  printf("  Number of entries in total: %6d \n", Nevents);



  //---------------------------------------------------------------------------------- 
  // Intermediate fits
  //---------------------------------------------------------------------------------- 

  printf("\n  Fitting ranges of the data: \n");
  TCanvas* c1 = new TCanvas("c1", "", 50, 20, 1200, 500);

  const int Nint = 24;
  Hist_Osc->Clone("Hist_Osc_InRanges");    // Make a copy of the data "for playing around"...
  TF1 *fitOscExp[Nint];                    // Make Nint fitting functions (same f(x), but different ranges)
  double tmean[Nint], etmean[Nint], tmin, tmax;
  double ampli[Nint], eampli[Nint];
  double expon[Nint], eexpon[Nint];
  double omega[Nint], eomega[Nint];
  double phase[Nint], ephase[Nint];
  double Chi2_fullrange = 0.0;

  for (int i=0; i < Nint; i++) {
    tmin = (36.165 / double(Nint)) * double(i);
    tmax = (36.165 / double(Nint)) * double(i+1);
    sprintf(name, "fitOscExp%02d", i);
    fitOscExp[i] = new TF1(name, "[0] + [1] * exp(-[2]*x) * cos([3] + [4]*x)", tmin, tmax);
    fitOscExp[i]->SetParameters(0.037, 0.15, 0.082, -1.0, 9.0);
    fitOscExp[i]->SetLineWidth(2);
    fitOscExp[i]->SetLineColor(i%5+2);      // Making the colors cycle in values 2-6
    fitOscExp[i]->SetNpx(1000);             // Granularity of drawing the function
    Hist_Osc_InRanges->Fit(name, "RQ+");

    // Record fitting result for interval:
    tmean[i] = (tmax + tmin) / 2.0;
    etmean[i] = 0.0;
    ampli[i] = fitOscExp[i]->GetParameter(1);
    eampli[i] = fitOscExp[i]->GetParError(1);
    expon[i] = fitOscExp[i]->GetParameter(2);
    eexpon[i] = fitOscExp[i]->GetParError(2);
    phase[i] = fitOscExp[i]->GetParameter(3);
    ephase[i] = fitOscExp[i]->GetParError(3);
    omega[i] = fitOscExp[i]->GetParameter(4);
    eomega[i] = fitOscExp[i]->GetParError(4);

    printf("  %2d fit:   %4.1f - %4.1f     Chi2 = %5.1f    Prob = %5.3f \n",
	   i, tmin, tmax, fitOscExp[i]->GetChisquare(), fitOscExp[i]->GetProb());
    Chi2_fullrange += fitOscExp[i]->GetChisquare();
  }

  printf("  Total Chi2 of all functions (i.e. full range):  %6.1f \n", Chi2_fullrange);
  Hist_Osc_InRanges->Draw("e");

  c1->Update();
  // c1->SaveAs("Fys2Lab_HarmonicOscillator_OscExpFit_InRanges.png");


  //---------------------------------------------------------------------------------- 
  // Plot the results of the interval fits:
  TCanvas* c2 = new TCanvas("c2", "", 100, 40, 900, 600);
  c2->Divide(2,2);

  c2->cd(1);
  TGraphErrors* g_ampli = new TGraphErrors(Nint, tmean, ampli, etmean, eampli);
  g_ampli->SetLineColor(4);
  g_ampli->SetLineWidth(2);
  g_ampli->Draw("AP");

  c2->cd(2);
  TGraphErrors* g_expon = new TGraphErrors(Nint, tmean, expon, etmean, eexpon);
  g_expon->SetLineColor(4);
  g_expon->SetLineWidth(2);
  g_expon->Draw("AP");

  c2->cd(3);
  TGraphErrors* g_phase = new TGraphErrors(Nint, tmean, phase, etmean, ephase);
  g_phase->SetLineColor(4);
  g_phase->SetLineWidth(2);
  g_phase->Draw("AP");

  c2->cd(4);
  TGraphErrors* g_omega = new TGraphErrors(Nint, tmean, omega, etmean, eomega);
  g_omega->SetLineColor(4);
  g_omega->SetLineWidth(2);
  g_omega->Draw("AP");

  c2->Update();
  // c2->SaveAs("Fys2Lab_HarmonicOscillator_OscExpFit_InRangesResults.png");
  /*
  */


//---------------------------------------------------------------------------------- 
// Basic and final fit:
//---------------------------------------------------------------------------------- 

  TCanvas* c0 = new TCanvas("c0", "", 150, 60, 1200, 500);
  // c0->Divide(1,2);

  // Fitting histogram with basic dampled harmonic oscillating function:
  // NOTE: Give reasonable starting values is essential!!!
  TF1 *fitBasic = new TF1("fitBasic", "[0] + [1]*exp(-[2]*x) * cos([3] + [4]*x)", 0.005, 36.165);
  fitBasic->SetParameters(0.038, 0.19, 0.082, 0.0, 8.97);
  fitBasic->SetLineColor(4);
  fitBasic->SetLineWidth(2);
  fitBasic->SetNpx(1000);

  // Fitting histogram with oscillating function, which has two dampling coefficients/exponentials (before and after
  // change around t=15s most likely due to turbulence), a slowly changing frequency, and a slow oscillation in amplitude:
  TF1 *fitFull = new TF1("fitFull", "[0] + ([1]*exp(-[2]*x) + [9]*exp(-[10]*x)) * cos([3]+[4]*x+[5]*x*x) * (1.0 + [6]*cos([7]*x+[8]))", 0.005, 36.165);
  fitFull->SetParameters(0.038, 0.19, 0.104, -1.0, 8.97, 0.0007, 0.06, 0.4, 1.0, 0.02, 0.01);
  fitFull->SetLineColor(2);
  fitFull->SetLineWidth(2);
  fitFull->SetNpx(1000);

  // Setting range, line width and color and axis titles of histograms:
  Hist_Osc->GetXaxis()->SetRangeUser(-0.005, 36.165);
  Hist_Osc->GetXaxis()->SetTitle("Time (seconds)");
  Hist_Osc->GetYaxis()->SetTitle("Oscillation (arbitrary scale)");


  // Fitting and drawing histograms with errors (and separate stat boxes):
  Hist_Osc->Clone("Hist_Osc_Full");
  Hist_Osc_Full->Fit("fitFull", "R");
  c0->Update();
  TPaveStats *st_full = (TPaveStats*)Hist_Osc_Full->FindObject("stats");
  st_full->SetY1NDC(0.60); st_full->SetY2NDC(0.98); st_full->SetTextColor(2);

  Hist_Osc->Clone("Hist_Osc_Basic");
  Hist_Osc_Basic->Fit("fitBasic", "R");
  c0->Update();
  TPaveStats *st_basic = (TPaveStats*)Hist_Osc_Basic->FindObject("stats");
  st_basic->SetY1NDC(0.13); st_basic->SetY2NDC(0.38); st_basic->SetTextColor(4);

  Hist_Osc_Basic->Draw("e");
  Hist_Osc_Full->Draw("e same");
  c0->Modified();
  // c0->SaveAs("Fys2Lab_HarmonicOscillator_OscExpFit.png");


  printf("\n  Result for comparable parameters of basic and full fit: \n");
  for (int i=0; i < 5; i++) {
    printf("  Par %1d:   Basic = %8.5f +- %8.5f    Full = %8.5f +- %8.5f \n", i,
	   fitBasic->GetParameter(i), fitBasic->GetParError(i),
	   fitFull->GetParameter(i), fitFull->GetParError(i));
  }
  printf("\n");



//---------------------------------------------------------------------------------- 
// Residuals resulting from the full fit:
//---------------------------------------------------------------------------------- 
  TCanvas* c3 = new TCanvas("c3", "", 200, 80, 1200, 500);

  for (int i=0; i < NeventsMax; i++) {
    t = Hist_Osc->GetBinCenter(i);
    Resi_Osc->SetBinContent(i, Hist_Osc->GetBinContent(i) - fitFull->Eval(t));
    Resi_Osc->SetBinError(i, Hist_Osc->GetBinError(i));
  }

  // Setting range, line width and color and axis titles of histograms:
  Resi_Osc->GetXaxis()->SetRangeUser(-0.005, 36.165);
  Resi_Osc->GetXaxis()->SetTitle("Time (seconds)");
  Resi_Osc->GetYaxis()->SetTitle("Oscillation residuals (arbitrary scale)");

  // Comparing constant fit to a harmonic oscillator fit with amplitude modulation:
  TF1 *fitConst = new TF1("fitConst", "[0]", 0.025, 36.165);
  fitConst->SetLineColor(2);
  Resi_Osc->Fit("fitConst", "RQ");


  // Dividing the fitting range, and simply fit with a harmonic oscillator:
  const int Nint2 = 12;
  TF1 *fitOsc[Nint2];
  double tmin2, tmax2;
  for (int i=0; i < Nint2; i++) {
    tmin2 = (36.165 / double(Nint2)) * double(i);
    tmax2 = (36.165 / double(Nint2)) * double(i+1);
    sprintf(name, "fitOsc%02d", i);
    fitOsc[i] = new TF1(name, "[0] + [1] * cos([2] + [3]*x)", tmin2, tmax2);
    fitOsc[i]->SetParameters(0.0, 0.002, 0.0, 9.0);
    fitOsc[i]->SetLineWidth(2);
    fitOsc[i]->SetLineColor(i%5+2);
    fitOsc[i]->SetNpx(1000);
    Resi_Osc->Fit(name, "RQ+");
  }

  Resi_Osc->Draw("e");

  c3->Update();
  // c3->SaveAs("Fys2Lab_HarmonicOscillator_OscExpFit_Residuals.png");
  /*
  */

}



// ----------------------------------------------------------------------------------- //
/*

Questions:
----------
 1) Look at the data file and see if you can by eye estimate the size of the uncertainty of the
    points. It is not easy, but you should be able to get it to within a factor of 2-3.

 2) Play with the fitting function and introduce terms to reduce the Chi2.
    See if you in a fit of the range [0.005,36.005] can get a better Chi2 than 6600.

 3) If you were to measure the oscillation frequency for (infinitely) small oscillations,
    what would your estimate be? Can you make a confidence interval?


Advanced questions:
-------------------
  1) Can you detect a change in the error as a function of time? I.e. does the errors
     needed to get a reasonable/constant Chi2 get better or worse if you fit small ranges
     in time.

*/
// ----------------------------------------------------------------------------------- //