// root.cern.ch/root/htmldoc/guides/spectrum/Spectrum.html#fitting
// root.cern.ch/doc/master/FitAwmi_8C.html

// Caution: in the fitted model all the peaks have the same sigma.

// This macro fits the source spectrum using the AWMI algorithm
// from the "TSpectrumFit" class ("TSpectrum" class is used to find peaks).


#include "TROOT.h"
#include "TMath.h"
#include "TRandom.h"
#include "TH1.h"
#include "TF1.h"
#include "TCanvas.h"
#include "TSpectrum.h"
#include "TSpectrumFit.h"
#include "TPolyMarker.h"
#include "TList.h"

#include <iostream>
using namespace std;


TH1F* FitAwmi_Create_Spectrum () 
{
  // This function creates well separated peaks with randomized means and areas
  // Note: TSpectrumFit assumes that all the peaks have the same sigma (!)

  Int_t NBins = 1000;
  Double_t xmin = -10., xmax = 10., xpos = xmin;

  TH1F* hspectrum = new TH1F ("hspectrum", "Simulated spectrum", NBins, xmin, xmax);
  hspectrum->SetStats (kFALSE);

  TF1 fgauss ("fgauss", "TMath::Gaus (x,  [0] ,  [1]  , 1 )", xmin, xmax);
  fgauss.SetParNames                 (  "mean", "sigma");

  gRandom->SetSeed(0);                // Make it really random
  Double_t sigma = hspectrum->GetBinWidth (1) * gRandom->Uniform (2., 6.);
  fgauss.SetParameter (1, sigma);

  Int_t NPeaks = 0;

  for (double xpos = xmin + 3.*sigma;  xpos + 3.*sigma < xmax ;  xpos += 6.*sigma)
  {
    NPeaks++;
    fgauss.SetParameter (0, xpos);
    Double_t Area = gRandom->Uniform (1., 11.);
    hspectrum->Add (&fgauss, Area, "");     // "" ... or ... "I"

    cout << "Created peak at " << xpos << " with [Amp, Area]:  ["
         << Area / sigma / TMath::Sqrt(TMath::TwoPi()) << " , "
         << Area << "]\n";
  }
  cout << "\nTotal number of created peaks = " << NPeaks
       << " with sigma = " << sigma << "\n\n";

  return hspectrum;
}




