library(stats)
library(MBESS)

#Function for obtaining the confidence interval for a correlation
#Input is the correlation (givenr), sample size for that correlation (n)
#and the desired confidence interval (confidence)
#output is (1)lower endpoint, (2)upper endpoint, (3)interval width
NICIr<- function(givenr, n, confidence){
		ndraws<- 100000
		dfnew <- n-2
		alpha <- 1-confidence
		tobs <- givenr*sqrt(dfnew)/sqrt(1-givenr* givenr)

		ptnum <- rnorm(n= ndraws,mean=0,sd=1) + tobs*sqrt(rchisq(n= ndraws,df=dfnew,ncp=0)/(dfnew))
		ptden <- sqrt(rchisq(n= ndraws,df= dfnew+1,ncp=0)/(dfnew))
		pt <-ptnum/ptden
		r <- pt/sqrt(pt* pt + dfnew)
		end <- quantile(r, probs=c(alpha/2, (1-alpha/2)))
		interval <- as.numeric(end[2] - end[1])
		c(end[1],end[2],interval)
	}


#Function for obtaining the confidence interval for a standardized mean difference
#Input is the SMD (givend), sample size for that correlation (n)
#and the desired confidence interval (confidence)
#output is (1)lower endpoint, (2)upper endpoint, (3)interval width
#Kelley's ci.smd function is generally used except for very small effect sizes where there
#is an apparent discontinuity when smd == 0.
NICId<- function(givend, n, confidence){
		ndraws<- 1000000
		dfnew <- 2*n-2
		if (abs(givend)<.01){
			alpha <- 1-confidence	
			tobs <- givend*sqrt(n/2)
			posteriorL <- (rnorm(n= ndraws,mean=0,sd=1) + tobs*sqrt(rchisq(n= ndraws,df= dfnew,ncp=0)/(dfnew)))
			pd <- posteriorL/sqrt(n/2)
			end <- quantile(pd, probs=c(alpha/2, (1-alpha/2)))
			interval <- as.numeric(end[2] - end[1])
			} else {
			ci <- ci.smd(smd= givend,n.1=n,n.2= n, conf.level=confidence)
			end <- c(ci$Lower.Conf.Limit.smd, ci$Upper.Conf.Limit.smd)
			interval <- as.numeric(end[2] - end[1])
				}
		c(end[1],end[2],interval)
	}

################################################
#Function for estimating average statistical power across the posterior distribution of the effect size estimate
# PPOWER(effect, n1, n2, DesiredPower, lowern, uppern, alpha, type, posterior)
#Variable definitions:
#  effect:  the observed effect size
#  n1			sample size for group 1 (or observed sample size for the correlation)
#  n2			sample size for group 2 (NA in the case of the correlation)
#  DesiredPower	value of desired average power
#  lowern		Lower sample size bounds to examine
#  uppern		Upper sample size bounds to examine
#  alpha		Type I error rate (e.g., .05)
#  type			effect size type.  Options are "d" and "r".
#  posterior	posterior distribution of the effect size.  May be user-supplied (i.e., from a HB analysis).
#				Otherwise specify as NA to generate a posterior from the noninformative prior. 
#
#Returns:
#  1.  Sample size required to achieve the specified level of DesiredPower
#  2.  Median 90% confidence interval for the effect size based on that sample size
#  3.  Quantiles for power.
#  4.  A graph of Average power and Confidence interval width across the range of lowern to uppern.
################################################

