#include "TStopwatch.h"
#include "TCanvas.h"
#include "TRandom3.h"
#include "TROOT.h"
#include "RooPlot.h"
#include "RooAbsPdf.h"
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooGlobalFunc.h"
#include "RooFitResult.h"
#include "RooRandom.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include "RooStats/LikelihoodInterval.h"
#include "RooStats/LikelihoodIntervalPlot.h"
#include "RooStats/BayesianCalculator.h"
#include "RooStats/MCMCCalculator.h"
#include "RooStats/MCMCInterval.h"
#include "RooStats/MCMCIntervalPlot.h"
#include "RooStats/ProposalHelper.h"
#include "RooStats/SimpleInterval.h"
#include "RooStats/FeldmanCousins.h"
#include "TLegend.h"
#include "RooStats/PointSetInterval.h"
#include "TStyle.h"
#include "RooStats/HypoTestResult.h"
#include "RooStats/HypoTestResult.h"
#include "RooStats/SamplingDistribution.h"
#include "TF1.h"
#include "TRandom.h"

using namespace RooFit;
using namespace RooStats;


/*
 * Problem: Have a known background of b=noffObs events and observe nonObs=s+b events. 
 * Test with the profilelikelihoodratio for the null hypothesis s=0 and for the alternative hypothesis s>0 
 * -> One Parameter (s)
 * -> Likelihood function for Null hyptohesis L0=TMath::Poisson(nonObs,b)
 * -> Likelihood function for alternative hypothesis L1=TMath::Poisson(nonObs,b+s) which is maximal at s=noffObs-nonObs.
 *
 * For the concrete example of nonObs=150. and noffObs=100. this means 
 * -> Null hypothesis -log(L0)=-log(TMath::Poisson(150.,100.)) ~ 14.24
 * -> Alternative hypothesis -log(L1)=-log(TMath::Poisson(150.,150)) ~ 3.22
 *
 * The problem is that in the GetHypoTest() function below (implemented in ProfileLikelihoodCalculator.cxx) the following values are calculated:
 * 
 * Variables in ProfileLikelihoodCalculator::GetHypoTest():
 * Double_t NLLatMLE = fFitResult->minNll(); (this variable should hold -log(L1) as it is the maximum likelihood for the best fit of the likelihood function to the data leaving all variables floating)
 * What I get is however ~18.56 instead of the expected 3.22
 *
 * For the variable 
 * Double_t NLLatCondMLE = fit2->minNll(); 
 * holding the -log(L0) for the null hypothesis 
 * in GetHypoTest() I get the value 40.4 instead of the expected value 14.24 for the null hypthesis likelihood
 *
 * What is going on here?
*/
double test(double nonObs=150., double noffObs=100.){
 
  RooWorkspace* wspace = new RooWorkspace("wspace");

  RooRealVar* non = (RooRealVar*) wspace->factory("non[0., 400.]");

  RooRealVar* b=(RooRealVar*) wspace->factory("b[0., 400.]");
  RooRealVar* s=(RooRealVar*) wspace->factory("s[0., 100.]"); 

  // make model
  wspace->factory("Poisson::on(non, sum::splusb(s,b))");

  wspace->factory("PROD::model(on)");

  wspace->defineSet("obs","non");
  wspace->defineSet("poi","s");
  wspace->defineSet("nuis","b");

  non->setVal(nonObs);
  b->setVal(noffObs);

  RooArgSet* ArgSet = new RooArgSet("args");   
  ArgSet->add(*non);
  ArgSet->add(*b);

  RooDataSet *data = new RooDataSet("modelData", "modelData", *ArgSet); 
  data->add(*non);
  data->add(*b);

  wspace->import(*data);

  wspace->var("non")->setVal(nonObs);
  wspace->var("non")->setConstant(true);

  wspace->var("b")->setVal(noffObs);
  wspace->var("b")->setConstant(true); 

  ModelConfig* modelConfig = new ModelConfig("FourBins");
  modelConfig->SetWorkspace(*wspace);
  modelConfig->SetPdf(*wspace->pdf("model"));
  modelConfig->SetParametersOfInterest(*wspace->set("poi"));
  modelConfig->SetNuisanceParameters(*wspace->set("nuis"));
  wspace->import(*modelConfig);

  // use ProfileLikelihood
  ProfileLikelihoodCalculator plc(*data, *modelConfig); 

  RooRealVar *xs=wspace->var("s");

  RooArgSet *nullpar = new RooArgSet("nullParams");
  nullpar->addClone(*xs);
  
  nullpar->setRealValue("s",0.);
  plc.SetNullParameters(*nullpar);

  HypoTestResult *htr = plc.GetHypoTest();
  
  double pval0=htr->NullPValue(); 

  std::cout<<"nullpval="<<pval0<<" sig="<<PValueToSignificance(pval0)<<std::endl;

  return PValueToSignificance(pval0);

}

