# -------------------------------------------------------------------------
# -------------------------------------------------------------------------
#
# Course:       STAT 438, Spring 2020
# Professor:    Dr. Debdeep Pati
# Grader:       Alexander Coulter
# Project:      Beta Distribution HPDI
# Date Created: 2020-02-20
# Last Update:  2020-02-20
#
# -------------------------------------------------------------------------
# -------------------------------------------------------------------------
#
# Function to calculate HPDI for beta distribution beta(a,b).
# Returns* interval in form (L, U).
# If a, b < 1, then HPDI is union of disjoint intervals, (0,L) and (U,1).
# *If a, b < 1, then HPDI returned as pair (U, L), i.e. larger number first.
# If a, b = 1, then returns interval centered at 0.5.
# If a or b <= 0, then returns NA.
# Can set parameters (default a = 1, b = 1), interval coverage (default 0.95), and tolerance (default 1e-9).
#
# -------------------------------------------------------------------------
# -------------------------------------------------------------------------
# Function:
BETA.HPDI = function(a=1,b=1,cover=0.95,tol=1e-9){
	if(a <= 0 | b <= 0){
		return(NA)
	} else if(a == 1 & b == 1){
		return(c((1-cover)/2, 1-(1-cover)/2))
	} else if(a <= 1 & b >= 1){
		return(c(0, qbeta(cover,a,b)))
	} else if(a >= 1 & b <= 1){
		return(c(qbeta(1-cover,a,b),1))
	} else if(a > 1 & b > 1){
		L = M = (a-1)/(a+b-2)
		j = 1
		i = 0
		while(abs(i - cover) > tol){
			if(i < cover) L = L - M/(2^j)
			if(i > cover) L = L + M/(2^j)
			U = uniroot(f = function(x){dbeta(x,a,b)-dbeta(L,a,b)},lower=M,upper=1,tol=1e-9)$root
			i = pbeta(U,a,b) - pbeta(L,a,b)
			j = j + 1
		}
		return(c(L,U))
	}
	if(a < 1 & b < 1){
		L = M = (a-1)/(a+b-2)
		j = 1
		i = 0
		while(abs(i - (1-cover)) > tol){
			if(i < (1 - cover)) L = L - M/(2^j)
			if(i > (1 - cover)) L = L + M/(2^j)
			U = uniroot(f = function(x){dbeta(x,a,b)-dbeta(L,a,b)},lower=M,upper=1,tol=1e-9)$root
			i = pbeta(U,a,b) - pbeta(L,a,b)
			j = j + 1
		}
		return(c(U,L))
	}
}

# -------------------------------------------------------------------------
# -------------------------------------------------------------------------
# Test Beta HPDI function - keep in mind above notes:
# (not run on "source" call)
if(FALSE){
	
	# a, b < 1 (HPDI is disjoint union):
	hpdi = BETA.HPDI(a=0.3,b=0.5,cover=0.95); hpdi
	pbeta(hpdi[2],0.3,0.5) + pbeta(hpdi[1],0.3,0.5,lower.tail=FALSE)
	
	# a, b = 1 (HPDI is centered at 0.5):
	hpdi = BETA.HPDI(a=1,b=1,cover=0.95); hpdi
	pbeta(hpdi[2],1,1) - pbeta(hpdi[1],1,1)
	
	# a, b > 1; coverage = 0.99:
	hpdi = BETA.HPDI(a=2,b=10,cover=0.99); hpdi
	pbeta(hpdi[2],2,10) - pbeta(hpdi[1],2,10)
	
	# a > 1, b = 1; coverage = 0.99:
	hpdi = BETA.HPDI(a=10,b=1,cover=0.99); hpdi
	pbeta(hpdi[2],10,1) - pbeta(hpdi[1],10,1)
	
	# a = -1:
	hpdi = BETA.HPDI(a=-1,b=1); hpdi

}