PPOWER <- function(effect, n1, n2, DesiredPower, lowern, uppern, alpha, type, posterior){
	post <- TRUE
	if (is.na(posterior[1])==TRUE)  ndraws<- 50000
	if (is.na(posterior[1])==TRUE)  post <- NA
	if (is.na(posterior[1])==FALSE) ndraws<- nrow(posterior)

	if (type=='d'){
		dfobs <- n1+n2-2
		tobs <- effect*sqrt((n1*n2)/(n1+n2))
		
		val <- matrix(NA,nrow=uppern-lowern, ncol=4)
	
		for (i in lowern:uppern-1 ){
			dfnew <- i*2 - 2
			if (is.na(post)==TRUE) posterior <- rnorm(n=ndraws,mean=0,sd=1)/sqrt((n1*n2)/(n1+n2))+ effect*sqrt(rchisq(n=ndraws,df=dfobs,ncp=0)/(dfobs))
			pptnum <- rnorm(n=ndraws,mean=0,sd=1) + posterior*sqrt((i*i)/(2*i))
			pptden <- sqrt(rchisq(n= ndraws,df= dfnew,ncp=0)/(dfnew))
			ppt <- abs(pptnum/pptden)

			dist <- ecdf(ppt)
			SApower <- 1 - dist(qt(alpha/2,df=dfnew,ncp=0,lower.tail=FALSE))
			val[i -lowern +1 ,1]<- effect
			val[i- lowern +1, 2]<- NICId(quantile(abs(posterior),probs=c(.50)), i, .90)[3]
			val[i- lowern +1, 3]<- i
			val[i- lowern +1, 4]<- SApower
		}
	} else if (type=='r'){
		dfobs <- n1-2
		tobs <- effect*sqrt(dfobs)/sqrt(1- effect* effect)
		ndraws<- 50000

		val <- matrix(NA,nrow=uppern-lowern + 1, ncol=4)
	
		for (i in lowern:uppern ){
			dfnew <- i - 2
			posteriorL <- (rnorm(n= ndraws,mean=0,sd=1) + tobs*sqrt(rchisq(n= ndraws,df=dfobs,ncp=0)/(dfobs)))/sqrt(rchisq(n= ndraws,df=dfobs+1,ncp=0)/(dfobs))
			if (is.na(post)==TRUE) posterior <- posteriorL/sqrt(posteriorL*posteriorL + dfobs)
			pptnum <- (rnorm(n=ndraws,mean=0,sd=1) + (posterior*sqrt(dfnew/(1-posteriorr^2)))*sqrt(rchisq(n=ndraws,df=dfnew+1,ncp=0)/(dfnew)))
			pptden <- sqrt(rchisq(n= ndraws,df= dfnew,ncp=0)/(dfnew))
			ppt <- abs(pptnum/pptden)

			dist <- ecdf(ppt)
			SApower <- 1 - dist(qt(alpha/2,df=dfnew,ncp=0,lower.tail=FALSE))
			val[i -lowern +1 ,1]<- effect
			val[i- lowern +1, 2]<- NICIr(quantile(abs(posterior),probs=c(.50)), i, .90)[3]
			val[i- lowern +1, 3]<- i
			val[i- lowern +1, 4]<- SApower
		}
	}
	power<- smooth.spline(x=val[,3],y=val[,4])
	CIW<- smooth.spline(x=val[,3],y=val[,2])

	if (max(val[,4]) < DesiredPower) cat("\n\nuppern is too low. Please provide a higher value. An n of ", uppern," yields expected power of ",max(val[,4]), ".\n\n\n")
	if (min(val[,4]) > DesiredPower) cat("\n\nlowern is too high. Please provide a lower value. An n of ", lowern," yields expected power of ",min(val[,4]), ".\n\n\n") 

	f <- function(x)(predict(power, x)$y - DesiredPower)
	min_n <- uniroot(f,interval=c(val[1,3],  val[NROW(val),3]),tol = 0.001)$root
	
	plot(power, xlab="Sample Size for Future Study",ylab="Expected Statistical Power",frame.plot=FALSE,type="l",lwd=2)
	segments(lowern, DesiredPower,min_n, DesiredPower,col= "grey40", lwd=1,lty="dashed")
	segments(min_n, min(val[,4]),min_n, DesiredPower,col= "grey40", lwd=1,lty="dashed")
	text(uppern-5,predict(power, uppern)$y,"Expected Power",adj=c(1,0),cex=.7)	

	par(new=TRUE)
	plot(CIW,type="l",lty="dashed",lwd=2,main=NA,col=c("grey60"),frame.plot=FALSE,axes=FALSE,xlab=NA, ylab=NA)
	axis(side=4,at=NULL,tick=TRUE,outer=FALSE)
	text(uppern-1,(max(val[,2])-min(val[,2]))/2+min(val[,2]),"Median Confidence Interval Width",srt=90)
	segments(min_n,  predict(CIW, uppern)$y,min_n, predict(CIW, min_n)$y,col= "grey40", lwd=1,lty="dashed")
	segments(min_n, predict(CIW, min_n)$y, uppern, predict(CIW, min_n)$y,col= "grey40", lwd=1,lty="dashed")
	text(lowern+5,predict(CIW, lowern)$y,"90% CI Width",adj=c(0,1),cex=.7)	

	#Quantiles for Power
	if (is.na(posterior[1])==TRUE){
		ndraws<- 10000000
		nq <- round(min_n)
		dfq <- 2*nq - 2
		if (type=='d') Q <- rnorm(n=ndraws,mean=0,sd=1)/sqrt((n1*n2)/(n1+n2))+ effect*sqrt(rchisq(n=ndraws,df=dfobs,ncp=0)/(dfobs))
		if (type=='r'){
			posteriorL <- (rnorm(n= ndraws,mean=0,sd=1) + tobs*sqrt(rchisq(n= ndraws,df=dfobs,ncp=0)/(dfobs)))/sqrt(rchisq(n= ndraws,df=dfobs+1,ncp=0)/(dfobs))
			Q <- posteriorL/sqrt(posteriorL*posteriorL + dfobs)
			} 
		} else {
			if (type=='d') Q <- posterior
			if (type=='r'){
				Q <- posterior/sqrt(posterior*posterior + dfobs)
			} 
			quant <- quantile(abs(Q), probs=c(.1,.2,.4,.5,.6,.8))
			PowerQ1<- power.t.test(n = round(min_n), delta = quant, sd = 1, sig.level = 0.05, power = NULL, type = c("two.sample"), alternative = c("two.sided"), strict = FALSE)$power
		}
	

	if (type=='d') cat("Required sample size per group is ",min_n," for average (expected) power of ", DesiredPower, ".\n\n")
	if (type=='r') cat("Required sample size is ",min_n," for average (expected) power of ", DesiredPower, ".\n\n")
	cat("Median 90% Confidence Interval Width based on this sample size is ", predict(CIW, min_n)$y, ".\n\n")
	cat("Quantiles for power are:\n ")
	print(PowerQ1)
}




