using namespace RooStats;

void UpperLimit()
{  
  // ** Some variable definition
  //
  double nobs = 0;
  double b = 0.03;
  double sigmaB = 0.40;
  double sigmaS = 0.40;
  double scaleS = 1;
  double scaleB = 1;
  int nToys = 1000;

  // ** Workspace definition
  //
  RooWorkspace *w = new RooWorkspace("w");

  w->factory("prod:sig(s[1,0,10],sSyst[1,0,3],scaleS[1])");
  w->factory("prod:bkg(b[1,0,10],bSyst[1,0,3],scaleB[1])");
  w->factory("sum:nexp(sig,bkg)");
  w->factory("Poisson:pdf(nobs[0,0,10],nexp)");  
  // Constraints
  w->factory("Gaussian:constraintB(bObs[1,0,3],bSyst,sigmaB[1])");
  w->factory("Gaussian:constraintS(sObs[1,0,3],sSyst,sigmaS[1])");
  w->factory("PROD:model(pdf,constraintB,constraintS)");
  
  
  // ** Setting parameter values
  //
  w->var("b")->setVal(b);

  w->var("bObs")->setVal(1);
  w->var("bObs")->setConstant(true);
  w->var("sigmaB")->setVal(sigmaB);
  w->var("sigmaB")->setConstant(true);
  
  w->var("sObs")->setVal(1);
  w->var("sObs")->setConstant(true);
  w->var("sigmaS")->setVal(sigmaS);
  w->var("sigmaS")->setConstant(true);
  
  w->var("scaleS")->setVal(scaleS);
  w->var("scaleS")->setConstant(true);
  w->var("scaleB")->setVal(scaleB);
  w->var("scaleB")->setConstant(true);
  
  
  // ** Model definition
  //
  w->defineSet("nuis","b,bSyst,sSyst");
  w->defineSet("glob","bObs,sObs");
  ModelConfig *mConf = new ModelConfig("S+B");
  mConf->SetWorkspace(*w);
  mConf->SetPdf(*w->pdf("model"));
  mConf->SetParametersOfInterest(RooArgSet(*w->var("s"))); // POI
  mConf->SetObservables(RooArgSet(*w->var("nobs"))); // OBS
  mConf->SetNuisanceParameters(*w->set("nuis")); // NUIS
  mConf->SetGlobalObservables(*w->set("glob")); // GLOBAL
  mConf->SetSnapshot(RooArgSet(*w->var("s")));

  w->import(*mConf);
  mConf->Print();

  
  // ** Dataset creation 
  //
  RooDataSet *myData = new RooDataSet("data","obsData", RooArgSet(*w->var("nobs")));
  // Seeting number of observed events
  w->var("nobs")->setVal(nobs);
  myData->add(RooArgSet(*w->var("nobs")));

  w->import(*myData, RooFit::Rename("obsData"));
  w->Print();
  

  // ** Write model in output file
  //
  w->writeToFile("output/CountingModel.root", true);

  
  //
  // ** Frequentist calculator
  //
  RooRandom::randomGenerator()->SetSeed(4357);
  // ** sb-model
  ModelConfig *sbModel = (ModelConfig*)w->obj("S+B");
  RooRealVar *poi = (RooRealVar*)sbModel->GetParametersOfInterest()->first();
  // ** b-model
  ModelConfig *bModel = (ModelConfig*)sbModel->Clone();
  bModel->SetName(sbModel->GetName()+TString("_with_poi_0 (i.e. B-only model)"));
  poi->setVal(0);
  bModel->SetSnapshot(RooArgSet(*poi));

  FrequentistCalculator *fc = new FrequentistCalculator(*myData, *bModel, *sbModel);
  fc->SetToys(nToys,nToys);
  //fc->UseSameAltToys();
  //fc->SetNToysInTails(500,500);
  
  //AsymptoticCalculator  *ac = new AsymptoticCalculator(*myData, *bModel, *sbModel);
  //ac->SetOneSided(true);

  HypoTestInverter calc(*fc);
  //HypoTestInverter calc(*ac);
  calc.SetConfidenceLevel(0.95);
  calc.UseCLs(true);
  calc.SetVerbose(true);
  
  ToyMCSampler* toymcs = (ToyMCSampler*)calc.GetHypoTestCalculator()->GetTestStatSampler();

  ProfileLikelihoodTestStat *profll = new ProfileLikelihoodTestStat(*mConf->GetPdf());
  profll->SetOneSided(true);
  profll->SetPrintLevel(2);
  //profll->SetStrategy(2);
  //profll->SetMinimizer("Minuit2");
  //profll->SetTolerance(1.E-10);
  //profll->SetReuseNLL(false);
  toymcs->SetTestStatistic(profll);

  if (!mConf->GetPdf()->canBeExtended()){
    toymcs->SetNEventsPerToy(1);
    cout << "Can not be extended" << endl;
  }

  int npoints = 30;
  int poimin = 0;
  int poimax = 6;

  calc.SetFixedScan(npoints,poimin,poimax);
  HypoTestInverterResult *r = calc.GetInterval();
  /*
  calc.SetCloseProof(1);
  SamplingDistribution *limDist = calc.GetUpperLimitDistribution(true, 5);
  std::cout << "expected up limit " << limDist->InverseCDF(0.5) << " +/- "
            << limDist->InverseCDF(0.16) << "  "
            << limDist->InverseCDF(0.84) << "\n";
  delete r;
  r = calc.GetInterval();
  */
  double upperLimit = r->UpperLimit();
  double upperLimitError = r->UpperLimitEstimatedError();

  TCanvas *cHT = new TCanvas("HypoTestInverter Scan","HypoTestInverter Scan");
  HypoTestInverterPlot *plotHT = new HypoTestInverterPlot("HTI_Result_Plot","HypoTest Scan Result",r);
  cHT->SetLogy(false);
  plotHT->Draw("CLb 2CL");
  cHT->Draw();

  std::cout << "Expected upper limits, using the B (alternate) model : " << std::endl;
  std::cout << " expected limit (median) " << r->GetExpectedUpperLimit(0)  << std::endl;
  std::cout << " expected limit (-1 sig) " << r->GetExpectedUpperLimit(-1) << std::endl;
  std::cout << " expected limit (+1 sig) " << r->GetExpectedUpperLimit(1)  << std::endl;
  std::cout << " expected limit (-2 sig) " << r->GetExpectedUpperLimit(-2) << std::endl;
  std::cout << " expected limit (+2 sig) " << r->GetExpectedUpperLimit(2)  << std::endl;
  std::cout << " observed " << upperLimit << std::endl;

  //
  // ** Profile likelihood ratio plots
  //
  int nEntries = r->ArraySize();
  TCanvas *cLR = new TCanvas("cLR","cLR",3000,3000);
  if (nEntries > 1){
    int ny = TMath::CeilNint(TMath::Sqrt(nEntries));
    int nx = TMath::CeilNint(double(nEntries) / ny);
    cLR->Divide(nx, ny);
  }
  for (int i = 0; i < nEntries; i++) {
    if (nEntries > 1) cLR->cd(i + 1);
    SamplingDistPlot *pl = plotHT->MakeTestStatPlot(i);
    pl->SetLogYaxis(true);
    pl->Draw();
  }
  cLR->Draw();
  gPad = cLR;

}