#include "RooRealVar.h"
#include "RooDataSet.h"
#include "RooGaussian.h"
#include "RooConstVar.h"
#include "RooPolynomial.h"
#include "RooAddPdf.h"
#include "RooProdPdf.h"
#include "RooFitResult.h"
#include "RooPlot.h"
#include "TCanvas.h"
#include "TAxis.h"
#include "TH1.h"
#include "RooStats/ModelConfig.h"
using namespace RooFit;
using namespace RooStats;

void workspace_build()
{
  RooRealVar mass("mass","mass", 4360., 6360., "MeV");
  RooRealVar mean_sig("mean_sig","mean_sig",5370);
  RooRealVar sigma_sig("sigma_sig","sigma_sig", 20.9);
  RooGaussian signal_model("signal_model","signal_model", mass, mean_sig, sigma_sig);
  RooRealVar n_sig("n_sig", "n_sig", 2000,0,10000);

  RooRealVar exponent("exponent","exponent", -1e-3, -1., 1.);
  RooRealVar n_bkg("n_bkg","n_bkg", 10000, 0, 100000);
  RooExponential bkg_only_model("bkg_only_model", "bkg_only_model", mass, exponent);

  RooAddPdf mass_model("mass_model","mass_model", RooArgList(signal_model, bkg_only_model), RooArgList(n_sig, n_bkg));

  //generate dataset
  RooDataSet* data = mass_model.generate(mass,Extended());
  data->SetName("data");


  //fit
  RooFitResult* r = mass_model.fitTo(*data,Save());
  r->Print("v");

  //Plot
  RooPlot* frame = mass.frame();
  data->plotOn(frame);
  mass_model.plotOn(frame);
  frame->Draw();

  RooWorkspace* w = new RooWorkspace("w", "workspace");
  w->import(*data);
  w->import(mass_model);
  w->writeToFile("workspace.root");
}

void modelconfig_build()
{
  TFile f("workspace.root");
  RooWorkspace* w = (RooWorkspace*)f.Get("w");
  RooAddPdf* mass_model = (RooAddPdf*)w->pdf("mass_model");
  RooRealVar* n_sig = (RooRealVar*)w->var("n_sig");
  RooRealVar* exponent = (RooRealVar*)w->var("exponent");
  RooRealVar* n_bkg = (RooRealVar*)w->var("n_bkg");
  RooRealVar* mass = (RooRealVar*)w->var("mass");

  ModelConfig modelConfig(w);
  modelConfig.SetPdf(*mass_model);
  modelConfig.SetParametersOfInterest(RooArgSet(*n_sig));
  modelConfig.SetNuisanceParameters(RooArgSet(*exponent,*n_bkg));
  modelConfig.SetObservables(*mass);
  modelConfig.SetName("ModelConfig");
  w->import(modelConfig);
  w->writeToFile("model.root");
  }



void example()
{
	workspace_build();
	modelconfig_build();
}

/*
gROOT->ProcessLine(".L $ROOTSYS/tutorials/roostats/StandardHypoTestInvDemo.C+");
StandardHypoTestInvDemo("model.root", "w", "ModelConfig", "", "data", 2, 3, true, 8, 0, 3000,500);
*/