#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooUniform.h"
#include "RooDataSet.h"
#include "RooDataHist.h"
#include "RooHistFunc.h"
#include "RooRealSumPdf.h"
#include "RooParamHistFunc.h"
#include "RooHistConstraint.h"
#include "RooProdPdf.h"
#include "RooMinuit.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TStyle.h"
#include "TPaveText.h"
#include "RooStats/HybridCalculatorOriginal.h"
#include "RooStats/HybridResult.h"
#include "RooStats/HybridPlot.h"

#include <sys/stat.h>

using namespace RooFit;
using namespace RooStats;

void DummyTest() {

  // Plotting style
  gStyle->SetTitleSize(0.06, "xy");
  gStyle->SetLabelSize(0.04, "xy");
  gStyle->SetTitleOffset(0.8, "x");
  gStyle->SetTitleOffset(1., "y");


  // Store the sum of squares of weights also
  // Error per bin will be computed as sqrt(sum of squares of weights); otherwise sqrt(bin content)
  // histo->GetBinError(bin-number) to get bin error
  TH1::SetDefaultSumw2(kTRUE);

  // Reading data
  TFile * StandardFile = new TFile("DummyData.root");
  RooWorkspace * w = (RooWorkspace*) StandardFile->Get("w");
  TH1D * Histo_Bin_Dummy2_th1 = (TH1D*)w->genobj("Histo_Bin_Dummy2");

  // Declare observable x
  RooRealVar x("x", "x", 0.0, 1.0);
  x.setBins(50);
  RooArgList observables(x); // variables to be generated, used much later for model

  // Define signal range in which events counts are to be defined
  x.setRange("sideband", 0.15, 0.5);
  x.setRange("extrapolate", 0.55, 1.0);
  x.setRange("fullRange", 0.0, 1.0);

  // Create Gaussian PDF for sideband fit g(x, mean, sigma)
  RooRealVar mean_left("#mu_{left}", "mean of gaussian", 0.35, 0.0, 1.0);
  RooRealVar sigma_left("#sigma_{left}", "width of gaussian", 0.5, 0, 10);
  RooGaussian gauss_left("gauss_left", "gaussian", x, mean_left, sigma_left);

  // Get data
  RooDataHist Histo_Bin_Dummy2("Histo_Bin_Dummy2", "Histo_Bin_Dummy2", x, Import(*Histo_Bin_Dummy2_th1));

  // Associated #N as expected number of events
  RooRealVar n_left("N_{left}", "number of left events", 100000, 0., 1e+08);
  RooExtendPdf ex_pdf_left("ex_pdf_left", "extended left p.d.f", gauss_left, n_left, "fullRange");

  // Fit dummy data to left gaussian/pdf
  RooFitResult *fit_result_left = ex_pdf_left.fitTo(Histo_Bin_Dummy2, Extended(kTRUE), Range("sideband"),
          Save(), PrintLevel(-1));

  // Now get the number of events in the signal region
  RooArgList pars(*ex_pdf_left.getParameters(RooArgSet(x)));
  float n_left_fit = ((RooRealVar *) pars.find("N_{left}"))->getVal();

  // Extract gaussian function (with fit parameters) as a TF1 function
  RooArgSet prodSet(ex_pdf_left);
  RooProduct unNormPdf("fitted Function", "fitted Function", prodSet);
  TF1 *f_left_gauss = unNormPdf.asTF(RooArgList(x), pars);

  // Extract the extrapolated left in the Signal region
  // and save it in a histo
  Double_t x_bin_width = 1.0 / 50.0;
  TH1F *hbkg_ex = new TH1F("hbkg_ex", "ROOT fit", 50, 0.0, 1.0);
  // For outside the signal region, take data from histogram
  // For the signal region compute the integral of the leakage
  Double_t integ_left_leakage_bin;
  Double_t integ_left_leakage_bin_error;
  Double_t integ_left_full = f_left_gauss->Integral(0.0, 1.0);
  for (int i = 1; i < hbkg_ex->GetNbinsX(); i++) {
    if (hbkg_ex->GetXaxis()->GetBinCenter(i) < 0.55) {
      hbkg_ex->SetBinContent(i, Histo_Bin_Dummy2_th1->GetBinContent(i));
      hbkg_ex->SetBinError(i, Histo_Bin_Dummy2_th1->GetBinError(i));
    } else {
      integ_left_leakage_bin = n_left_fit * f_left_gauss->Integral(hbkg_ex->GetXaxis()->GetBinLowEdge(i), hbkg_ex->GetXaxis()->GetBinUpEdge(i), 0) / integ_left_full;
      integ_left_leakage_bin_error = n_left_fit * f_left_gauss->IntegralError(hbkg_ex->GetXaxis()->GetBinLowEdge(i), hbkg_ex->GetXaxis()->GetBinUpEdge(i), 0, fit_result_left->covarianceMatrix().GetMatrixArray()) / integ_left_full;
      hbkg_ex->SetBinContent(i, integ_left_leakage_bin);
      hbkg_ex->SetBinError(i, integ_left_leakage_bin_error);
    }
  }
  // Coloring extrapolated bkg red
  hbkg_ex->SetFillColor(kRed);
  hbkg_ex->SetMarkerColor(kRed);
  hbkg_ex->SetLineColor(kRed);

  // Gauss plot
  TCanvas *c01 = new TCanvas("c01", "c01", 80, 0, 1440, 1200);
  RooPlot *x_frame_left = x.frame(Title("RooFit fit"));
  Histo_Bin_Dummy2.plotOn(x_frame_left);
  ex_pdf_left.plotOn(x_frame_left);
  ex_pdf_left.paramOn(x_frame_left, Layout(0.55));
  x_frame_left->SetTitle("");
  x_frame_left->SetMinimum(1e-1);
  x_frame_left->SetMaximum(1e6);
  x_frame_left->Draw();
  hbkg_ex->Draw("Esame");
  c01->SetLogy();
  c01->SetGridx();
  c01->SetGridy();
  c01->SetLeftMargin(0.13);
  c01->SetBottomMargin(0.12);
  c01->SaveAs("test1.png");

  // Use the extrapolate Bkgr shape to subtract the data in the Signal region
  TH1D *hdiff = (TH1D *) Histo_Bin_Dummy2_th1->Clone("hdiff");
  hdiff->Add(hbkg_ex, -1.0);
  // Set negative bin contents to zero...and bins outside of
  // signal region also to zero
  for (int i = 1; i < hdiff->GetNbinsX(); i++) {
    if (hdiff->GetBinLowEdge(i) < 0.55) {
      hdiff->SetBinContent(i, 0.0000);
      hdiff->SetBinError(i, 0.00001);
    }
    if (hdiff->GetBinContent(i) < 0) {
      hdiff->SetBinContent(i, 0.0000);
      hdiff->SetBinError(i, 0.00001);
    }
  }
  // After subtraction one has to reset the number of signal entries by hand (over-/underflow bin problem)
  // underflow bin: 0; overflow bin: N + 1
  // Integral with error gives the direct number of NR leftover entries with error
  Double_t n_signal_leftover_error;
  Double_t n_signal_leftover = hdiff->IntegralAndError(1, hdiff->GetNbinsX(), n_signal_leftover_error);
  hdiff->SetEntries(n_signal_leftover);

  // Signal leftover plot
  TCanvas *c02 = new TCanvas("c02", "c02", 80, 0, 1440, 1200);
  c02->cd(2);
  hdiff->SetTitle("");
  hdiff->SetName("Data minus left");
  // hdiff->SetMinimum(0.0);
  hdiff->GetXaxis()->SetTitle("x");
  hdiff->SetTitleSize(0.06, "xy");
  hdiff->SetLabelSize(0.04, "xy");
  hdiff->SetTitleOffset(0.8, "x");
  hdiff->SetTitleOffset(0.9, "y");
  //hdiff->SetStats();
  hdiff->Draw();
  c02->SetGridx();
  c02->SetGridy();
  c02->SetLeftMargin(0.13);
  c02->SetBottomMargin(0.12);
  c02->SaveAs("test2.png");

  // Perform joint fit of the Bkg + signal with both templates
  // Use the bkgr extrapolation and Signal shape (diff) as histogram templates
  // to fit the Signal region
  RooDataHist hbkgX("hbkgX", "hbkgX", x, Import(*hbkg_ex));
  RooDataHist hsigX("hsigX", "hsigX", x, Import(*hdiff));

  // Construct parameterized histogram shapes for signal and background
  // this is only used to propagate the bin fluctuations stat error
  RooParamHistFunc pdf_parametrized_histo_sig("pdf_parametrized_histo_sig", "pdf_parametrized_histo_sig", hsigX);
  RooParamHistFunc pdf_parmetrized_histo_bkg("pdf_parmetrized_histo_bkg", "pdf_parmetrized_histo_bkg", hbkgX);
  // Initialise amplitudes of the sig/bkg pdfs respectively
  RooRealVar Amplitude_sig("Amplitude_sig", "Amplitude_sig", 1, 0.01, 5000);
  RooRealVar Amplitude_bkg("Amplitude_bkg", "Amplitude_bkg", 1, 0.01, 5000);

  // Construct the sum of these as unconstrained model
  RooRealSumPdf model_unconstrained("summed_pdf_parametrized_histo", "summed_pdf_parametrized_histo",
          RooArgList(pdf_parmetrized_histo_bkg, pdf_parametrized_histo_sig),
          RooArgList(Amplitude_bkg, Amplitude_sig),
          true);

  // Initialise constraints
  RooHistConstraint histo_constraint_sig("histo_constraint_sig", "histo_constraint_sig", pdf_parametrized_histo_sig);
  RooHistConstraint histo_constraint_bkg("histo_constraint_bkg", "histo_constraint_bkg", pdf_parmetrized_histo_bkg);

  // Add constraints from templates
  RooProdPdf model("model", "model", RooArgSet(histo_constraint_bkg, histo_constraint_sig),
          Conditional(model_unconstrained, x));

  // Import data again and fit model to it in order to constrain the respective amplitudes of the sig/bkg pdfs
  RooFitResult *fit_result_signal = model.fitTo(Histo_Bin_Dummy2, Range("fullRange"), Save(), PrintLevel(-1));

  // Model template plot
  TCanvas *c03 = new TCanvas("c03", "c03", 80, 0, 1440, 1200);
  RooPlot *xframe02 = x.frame(Title("RooFit fit"));
  Histo_Bin_Dummy2.plotOn(xframe02, Name("Data"));
  model.plotOn(xframe02, VisualizeError(*fit_result_signal, 1), Name("Model with errors"));  // CASE 2
  // model.plotOn(xframe02, Name("Model with errors"));  // CASE 1
  Histo_Bin_Dummy2.plotOn(xframe02);
  model.plotOn(xframe02, Name("left gauss"), Components(pdf_parmetrized_histo_bkg), LineColor(kRed));
  model.plotOn(xframe02, Name("signal template"), Components(pdf_parametrized_histo_sig), LineColor(kGreen));
  xframe02->SetTitle("");
  xframe02->SetMinimum(1e-1);
  xframe02->SetMaximum(1e6);
  xframe02->Draw();
  c03->SetLogy();
  c03->SetGridx();
  c03->SetGridy();
  c03->Modified();
  c03->Update();
  TLegend *leg03 = new TLegend(0.50, 0.64, 0.93, 0.93);
  leg03->AddEntry("Data", "Data", "LEP");
  leg03->AddEntry("Left gauss", "Left gaussian template", "F");
  leg03->AddEntry("Signal template", "Signal leftover template", "F");
  leg03->AddEntry("Model with errors", "Model with errors", "F");
  leg03->SetTextSize(0.035);
  leg03->Draw();
  c03->SetLeftMargin(0.13);
  c03->SaveAs("test3.png");
}