#include <iostream>

#include "RooWorkspace.h"
#include "RooRealVar.h"
#include "RooPoisson.h"
#include "RooDataSet.h"
#include "RooProdPdf.h"
#include "RooFitResult.h"
#include "RooMsgService.h"

void testWSsimpleWithCons(bool useCons = false, bool rmCons = false, bool detail = false) {
  
  RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL);
  
  RooWorkspace *w = new RooWorkspace("w","w");
  w->factory("Gaussian::cons(nui[0,-5,5],glo[0,-5,5],1.)");
  w->var("glo")->setConstant();
  w->factory("prod::ns(s[35.,0.,50.],sum::expns(1,prod::rds(nui,inc[0.05])))");
  w->factory("Poisson::pb(nobs[115.,0,300],sum::sb(ns,b[100.]))");
  w->factory("PROD::model(pb,cons)");
  RooRealVar *s = w->var("s");
  RooRealVar *n = w->var("nui");
  RooPoisson *pb  = (RooPoisson *)w->pdf("pb");
  RooProdPdf *pdf = (RooProdPdf *)w->pdf("model");
  pb->setNoRounding(true);
  RooRealVar *obs = w->var("nobs");
  RooDataSet *da = new RooDataSet("da","da",*obs);
  da->add(*obs);
  if (detail)
    da->Print("v");
  
  double minnll[3] = { 0, 0, 0 };
  for (int i = 0; i < 3; i++) {
    s->setVal(35.);
    s->setError(0);
    n->setVal(0);
    n->setError(0);
    RooFitResult *fitR(nullptr);
    if (rmCons)
      w->removeSet("CACHE_CONSTR_OF_PDF_model_FOR_OBS_nobs");
    if (useCons)
      fitR = pdf->fitTo(*da,RooFit::PrintLevel(-1),RooFit::Save(),RooFit::Minos(*s),RooFit::Constrain(*n));
    else
      fitR = pdf->fitTo(*da,RooFit::PrintLevel(-1),RooFit::Save(),RooFit::Minos(*s));
    if (detail)
      fitR->Print("v");
    minnll[i] = fitR->minNll();
    delete fitR;
  }
  std::cout << "With useConstrain = " << useCons << ", the minNLL for the three fits : " << minnll[0] << " " << minnll[1] << " " << minnll[2] << std::endl;
  
}

void go() {
  std::cout << "Without using Constrain(*n) option in fitTo" << std::endl;
  testWSsimpleWithCons();
  std::cout << "Using Constrain(*n) option in fitTo" << std::endl;
  testWSsimpleWithCons(true);
  std::cout << "Without using Constrain(*n) option in fitTo but removing the cached set" << std::endl;
  testWSsimpleWithCons(false,true);
}