################################################
#Function for producing a sample from the posterior distribution using Hierarchical Bayesian MCMC
#H(estimate, sd, df, n.burnin, n.thin, k)
#Variable definitions:
#  estimate:  	the observed effect size estimates.  Note that the first estimate is the one of interest
#  sd:			The standard deviation associated with the effect size estimate
#  df:			The degrees of freedome associated with the effect size estimate
#  n.burnin:	How many iterations to conduct before starting to draw samples from the posterior
#  n.thin		How many iterations between samples from the posterior
#  k			The desired number of samples from the posterior
#
#Returns:
#  1.  a vector of draws from the posterior distribution of the first effect size estimate.
################################################

HB <- function(estimate, sd, df, n.burnin, n.thin, k){
	
	summary(estimate)
	J <- NROW(estimate)
	y <- estimate
	sigma.y <- sd
	n.thin<-5
	n.burnin<-10000
	k<- 50000
	n.iter <- (k*n.thin) + n.burnin
	
	theta.update <- function (){
		theta.hat <- (mu/tau^2 + y/sigma.y^2)/(1/tau^2 + 1/sigma.y^2)
		V.theta <- 1/(1/tau^2 + 1/sigma.y^2)
		rnorm (J, theta.hat, sqrt(V.theta))
	}
	
	mu.update <- function(){
		rnorm (1, mean(theta), tau/sqrt(J))
	}

	tau.update <- function(){
		sqrt(sum((theta-mu)^2/rchisq(1,J-1)))
	}
	
	sims <- array (NA, c((n.iter-n.burnin)/n.thin, J))
	mu <- rnorm (1, mean(y), sd(y))
	tau <- runif (1, 0, sd(y))
	
	for (t in 1:n.iter){
		theta <- theta.update()
		mu <- mu.update()
		tau <- tau.update()
		if ( (t >n.burnin)&((t - n.burnin) %% n.thin == 0)) sims[(t-n.burnin)/n.thin,] <- c(theta)
	}

	#Note that the study of interest is the first effect size.
	as.matrix(sims[,1])

	}