/*
 *Problem: Have known number of background events (for example b=22) and dice a number of measured events (nobs) 
  from a Poisson distribution of mean b. 
  A hypothesis test is used to test for 
  1) Null hypothesis, i.e. nobs compatible with background b
  2) Alternative hypothesis, i.e. nobs=s+b with signal events s>0
  
  -> The likelihood function for this problem is 
   
      L=TMath::Poisson(nobs,b+s)

  -> For the null hypothesis (s=0)
 
      L0=TMath::Poisson(nobs,b)
   
  -> For the alternative hypothesis (maximizing the likelihood function with s=nobs-b)
  
      L1=TMath::Poisson(nobs,nobs)
 
   
  -> The log-likelihood ratio 

      ll=sqrt(-2.*log(L0/L1))=sqrt(-2.*(log(L0) - log(L1)) )
     
     is distributed like the absolute value of a standard normal distributed variable if the null hypothesis is true. 
     Thus, the the significance of the excess s=nobs-b for s>0 can be defined as ll and should be standard normal distributed 
     if the null hypothesis is true.



 On the other side, the significance can also be calculated using the 
 GetHypotest() function of the ProfileLikelihoodCalculator in RooStat.
 This significance obtained from the ProfileLikelihoodCalculator should also be standard normal distributed. 
 However, a test shows that this is not the case but the width of the significance distribution coming out of the 
 ProfileLikelihoodCalculator routine is 1.38+-0.03 which is significantly different from 1. The latter problem does not 
 occur if the the significance distribution is calculated with ll defined above. 

 I can see two possible reasons for this behaviour:
  1) I used the ProfileLikelihoodCalculator in a wrong way 
  2) The GetHypothesis() function of the ProfileLikelihoodCalculator routine is not giving back the right 
     significance. A possible reason for this is discussed above. 
  
 It would be very helpful for me if some expert could try to comment on this problem...
*/
void distribution(int times=1000){

TH1D *sigRoo=new TH1D("RooStat","",1000,-5.,5.);
TH1D *mySig=new TH1D("mySig","",1000,-5.,5.);

TRandom3 *randi=new TRandom3();
randi->SetSeed(0);

double b,nobs;
double sr,ms;
double temp,sign;

b=22.;

for(int i=0;i<times;i++){

nobs=randi->Poisson(b);
sign=1.;

if( (nobs-b)<=0.) {
 i--;
 continue;
}

ms=sqrt(-2.*(log( TMath::Poisson(nobs,b)) - log( TMath::Poisson(nobs,nobs))));
sr=test(nobs,b);

sigRoo->Fill(sr);
mySig->Fill(ms);
}


cout<<"RooStats sig fit:"<<endl;
sigRoo->Fit("gaus");
sigRoo->GetFunction("gaus")->SetLineColor(kGreen);
sigRoo->SetLineColor(kGreen);
sigRoo->SetMarkerColor(kGreen);

cout<<"My sig fit:"<<endl;
mySig->Fit("gaus");
mySig->GetFunction("gaus")->SetLineColor(kRed);
mySig->SetLineColor(kRed);
mySig->SetMarkerColor(kRed);

sigRoo->GetXaxis()->SetTitle("Significance");
sigRoo->GetYaxis()->SetTitle("Entries");

gStyle->SetOptStat(0000);
TLegend *leg=new TLegend(0.75,0.8,0.89,0.89);
leg->AddEntry(sigRoo,"RooStat","L");
leg->AddEntry(mySig,"MyModel","L");

double sro=sigRoo->GetFunction("gaus")->GetParameter(2);
double smy=mySig->GetFunction("gaus")->GetParameter(2);

double sroe=sigRoo->GetFunction("gaus")->GetParError(2);
double smye=mySig->GetFunction("gaus")->GetParError(2);

TLegend *leg2=new TLegend(0.11,0.8,0.3,0.89);
char fname[50];
int forget=sprintf(fname,"#sigma=%4.3f +- %4.3f",sro,sroe);
leg2->AddEntry(sigRoo,fname,"L");
forget=sprintf(fname,"#sigma=%4.3f +- %4.3f",smy,smye);
leg2->AddEntry(mySig,fname,"L");

TCanvas *cc2=new TCanvas();
cc2->cd();
sigRoo->Draw("E1");
mySig->Draw("E1,same");
leg->Draw("same");
leg2->Draw("same");
}