/// \file
/// \ingroup tutorial_roofit_main
/// \notebook
/// Special pdf's: using a pdf defined by a sum of real-valued amplitude components
///
/// \macro_image
/// \macro_code
/// \macro_output
///
/// \date July 2008
/// \author Wouter Verkerke

#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooFormulaVar.h"
#include "RooRealSumPdf.h"
#include "RooPolyVar.h"
#include "RooProduct.h"
#include "TH1.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "RooPlot.h"
using namespace RooFit;

void rf704_amplitudefit()
{
   // S e t u p   2 D   a m p l i t u d e   f u n c t i o n s
   // -------------------------------------------------------

   // Observables
   RooRealVar t("t", "time", -1., 15.);
   RooRealVar cosa("cosa", "cos(alpha)", -1., 1.);

   RooRealVar tau("tau", "#tau", 1.5);
   RooRealVar deltaGamma("deltaGamma", "deltaGamma", 0.3);
   RooFormulaVar coshG("coshGBasis", "exp(-@0/ @1)*cosh(@0*@2/2)", {t, tau, deltaGamma});
   RooFormulaVar sinhG("sinhGBasis", "exp(-@0/ @1)*sinh(@0*@2/2)", {t, tau, deltaGamma});

   // Construct polynomial amplitudes in cos(a)
   RooPolyVar poly1("poly1", "poly1", cosa, RooArgList{0.5, 0.2, 0.2}, 0);
   RooPolyVar poly2("poly2", "poly2", cosa, RooArgList{1.0, -0.2, 3.0}, 0);

   // Construct 2D amplitude as uncorrelated product of amp(t)*amp(cosa)
   RooProduct ampl1("ampl1", "amplitude 1", {poly1, coshG});
   RooProduct ampl2("ampl2", "amplitude 2", {poly2, sinhG});

   // C o n s t r u c t   a m p l i t u d e   s u m   p d f
   // -----------------------------------------------------

   // Amplitude strengths
   RooRealVar f1("f1", "f1", 1, 0, 2);
   RooRealVar f2("f2", "f2", 0.5, 0, 2);

   // Construct pdf
   RooRealSumPdf pdf("pdf", "pdf", RooArgList(ampl1, ampl2), RooArgList(f1, f2));

   // Generate some toy data from pdf
   std::unique_ptr<RooDataSet> data{pdf.generate({t, cosa}, 10000)};

   // Fit pdf to toy data with only amplitude strength floating
   pdf.fitTo(*data, PrintLevel(-1));

   // P l o t   a m p l i t u d e   s u m   p d f
   // -------------------------------------------

   // Make 2D plots of amplitudes
   TH1 *hh_cos = ampl1.createHistogram("hh_cos", t, Binning(50), YVar(cosa, Binning(50)));
   TH1 *hh_sin = ampl2.createHistogram("hh_sin", t, Binning(50), YVar(cosa, Binning(50)));
   hh_cos->SetLineColor(kBlue);
   hh_sin->SetLineColor(kRed);

   // Make projection on t, plot data, pdf and its components
   // Note component projections may be larger than sum because amplitudes can be negative
   RooPlot *frame1 = t.frame();
   data->plotOn(frame1);
   pdf.plotOn(frame1);
   pdf.plotOn(frame1, Components(ampl1), LineStyle(kDashed));
   pdf.plotOn(frame1, Components(ampl2), LineStyle(kDashed), LineColor(kRed));

   // Make projection on cosa, plot data, pdf and its components
   // Note that components projection may be larger than sum because amplitudes can be negative
   RooPlot *frame2 = cosa.frame();
   data->plotOn(frame2);
   pdf.plotOn(frame2);
   pdf.plotOn(frame2, Components(ampl1), LineStyle(kDashed));
   pdf.plotOn(frame2, Components(ampl2), LineStyle(kDashed), LineColor(kRed));

   TCanvas *c = new TCanvas("rf704_amplitudefit", "rf704_amplitudefit", 800, 800);
   c->Divide(2, 2);
   c->cd(1);
   gPad->SetLeftMargin(0.15);
   frame1->GetYaxis()->SetTitleOffset(1.4);
   frame1->Draw();
   c->cd(2);
   gPad->SetLeftMargin(0.15);
   frame2->GetYaxis()->SetTitleOffset(1.4);
   frame2->Draw();
   c->cd(3);
   gPad->SetLeftMargin(0.20);
   hh_cos->GetZaxis()->SetTitleOffset(2.3);
   hh_cos->Draw("surf");
   c->cd(4);
   gPad->SetLeftMargin(0.20);
   hh_sin->GetZaxis()->SetTitleOffset(2.3);
   hh_sin->Draw("surf");
}
