#include <RooRealVar.h>
#include <RooPolynomial.h>
#include <RooPlot.h>
#include <RooDataSet.h>
#include <RooGlobalFunc.h>
#include <RooFitResult.h>
#include <RooMsgService.h>

#include <TCanvas.h>
#include <TLegend.h>

void rf612_recoverFromInvalidParameters() {

  // Create a fit model:
  // The polynomial is notoriously unstable, because it can quickly go negative.
  // Since PDFs need to be positive, one often ends up with an unstable fit model.
  RooRealVar x("x", "x", -15, 15);
  RooRealVar a1("a1", "a1", -0.5, -10., 20.);
  RooRealVar a2("a2", "a2", 0.2, -10., 20.);
  RooRealVar a3("a3", "a3", 0.01);
  RooPolynomial pdf("pol3", "c + a1 * x + a2 * x*x + 0.01 * x*x*x", x, RooArgSet(a1, a2, a3));

  // Create toy data with all-positive coefficients:
  std::unique_ptr<RooDataSet> data(pdf.generate(x, 10000));

  // For plotting.
  // We create pointers to the plotted objects. We want these objects to leak out of the function,
  // so we can still see them after it returns.
  TCanvas* c = new TCanvas();
  RooPlot* frame = x.frame();
  data->plotOn(frame, RooFit::Name("data"));

  // Plotting a PDF with disallowed parameters doesn't work. We would get a lot of error messages.
  // Therefore, we disable plotting messages in RooFit's message streams:
  RooMsgService::instance().getStream(0).removeTopic(RooFit::Plotting);
  RooMsgService::instance().getStream(1).removeTopic(RooFit::Plotting);


  // RooFit before ROOT 6.24
  // --------------------------------
  // Before 6.24, RooFit wasn't able to recover from invalid parameters. The minimiser just errs around
  // the starting values of the parameters without finding any improvement.

  // Set up the parameters such that the PDF would come out negative. The PDF is now undefined.
  a1.setVal(10.);
  a2.setVal(-1.);

  // Perform a fit:
  std::unique_ptr<RooFitResult> fitWithoutRecovery{pdf.fitTo(*data, RooFit::Save(),
      RooFit::RecoverFromUndefinedRegions(0.), // This is how RooFit behaved prior to ROOT 6.24
      RooFit::PrintEvalErrors(-1), // We are expecting a lot of evaluation errors. -1 switches off printing.
      RooFit::PrintLevel(-1))};

  pdf.plotOn(frame, RooFit::LineColor(kRed), RooFit::Name("noRecovery"));



  // RooFit since ROOT 6.24
  // --------------------------------
  // The minimiser gets information about the "badness" of the violation of the function definition. It uses this
  // to find its way out of the disallowed parameter regions.
  std::cout << "\n\n\n-------------- Starting second fit ---------------\n\n" << std::endl;

  // Reset the parameters such that the PDF is again undefined.
  a1.setVal(10.);
  a2.setVal(-1.);

  // Fit again, but pass recovery information to the minimiser:
  std::unique_ptr<RooFitResult> fitWithRecovery{pdf.fitTo(*data, RooFit::Save(),
      RooFit::RecoverFromUndefinedRegions(1.), // The magnitude of the recovery information can be chosen here.
                                               // Higher values mean more aggressive recovery.
      RooFit::PrintEvalErrors(-1), // We are still expecting a few evaluation errors.
      RooFit::PrintLevel(0))};

  pdf.plotOn(frame, RooFit::LineColor(kBlue), RooFit::Name("recovery"));



  // Collect results and plot.
  // --------------------------------
  // We print the two fit results, and plot the fitted curves.
  // The curve of the fit without recovery cannot be plotted, because the PDF is undefined if a2 < 0.
  fitWithoutRecovery->Print();
  std::cout << "Without recovery, the fitter encountered " << fitWithoutRecovery->numInvalidNLL()
      << " invalid function values. The parameters are unchanged." << std::endl;

  fitWithRecovery->Print();
  std::cout << "With recovery, the fitter encountered " << fitWithRecovery->numInvalidNLL()
      << " invalid function values, but the parameters are fitted." << std::endl;

  TLegend* legend = new TLegend(0.5, 0.7, 0.9, 0.9);
  legend->SetBorderSize(0);
  legend->SetFillStyle(0);
  legend->AddEntry("data", "Data", "P");
  legend->AddEntry("noRecovery", "Without recovery (cannot be plotted)", "L");
  legend->AddEntry("recovery", "With recovery", "L");
  frame->Draw();
  legend->Draw();
  c->Draw();
}