############################################################################################
# reml(effect, var)  returns the REML  random effects variance (tau) and meta-analytic mean (mu) estimates
#    REML estimation of the random effects variance based on Equation [18] from:
#    Viechtbauer, W. (2005). Bias and efficiency of meta-analytic variance estimators in the random-effects model.
#    Journal of Educational and Behavioral Statistics, 30, 261-293. 
#
#Returns the REML estimates of the variance and mean of the random effects
############################################################################################

reml <- function(yi, vi) {
	maxiter	=  200
	change	<- 1
	iter	<- 0
		
	# starting value for random effects variance
	vart		<- var(yi) - 1/length(yi)*sum(vi)				   
	# initialize to 0 in case original estimate is negative
	vart[vart < 0] <- 0					

		while (change > 0.00001) {
			iter	<- iter + 1
			varm	<- vart
			wi		<- 1/(vi + vart)
			
			#ML estimate of the weighted mean
			mu <- sum((yi*wi)/sum(wi))

			vart <- sum(wi*wi*((yi-mu)*(yi-mu)-vi))/sum(wi*wi)+1/sum(wi)
			
			change	<- abs(varm - vart)
			if (iter > maxiter) {
				break
			}
		}	

	#setting variance to 0 if negative
	vart[vart < 0] <- 0
	
	output <- list(vart, mu)
	names(output) <- c("vart", "mu")
	output
}


#Function called by the EB function.  See footnote 5 in Biesanz & Schrager(2010)
posteriordelta <- function(d, w, df, bootmu, boottau2, N){
	grid <- seq(from=-25,to=25,by=.02)
	dpost <- dt(d/w,df=df,ncp=grid)*dnorm(grid*w,mean= bootmu,sd=sqrt(boottau2))
	weight <- dpost/sum(dpost)

	sample(grid,size=N,replace=TRUE,prob= weight)*w
	}

EB <- function(estimate, sd, n1, n2,  k, mu, tau2, metan1, metan2){

	df <- n1+n2-2
	w <- sqrt(1/n1 + 1/n2)
	se <- ((n1+n2)/(n1*n2)+(estimate*estimate)/(2*(n1+n2-2)))*((n1+n2)/(n1+n2-2))
	PostSampleN <- 10	#number of posterior draws to take from each iteration

	EBsims <- array (NA, 0, 1)	#array of simulated values from the posterior distribution
	m <- NROW(metan1)  				#number of studies in the meta-analysis
	niter <- k/PostSampleN
	for (i in 1: niter){
		thetastar <- rnorm(m, mean=mu, sd=sqrt(tau2))
		dstar <- rt(m, (metan1 + metan2-2), thetastar/sqrt(1/metan1 +1/metan2))*sqrt(1/metan1 +1/metan2)
		dstarvar <-((metan1 + metan2)/(metan1* metan2)+(dstar* dstar)/(2*(metan1 + metan2-2)))*((metan1 + metan2)/(metan1 + metan2-2))

		dstar <-c(estimate, dstar)
		dstarvar <-c(se*se, dstarvar)
		metastar <- reml(dstar, dstarvar)

		mustar <- metastar$mu
		tau2star <- metastar$vart
	
		EBsims <- c(EBsims,posteriordelta(estimate,w,df, mustar,tau2star, PostSampleN))
	}
	as.matrix(EBsims)
}