void FitAwmi () 
{
  /*****     Stage 1: creating spectrum   *****/

  TH1F* hspec = FitAwmi_Create_Spectrum ();

  TCanvas* cFit = new TCanvas ("cFit", "cFit", 10, 10, 1000, 700);
  hspec->Draw ("L");
  cFit->Modified();
  cFit->Update();

  cout << "* Press Enter to [1] search for peak positions and (later) [2] fit.";
  cin.ignore();


  /*****     Stage 2:  peak search        *****/

  Int_t i, NFound, binnum;
  Int_t NBins = hspec->GetNbinsX ();
  Double_t* source = new Double_t [NBins];
  Double_t* destin = new Double_t [NBins];

  for (i = 0; i < NBins; i++)  source[i] = hspec->GetBinContent (i+1);

  TSpectrum* Analyser1dim = new TSpectrum (); // note: default maxpositions = 100

  // Searching for Candidate Peaks' positions

  NFound = Analyser1dim->SearchHighRes (source, destin, NBins, 
                                        2., 2., kFALSE, 10000, kFALSE, 0);
/*
  Int_t TSpectrum::SearchHighRes (
     Double_t* source,           : pointer to the vector of source spectrum
     Double_t* destVector,       : pointer to the vector of resulting deconvolved spectrum
     Int_t     ssize,            : length of source spectrum
     Double_t  sigma,            : sigma of searched peaks (see manual)
     Double_t  threshold,        : threshold value in % for selected peaks. Peaks with
                                   amplitude < threshold*highest_peak/100 are ignored, see manual
     bool      backgroundRemove, : logical variable, set if the removal of background 
                                   before deconvolution is desired
     Int_t     deconIterations,  : number of iterations in deconvolution operation
     bool      markov,           : true = first the source spectrum is replaced by new spectrum
                                   calculated using Markov chains method
     Int_t     averWindow        : (Markov: ) averaging window of searched peaks, (see manual)
  )

  The goal of this function is to identify automatically the peaks in spectrum with the
  presence of the continuous background and statistical fluctuations - noise.
  
  1. Background is removed (if desired), 
  2. Markov smoothed spectrum is calculated (if desired), 
  3. Response function is generated according to given sigma + Deconvolution is carried out.

  See root.cern.ch/doc/v608/classTSpectrum.html#a5ba181a45b117934b48c4ef5f78d0b2b
*/

  cout << "\n* SearchHighRes has found: " << NFound << " peaks.\n\n";


  /*****     Stage 3: preparation for fitting of peaks    *****/

  // Filling in the initial estimates of the input parameters

  Double_t* Amplit = new Double_t [NFound];
  Double_t* Posit  = Analyser1dim->GetPositionX (); // returns X positions of found peaks

  Bool_t* FixPosit  = new Bool_t [NFound];          // should these parameters be fixed? [T/F]
  Bool_t* FixAmplit = new Bool_t [NFound];

  for (i = 0; i < NFound; i++) 
  {
    FixAmplit[i] = kFALSE;
    FixPosit [i] = kFALSE;

    binnum = 1 + Int_t (Posit[i] + 0.5);            // the "nearest" bin
    Amplit   [i] = hspec->GetBinContent (binnum);   // estimating the peak amplitude
  }

  cout << "* Press Enter to fit the spectra using the AWMI approach.";
  cin.ignore();


  // Creation of the fitting tool

  TSpectrumFit* Fitter1dim = new TSpectrumFit (NFound);

/*
   Class for fitting 1D spectra using AWMI (algorithm without matrix inversion)
   and conjugate gradient algorithms for symmetrical matrices
   (Stiefel-Hestens method). AWMI method allows to fit simultaneously 100s up
   to 1000s peaks. Stiefel method is very stable, it converges faster,
   but is more time consuming.

   See: root.cern.ch/doc/v608/classTSpectrumFit.html
*/

  Fitter1dim->SetFitParameters (0, NBins-1 , 1000, 0.1, Fitter1dim->kFitOptimChiCounts,
                                Fitter1dim->kFitAlphaHalving, Fitter1dim->kFitPower2,
                                Fitter1dim->kFitTaylorOrderFirst );
/* 
  void TSpectrumFit::SetFitParameters (
      Int_t    xmin,             : Fitting region
      Int_t    xmax,             : Fitting region
      Int_t    numberIterations, : # of desired iterations in the fit
      Double_t alpha,            : Convergence coeff. = (0,1], for details see references
      Int_t    statisticType,    : Type of statistics, possible values:
                       kFitOptimChiCounts     = chi2 stats with counts as weight coefs.
                       kFitOptimChiFuncValues = chi2 stats with function values as weight coefs.
                       kFitOptimMaxLikelihood
      Int_t    alphaOptim,       : Optimization of convergence algorithm, possible values:
                                                        kFitAlphaHalving, kFitAlphaOptimal
      Int_t    power,            : Possible values kFitPower2,4,6,8,10,12, see Refs. 
                                   Only for Awmi fitting function.
      Int_t    fitTaylor         : Order of Taylor expansion. Only for Awmi fitting function.
                                   Possible values: kFitTaylorOrderFirst, kFitTaylorOrderSecond
   )
*/

  Fitter1dim->SetPeakParameters (2., kFALSE, Posit, FixPosit, Amplit, FixAmplit);

/* 
  void TSpectrumFit::SetPeakParameters	(
      Double_t        sigma,        : initial value of sigma parameter
      Bool_t          fixSigma,     : true if sigma should be fixed
      const Double_t* positionInit, : array of initial values of peaks positions
      const Bool_t*   fixPosition,  : array for peaks. True if a given position should be fixed
      const Double_t* ampInit,      : array of initial values of peaks amplitudes
      const Bool_t*   fixAmp        : array for amplitudes. True if a given amp should be fixed
  )
  See: 
*/

  // Fitter1dim->SetBackgroundParameters (source[0], kFALSE, 0., kFALSE, 0., kFALSE);


  /*****     Stage 4: do the fitting                 *****/


  Fitter1dim->FitAwmi (source);


  /*****     Stage 5: extracting found parameters    *****/

  Double_t* Positions        = Fitter1dim-> GetPositions ();       // Caution: Xscale is Bin No.
  Double_t* PositionsErrors  = Fitter1dim-> GetPositionsErrors (); // Caution: Xscale is Bin No.
  Double_t* Amplitudes       = Fitter1dim-> GetAmplitudes ();
  Double_t* AmplitudesErrors = Fitter1dim-> GetAmplitudesErrors ();
  Double_t* Areas            = Fitter1dim-> GetAreas ();           // Caution: Xscale is Bin No.
  Double_t* AreasErrors      = Fitter1dim-> GetAreasErrors ();     // Caution: Xscale is Bin No.

  Double_t x1 = hspec->GetBinCenter(1), dx = hspec->GetBinWidth (1);

  Double_t sigma, sigmaErr;
  Fitter1dim->GetSigma (sigma, sigmaErr);                          // Caution: Xscale is Bin No.

  // Current TSpectrumFit needs a sqrt(2) correction factor for sigma

  sigma /= TMath::Sqrt2(); sigmaErr /= TMath::Sqrt2();

  // Convert "bin numbers" into "x-axis values"

  sigma *= dx; sigmaErr *= dx;

  cout << "\n * Overall sigma found by fit: " << sigma << " (+-" << sigmaErr << ")" << endl;
  cout << " * Fit chi^2 = " << Fitter1dim->GetChi () << "\n\n";

  for (i = 0; i < NFound; i++)
  {
    // Convert "bin numbers" into "x-axis values"
    Positions      [i] = x1 + Positions[i] * dx;
    PositionsErrors[i] *= dx;
    Areas          [i] *= dx;
    AreasErrors    [i] *= dx;

    cout << "Found: "
         << Positions [i] << " (+-" << PositionsErrors [i] << ") "
         << Amplitudes[i] << " (+-" << AmplitudesErrors[i] << ") "
         << Areas     [i] << " (+-" << AreasErrors     [i] << ")" << endl;
  }

  hspec->Draw ("SAME HIST");

  TPolyMarker *pm = new TPolyMarker (NFound, Positions, Amplitudes);
  hspec->GetListOfFunctions()->Add (pm);
  pm->SetMarkerStyle (23);
  pm->SetMarkerColor (kRed);
  pm->SetMarkerSize (1);
}