from math import *
from util import pyroot_util
import ROOT

from ROOT import TFeldmanCousins
from ROOT.Math import poisson_cdf
from ROOT.TMath import Poisson as poisson, ChisquareQuantile

def limit( k , mu0, cl=0.90) :
	"given the background is mu0, compute the limit on signal(!) resulting from observering k events "

	return 0.5 * ChisquareQuantile( cl, 2*k+2 ) - mu0


def limit_FC( k , mu0, cl=0.90) :
	"given the background is mu0, compute the limit on signal(!) resulting from observering k events "

	# return 0.5 * ChisquareQuantile( cl, 2*k+2 ) - mu0 #original
	f = TFeldmanCousins(cl)
	fmumax=f.GetMuMax()
	pp_fmumax=ROOT.Math.poisson_pdf(int(fmumax),mu0)
	while pp_fmumax>1e-15:
	     f.SetMuMax(fmumax*2);
	     fmumax=f.GetMuMax();              
	     pp_fmumax=ROOT.Math.poisson_pdf(int(fmumax),mu0)

	ul=f.CalculateUpperLimit(k,mu0)
	return ul


def mean_limit( mu0, cl = 0.90, FC=0 ) :

	c = 0

	kmax = int( max( 20, mu0 + 5 * sqrt(mu0) ) )

	for k in range( kmax ) :
		# print ( "mean", k, mu0, poisson( k,mu0) , limit( k, mu0, cl ) )
		if (FC==1): l= limit_FC( k, mu0, cl )
		else: l= limit( k, mu0, cl )
		c += poisson( k, mu0 ) * l

	return c



def get_mu_lds(nbkg,alpha,beta):

    nmax=100000
    #stepsize=0.1
    #nmax = int( max( 20, nbkg + 5 * sqrt(nbkg) ) )
    #print("nmax=",nmax,int(nmax/stepsize))
    for n_obs in range(nmax):
        pvalue=1-ROOT.Math.poisson_cdf(n_obs,nbkg)
        if pvalue < alpha:
           n_crit=1+n_obs #original
           # n_crit=n_obs # new
           break

    #print("nbkg,ncrit:",nbkg,n_crit)
    #for cont in range(int(nmax/stepsize)):
    for cont in range(nmax):

        mu_lds=cont*0.1 # optimize step size
        p=1-ROOT.Math.poisson_cdf(n_crit-1,nbkg+mu_lds) #original
        #p=1-ROOT.Math.poisson_cdf(n_crit,nbkg+mu_lds) # new
        #print("cont,mu_lds,p:",cont,mu_lds,p)
        if p > (1-beta): break
    return mu_lds

def get_model_discovery_potential(nsrc,nbkg,alpha,beta):
     mu_lds=get_mu_lds(nbkg,alpha,beta)
     mdp=mu_lds/nsrc
     return  mdp