using namespace RooFit;
using namespace RooStats;

int nScanPoints = 50;

void vaccine_PL(int Iv_o = 5, int Ip_o = 90, int Nv_o = 15000, int Np_o = 15000, double confLevel = 0.683) {
  
  // vaccine efficiency - parameter of interest (poi)
  RooRealVar eps("eps","eps",double(Ip_o-Iv_o)/double(Ip_o),0,1);
  // incidence - nuisance parameter (np)
  RooRealVar inc("inc","inc",double(Ip_o)/double(Np_o),0,1);

  // Number of infected after vaccination
  RooRealVar Iv("Iv","Iv",Iv_o,0,Nv_o);
  // Number of infected without vaccination (placebo)
  RooRealVar Ip("Ip","Ip",Ip_o,0,Np_o);

  // Number of persons vaccined
  RooRealVar Nv("Nv","Nv",Nv_o);
  // Number of persons without vaccination (placebo)
  RooRealVar Np("Np","Np",Np_o);


  // build pdf for Iv
  RooGenericPdf IvPdf("IvPdf","Binomial Pdf for vaccine efficiency","TMath::Binomial(Ip,(Ip-Iv))*TMath::Power(eps,(Ip-Iv))*TMath::Power((1-eps),Iv)",RooArgList(eps,Ip,Iv));
  RooAbsReal* intOfIvPdf = IvPdf.createIntegral(Iv);

  // build pdf for Ip
  RooGenericPdf IpPdf("IpPdf","Binomial Pdf for incidence","TMath::Binomial(Np,Ip)*TMath::Power(inc,Ip)*TMath::Power((1-inc),Np-Ip)",RooArgList(inc,Np,Ip));
  RooAbsReal* intOfIpPdf = IpPdf.createIntegral(Ip);

  // plot of IvPdf over eps
  RooPlot* epsframe = eps.frame();
  IvPdf.plotOn(epsframe,RooFit::Name("IvPdf"));
  epsframe->SetTitle(";Vaccine Efficacy;Likelihood");
  epsframe->Draw();
  gPad->SetTicks();

  // let's build the product pdf for both poi and np
  RooProdPdf model("model","Vaccination efficacy model",IvPdf,IpPdf);
  // we have to make sure it is properly normalised
  RooAbsReal* intmodel = model.createIntegral(Ip,Iv);

  // we now create a dataset for the observation we made
  // we need to create a dataset to store our observation
  RooDataSet data("data","data",RooArgSet(Iv,Ip,Nv,Np));
  Iv = Iv_o;
  Ip = Ip_o;
  Nv = Nv_o;
  Np = Np_o;
  data.add(RooArgSet(Iv,Ip,Nv,Np));

  // let's save everything in a workspace
  RooWorkspace *w = new RooWorkspace("w");
  // importing pdfs
  w->import(model);
  // importing data
  w->import(data);


  // we can now create the model config object
  ModelConfig modelConfig("P(Iv|Ip,eps)");
  // let's define workspace, pdfs, poi, np, observable
  modelConfig.SetWorkspace(*w);
  modelConfig.SetPdf("model");
  modelConfig.SetParametersOfInterest("eps");
  modelConfig.SetNuisanceParameters("inc");
  modelConfig.SetObservables(RooArgSet(Iv,Ip));


  // let's add the ModelConfig to the workspace
  w->import(modelConfig);

  // let's fix the constants
  w->var("Np")->setConstant(1);
  w->var("Nv")->setConstant(1);

  // let's print its content
  w->Print();

  // let's save it to a file
  TFile *fOutput = new TFile("Workspace_vaccine_PL.root","RECREATE");
  w->Write();
  fOutput->Write();
  fOutput->Close();



  ProfileLikelihoodCalculator pl(data,modelConfig);
  pl.SetConfidenceLevel(confLevel);
  LikelihoodInterval *interval = pl.GetInterval();

  RooRealVar *firstPOI = (RooRealVar *)modelConfig.GetParametersOfInterest()->first();
  cout << "\n>>>> RESULT : " << confLevel * 100 << "% interval on " << firstPOI->GetName() << " is : ["
    << interval->LowerLimit(*firstPOI) << ", " << interval->UpperLimit(*firstPOI) << "]\n " << endl;

  LikelihoodIntervalPlot plot(interval);
  plot.SetNPoints(nScanPoints);
  plot.Draw("tf1");  

}