###This file contains all functions required for generate_results.R

################################################################################
#Functions for finding the Bayes expected gain of the fixed sampling methods
################################################################################
##FEgain
################################################################################
##Function for expected gain of fixed enrichment trial
#Inputs: alpha  = FWER
#			 lambda = proportion of the population in sub-pop
#			 info   = ratio of variance and sample size
#			 n      = number of simulations, single value
#			 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#			 mu     = prior mean vector
#			 sigma  = prior variance vector
#(vectors should be of length 2, entry 1 in sub-pop 1, entry 2 in sub-pop 2)
#Ouput as list: 	prob1 = probability of rejecting H_01
#			gain  = Bayes expected gain of the trial

#Fixed Fnrichment Bayes expected gain (by simulation)
FEgain <- function(alpha,lambda,info,n,delta,theta1){
	#probability of rejection
	prob1 <- 1-pnorm(sqrt(1/info)*qnorm(1-alpha)+delta[1],theta1,sqrt(1/(info)))
	#expected gain (mean of realised gains)
	gain  <- mean(prob1*lambda*theta1)
	#output both probabilities and the gain
	list(mean(prob1),gain)
}

################################################################################
##MHgain
################################################################################
##Fixed sampling trial testing multiple hypotheses Bayes expected gain
#Inputs: alpha  = FWER
#			 lambda = proportion of the population in sub-pop
#			 info   = ratio of variance and sample size
#			 n      = number of simulations, single value
#			 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#			 mu     = prior mean vector
#			 sigma  = prior variance vector
#(vectors should be of length 2, entry 1 in sub-pop 1, entry 2 in sub-pop 2)
#Output as list: 	prob = c(prob1,prob3,probb)
#     prob1 = probability of rejecting H_01
#			prob3 = probability of rejecting H_03
#			probb = probability of rejecting both
#     ##note probabilities involving H_03 not present in fixed enrichment trial
#			gain  = Bayes expected gain of the trial
MHgain <- function(alpha,lambda,info,n,delta,theta1,theta2){
	theta3 <- lambda*theta1 + (1-lambda)*theta2
	#first stage of the trial	
	#simulate observations from the first stage of the trial
	thetahat1 <- rnorm(n,theta1,sqrt(1/(lambda*info)))
	thetahat2 <- rnorm(n,theta2,sqrt(1/((1-lambda)*info)))
	thetahat3 <- lambda*thetahat1 + (1-lambda)*thetahat2
	#find corresponding p-values
	p1 <- 1-pnorm(thetahat1,mean=delta[1],sd=sqrt(1/(lambda*info)))
	p3 <- 1-pnorm(thetahat3,mean=delta[2],sd=sqrt(1/(info)))
	#the p-value under the intersection hypothesis (given be sime's method)
	twomin <- 2*((p1<p3)*p1 + (1-(p1<p3))*p3)
	maxp <- (p1>p3)*p1 + (1-(p1>p3))*p3
	p13 <- (twomin<maxp)*twomin + (1-(twomin<maxp))*maxp

	prob1 <- mean((p1<alpha)*(p3>alpha)*(p13<alpha))
	prob3 <- mean((p1>alpha)*(p3<alpha)*(p13<alpha))
	probb <- mean((p1<alpha)*(p3<alpha)*(p13<alpha))
	gain  <- mean(lambda*theta1*(p1<alpha)*(p3>alpha)*(p13<alpha)+
			theta3*(p3<alpha)*(p13<alpha))

	list(prob=c(prob1,prob3,probb),gain)
}

################################################################################
#Functions for finding the grid of Bayes optimal decisions
################################################################################

################################################################################
##Sub grid
################################################################################
#This function is used to create a sub-grid of points at which to 
#evaluate the optimal decision, this sub-grid adds new columns and rows
#at every midpoint of the old rows and columns
#If all surrounding points have the same value we assign that value to
#the new point, if not we leave it to be evaluated later. This is done
#in order to avoid additional computation where it is not necessary
#between the previous columns and rows
#Inputs:	original = the original gird
#		rows	 = the x axis values of the old grid
#		columns	 = the y axis values of the old grid
#Outputs as list:	grid = the new grid
#		x    = the new x axis values
#		y    = the new y axis values
subgrid <- function(original,rows,cols){
	#new rows and columns
	rows1 <- seq(from=min(rows),to=max (rows),
		   length.out=2*length(rows)-1)
	cols1 <- seq(from=min(cols),to=max(cols),
		   length.out=2*length(cols)-1)

	#set up the new matrix of points
	grid1 <- matrix(rep(-1,length(rows1)*length(cols1)),
		   nrow=length(rows1))

	#where the point already existed give it the same value
	grid1[which(rows1%in%rows),which(cols1%in%cols)] <- original

	#where the new points are in an existing row or column check
	#the next point and the point behind in that row/column
	grid1[which(rows1%in%rows),which(!cols1%in%cols)] <- (
		grid1[which(rows1%in%rows),which(!cols1%in%cols)-1] ==
		grid1[which(rows1%in%rows),which(!cols1%in%cols)+1])* 
		grid1[which(rows1%in%rows),which(!cols1%in%cols)+1] +
		(grid1[which(rows1%in%rows),which(!cols1%in%cols)-1] !=
		grid1[which(rows1%in%rows),which(!cols1%in%cols)+1])* 
		grid1[which(rows1%in%rows),which(!cols1%in%cols)]

	grid1[which(!rows1%in%rows),which(cols1%in%cols)] <- (
		grid1[which(!rows1%in%rows)-1,which(cols1%in%cols)] ==
		grid1[which(!rows1%in%rows)+1,which(cols1%in%cols)])* 
		grid1[which(!rows1%in%rows)+1,which(cols1%in%cols)]+
		(grid1[which(!rows1%in%rows)-1,which(cols1%in%cols)] !=
		grid1[which(!rows1%in%rows)+1,which(cols1%in%cols)])* 
		grid1[which(!rows1%in%rows),which(cols1%in%cols)]

	#if the new point lies in a new row and column check the four
	#surrounding points
	grid1[which(!rows1%in%rows),which(!cols1%in%cols)] <- (
		grid1[which(!rows1%in%rows)+1,which(!cols1%in%cols)+1]==
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)-1]&
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)-1]==
		grid1[which(!rows1%in%rows)+1,which(!cols1%in%cols)-1]&
		grid1[which(!rows1%in%rows)+1,which(!cols1%in%cols)-1]==
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)+1])*
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)+1]+
		(1-(grid1[which(!rows1%in%rows)+1,which(!cols1%in%cols)+1]==
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)-1]&
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)-1]==
		grid1[which(!rows1%in%rows)+1,which(!cols1%in%cols)-1]&
		grid1[which(!rows1%in%rows)+1,which(!cols1%in%cols)-1]==
		grid1[which(!rows1%in%rows)-1,which(!cols1%in%cols)+1]))*
		grid1[which(!rows1%in%rows),which(!cols1%in%cols)]

	#output the result
	list(grid=grid1,x=rows1,y=cols1)
}

################################################################################
##decisiongrid
################################################################################
##An iterative function for finding a grid approximation to the Bayes optimal 
##decision function to plot
#Inputs: original = the original grid which will be used as the starting point
#		    note this will provide the extremes of the plot, the density
#		    of the original grid may need to be altered to ensure the
#		    plot is produced correctly
#	 rows	  = the starting values of theta1
#	 cols   = the starting values of theta2
#	 maxit  = how many times will the grid process be iterated
#	 alpha  = FWER
#	 lambda = proportion of the population in sub-pop
#	 tau    = proportion of the sample in the first stage
#	 info   = ratio of variance and sample size
#	 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#	 m      = simulation size
#  decis  = function to be used to compute the optimal decision
#  dparm  = additional parameters to be passed to the decision function as a
#           list
#Outputs as list: grid   = the grid of optimal decisions
#	 rows	  = values of theta1 associated with the grid
#	 cols   = values of theta2 associated with the grid
decisiongrid <- function(original,rows,cols,maxit=8,
					alpha,lambda,tau,info,m,delta,decis,dparm){
	#split info into the stages
	temp <- c(info*tau,info*(1-tau))
	info <- temp
	rm(temp)

	#rename the variables ready for the loop
	newgrid <- original
	newx <- rows
	newy <- cols
	#we keep track of how many grid points we assess the optimal decision at
	assess <- 0

	for(iter in (1:maxit))
	{
		#use the subgrid function to create the next level of grid
		update <- subgrid(newgrid,newx,newy)
		newgrid <- update$grid
		newx <- update$x
		newy <- update$y

		#approximate the Bayes optimal decision at the required points
		#first we find the values at which we must approximate the decision
		ytemp <- c(which(newgrid==-1)/(dim(newgrid)[1]))
		yfind <- floor(ytemp) + ceiling(ytemp-floor(ytemp))
		xfind <- c(which(newgrid==-1) - (yfind-1)*(dim(newgrid)[1]))
		xfind[which(xfind==0)] <- dim(newgrid)[1]
				
		#add the number of points which must be assessed
		assess <- assess + length(newgrid[which(newgrid==-1)])
		print(assess)

		#assess the new points 
		temp <- decis(newx[xfind],newy[yfind],
				alpha,lambda,tau,info,delta,m,dparm)

		newgrid[which(newgrid==-1)] <- (temp[1,]>temp[2,])+
			    (temp[1,]>(1+pctd)*temp[2,])-
			    2*(temp[2,]>(1+pctd)*temp[1,])
 	}
	#list for output
	list(grid=newgrid,rows=newx,cols=newy)
}

################################################################################
##cgrid
################################################################################
##compute the grid of optimal decisions to be used in trial simulation, this 
##function computes the intial grid then calls decisiongrid to iterate
#Inputs: thetahat1 = vector of initial gridvalues for thetahat1
#  thetahat2 = vector of initial gridvalues for thetahat2
#  maxit  = how many times will the grid process be iterated
#	 alpha  = FWER
#	 lambda = proportion of the population in sub-pop
#	 tau    = proportion of the sample in the first stage
#	 info   = ratio of variance and sample size
#	 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#	 m      = simulation size
#  decis  = function to be used to compute the optimal decision
#  dparm  = additional parameters to be passed to the decision function as a
#           list
#Outputs as list: grid   = the grid of optimal decisions
#	 rows	  = values of theta1 associated with the grid
#	 cols   = values of theta2 associated with the grid
cgrid <- function(thetahat1,thetahat2,maxit,pctd,
				alpha,lambda,tau,info,m,delta,decis,dparm){
	#find all the points for the intial grid
	thetahat <- expand.grid(thetahat1,thetahat2)
 
	#split info into the stages
	temp <- c(info*tau,info*(1-tau))
	sinfo <- temp
	rm(temp)

	#find the Bayes optimal decision at each of the intial grid points
	init <- decis(thetahat[,1],thetahat[,2],alpha,lambda,tau,sinfo,delta,m,dparm)
	
	#evaluate whether the decision is clear or whether is is close to another decision
	initmat <- matrix((init[1,]>init[2,])+
		    (init[1,]>(1+pctd)*init[2,])-
		    2*(init[2,]>(1+pctd)*init[1,]),nrow=length(thetahat1))

	#now run the algorithm to find a more detailed decision grid
	decisions <- decisiongrid(initmat,thetahat1,thetahat2,maxit,
					alpha,lambda,tau,info,m,delta,decis,dparm)
	decisions
}

###############################################################################
##bayesopt
################################################################################
##function computing the bayesoptimal decision, calls Ogain for gain of each
##posterior simulation
#Inputs: thetahat1 = value of thetahat1 to compute bayes optimal decision
#  thetahat2 = value of thetahat2 to compute bayes optimal decision
#	 alpha  = FWER
#	 lambda = proportion of the population in sub-pop
#	 tau    = proportion of the sample in the first stage
#	 info   = ratio of variance and sample size
#	 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#	 m      = simulation size
#  dparm  = additional parameters to be passed to the decision function as a
#           list, this contains
#         pmu  = vector of prior means
#         psig = vector of prior variances
#         rho  = prior correlation
#         a    = repackaged prior information for computation
#         b    = repackaged prior information for computation
#         c    = repackaged prior information for computation
#Output:  gain = Bayes expected gain of each decision (enrich or not)
bayesopt <- function(thetahat1,thetahat2,alpha,lambda,tau,info,delta,m,dparm){	
  pmu <- dparm[[1]]
  psig <- dparm[[2]]
  rho <- dparm[[3]]
  a <- dparm[[4]]
  b <- dparm[[5]]
  c <- dparm[[6]]
  
  #drawing from the posterior, to choose an exact theta simply set sigma = 0
  
  postmu_1=a[1] + b[1,1]*thetahat1 + b[1,2]*thetahat2
  postmu_2=a[2] + b[2,1]*thetahat1 + b[2,2]*thetahat2
  postsig=sqrt(c(c[1,1],c[2,2]))
  postrho=c[1,2]/sqrt(c[1,1]*c[2,2])
  
  theta1 <- rnorm(n,postmu_1,postsig[1])
  theta2 <- rnorm(n,postmu_2+(postsig[2]/postsig[1])*postrho*(theta1-postmu_1),
                  sqrt(1-postrho^2)*postsig[2])
  
  #find rejection regions from first stage of the trial
  thetahat3 <- lambda*thetahat1 + (1-lambda)*thetahat2
  #find corresponding p-values
  p1_1 <- 1-pnorm(thetahat1,mean=delta[1],sd=sqrt(1/(lambda*info[1])))
  p3_1 <- 1-pnorm(thetahat3,mean=delta[2],sd=sqrt(1/(info[1])))
  #the p-value under the intersection hypothesis (given be sime's method)
  twomin <- 2*((p1_1<p3_1)*p1_1 + (1-(p1_1<p3_1))*p3_1)
  maxp <- (p1_1>p3_1)*p1_1 + (1-(p1_1>p3_1))*p3_1
  p13_1 <- (twomin<maxp)*twomin + (1-(twomin<maxp))*maxp
  #and the Z-values
  Z1_1 <- (thetahat1-delta[1])*sqrt(lambda*info[1])
  Z3_1 <- (thetahat3-delta[2])*sqrt(info[1])
  Z13_1<- qnorm(1-p13_1)
  
  #testing weights
  w1 <- sqrt(tau)
  w2 <- sqrt(1-tau)
  
  R1 <- (qnorm(1-alpha) - w1*Z1_1)/w2
  R3 <- (qnorm(1-alpha) - w1*Z3_1)/w2
  R13 <- (qnorm(1-alpha) - w1*Z13_1)/w2
  
  gain <- Ogain(R1,R3,R13,alpha,lambda,tau,info,delta,theta1,theta2,m)
  c(mean(gain[[1]]),mean(gain[[2]]))
  
}

#Vectorize is used to allow the arguements thetahat1 and thetahat2 to be input
#as vectors, the resulting output will be a 2 x length(thetahat1) matrix of 
#Bayes expected gain
vbayesopt <- Vectorize(bayesopt,vectorize.args=c("thetahat1","thetahat2"))

################################################################################
##Ogain
################################################################################
##Function finding the oracle gain for an AE trial under each decision (so named
##because it is computed with knowledge of true theta)
#Inputs: thetahat_1
#	 		 thetahat_2
#			 thetahat_3
#			 alpha  = FWER
#			 lambda = proportion of the population in sub-pop
#			 tau    = proportion of the sample in the first stage
#			 info   = ratio of variance and sample size
#			 n      = number of simulations, single value
#			 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#			 theta  = true theta values
#			 m      = simulation size
#Outputs as list: enrich = gain of enriching
#			  full   = gain of continuing in the full population
Ogain <- function(R1,R3,R13,alpha,lambda,tau,info,delta,theta1,theta2,m){
  #from the first stage obs
  #first the enriched trial
  thetahat1_e <- rnorm(m,theta1,sqrt(1/info[2]))
  Z1_e <- (thetahat1_e-delta[1])*sqrt(info[2])
  enrich <- lambda*theta1*(Z1_e > R1)*(Z1_e > R13)
  
  #now the full population trial
  #simulate observations from the second stage of the trial
  thetahat1_2 <- rnorm(m,theta1,sqrt(1/(lambda*info[2])))
  thetahat2_2 <- rnorm(m,theta2,sqrt(1/((1-lambda)*info[2])))
  thetahat3_2 <- lambda*thetahat1_2 + (1-lambda)*thetahat2_2
  #find corresponding p-values
  p1_2 <- 1-pnorm(thetahat1_2,mean=delta[2],sd=sqrt(1/(lambda*info[2])))
  p3_2 <- 1-pnorm(thetahat3_2,mean=delta[2],sd=sqrt(1/(info[2])))
  #the p-value under the intersection hypothesis (given be sime's method)
  twomin <- 2*((p1_2<p3_2)*p1_2 + (1-(p1_2<p3_2))*p3_2)
  maxp <- (p1_2>p3_2)*p1_2 + (1-(p1_2>p3_2))*p3_2
  p13_2 <- (twomin<maxp)*twomin + (1-(twomin<maxp))*maxp
  #and the Z-values
  Z1_2 <- (thetahat1_2-delta[1])*sqrt(lambda*info[2])
  Z3_2 <- (thetahat3_2-delta[2])*sqrt(info[2])
  Z13_2<- qnorm(1-p13_2)
  
  
  full  <- lambda*theta1*(Z1_2 > R1)*(Z3_2 < R3)*(Z13_2 > R13)+
    (lambda*theta1+(1-lambda)*theta2)*(Z3_2 > R3)*(Z13_2 > R13)
  
  list(enrich,full)
}

################################################################################
#Functions for plotting the grid of Bayes optimal decisions
################################################################################
##plotdecis
################################################################################
##function to plot the Bayes optimal decision
#Input: decis = list containing decision grid and corresponding thetahat values
#Output: Plot of the Bayes optimal decision rule
plotdecis <- function(decis){
#find co-ordinates of each decision to make the plot clearer
ytemp <- c(which(decis$grid<=0.5)/(dim(decis$grid)[1]))
ytemp <- floor(ytemp) + ceiling(ytemp-floor(ytemp))
xtemp <- c(which(decis$grid<=0.5)- (ytemp-1)*(dim(decis$grid)[1]))
xtemp <- decis$rows[xtemp]
ytemp <- decis$cols[ytemp]
x1 <- xtemp[which(abs(xtemp - round(xtemp)) < .Machine$double.eps & 
                    abs(ytemp - round(ytemp)) < .Machine$double.eps )]
y1 <- ytemp[which(abs(xtemp - round(xtemp)) < .Machine$double.eps & 
                    abs(ytemp - round(ytemp)) < .Machine$double.eps )]
rm(xtemp)
rm(ytemp)

ytemp <- c(which(decis$grid>=0.5)/(dim(decis$grid)[1]))
ytemp <- floor(ytemp) + ceiling(ytemp-floor(ytemp))
xtemp <- c(which(decis$grid>=0.5)- (ytemp-1)*(dim(decis$grid)[1]))
temp  <- expand.grid(decis$rows,decis$cols)
xtemp <- decis$rows[xtemp]
ytemp <- decis$cols[ytemp]
x2 <- xtemp[which(abs(xtemp - round(xtemp)) < .Machine$double.eps & 
                    abs(ytemp - round(ytemp)) < .Machine$double.eps )]
y2 <- ytemp[which(abs(xtemp - round(xtemp)) < .Machine$double.eps & 
                    abs(ytemp - round(ytemp)) < .Machine$double.eps )]
rm(xtemp)
rm(ytemp)

temp1 <- decis$grid
temp1[which(decis$grid<0)] <- 0
temp1[which(decis$grid>0)] <- 1

temp1 <- subgrid(temp1,decis$rows,decis$cols)

ytemp <- c(which(temp1$grid ==-1)/(dim(temp1$grid)[1]))
yfind <- floor(ytemp) + ceiling(ytemp-floor(ytemp))
xfind <- c(which(temp1$grid ==-1) - (yfind-1)*(dim(temp1$grid)[1]))
xfind[which(xfind==0)] <- dim(temp1)[1]

rm(ytemp)

plot(temp1$x[xfind],temp1$y[yfind],xlab="",ylab="",
     type="l",lwd=1,col="black",
     xlim=c(min(decis$rows),max(decis$rows)),
     ylim=c(min(decis$cols),max(decis$cols)))

temp2 <- decis$grid
temp2[which(decis$grid!=-2)] <- 0

temp2 <- subgrid(temp2,decis$rows,decis$cols)

ytemp <- c(which(temp2$grid ==-1)/(dim(temp2$grid)[1]))
yfind <- floor(ytemp) + ceiling(ytemp-floor(ytemp))
xfind <- c(which(temp2$grid ==-1) - (yfind-1)*(dim(temp2$grid)[1]))
xfind[which(xfind==0)] <- dim(temp2)[1]

rm(ytemp)

points(temp2$x[xfind],temp2$y[yfind],type="l",lwd=1,col="blue")

temp3 <- decis$grid
temp3[which(decis$grid!=2)] <- 0

temp3 <- subgrid(temp3,decis$rows,decis$cols)

ytemp <- c(which(temp3$grid ==-1)/(dim(temp3$grid)[1]))
yfind <- floor(ytemp) + ceiling(ytemp-floor(ytemp))
xfind <- c(which(temp3$grid ==-1) - (yfind-1)*(dim(temp3$grid)[1]))
xfind[which(xfind==0)] <- dim(temp2)[1]

rm(ytemp)

points(temp3$x[xfind],temp3$y[yfind],type="l",lwd=1,col="red")

}

##addleg
##function to add a legend to the plotdecis plot
#Input: thetahat1 = points to add legend for thetahat1
# thetahat2 = points to add to legend for thetahat 2

addleg <- function(thetahat1,thetahat2,alpha,lambda,tau,info,delta,m,dparm){

decis <- vdgrid(thetahat1,thetahat2,alpha,lambda,tau,info,delta,m,dparm)
x1 <- thetahat1[which(decis==1)]
y1 <- thetahat2[which(decis==1)]

x2 <- thetahat1[which(decis==3)]
y2 <- thetahat2[which(decis==3)]

points(x1,y1,pch=19,cex=0.6,col="blue")
points(x2,y2,pch=2,cex=0.6,col="black")
legend("topright",c("Continue","Enrich"),pch=c(2,19),col=c("black","blue"),bg="white")
mtext((expression(hat(theta)[2]^(1))),side=2,las=1,line=2)
mtext((expression(hat(theta)[1]^(1))),side=1,las=1,line=3)
}

################################################################################
#Functions for simulating to estimate trial performance
################################################################################

################################################################################
##simAE
################################################################################
##Function simulating Adaptive ENrichment trials
#Inputs: alpha  = FWER
#			 lambda = proportion of the population in sub-pop
#			 tau    = proportion of the sample in the first stage
#			 info   = ratio of variance and sample size
#			 n      = number of simulations, single value
#			 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#			 theta  = vector of theta values to simulate from
#				    (note supply a single value for true parameter)
#			 decis  = function to give the interim decision
#			 dparm  = list of decision function parameters
#(vectors should be of length 2, entry 1 in sub-pop 1, entry 2 in sub-pop 2)
#Ouput as list: 	prob1 = probability of rejecting H_01
#			prob3 = probability of rejecting H_03
#			probb = probability of rejecting both
#			gain  = Bayes expected gain of the trial
simAE <- function(alpha,lambda,tau,info,n,delta,theta1,theta2,decis,dparm){
#first some preperation
	#split info into the stages
	temp <- c(info*tau,info*(1-tau))
	info <- temp
	rm(temp)
#first stage of the trial	
	#simulate observations from the first stage of the trial
	thetahat1_1 <- rnorm(n,theta1,sqrt(1/(lambda*info[1])))
	thetahat2_1 <- rnorm(n,theta2,sqrt(1/((1-lambda)*info[1])))
	thetahat3_1 <- lambda*thetahat1_1 + (1-lambda)*thetahat2_1
	#find corresponding p-values
	p1_1 <- 1-pnorm(thetahat1_1,mean=delta[1],sd=sqrt(1/(lambda*info[1])))
	p3_1 <- 1-pnorm(thetahat3_1,mean=delta[2],sd=sqrt(1/(info[1])))
	#the p-value under the intersection hypothesis (given be sime's method)
	twomin <- 2*((p1_1<p3_1)*p1_1 + (1-(p1_1<p3_1))*p3_1)
	maxp <- (p1_1>p3_1)*p1_1 + (1-(p1_1>p3_1))*p3_1
	p13_1 <- (twomin<maxp)*twomin + (1-(twomin<maxp))*maxp
	#and the Z-values
	Z1_1 <- (thetahat1_1-delta[1])*sqrt(lambda*info[1])
	Z3_1 <- (thetahat3_1-delta[2])*sqrt(info[1])
	Z13_1<- qnorm(1-p13_1)

#find the interim decision, return vector d where 1 = subpop, 3 = full pop
#plan simulation for stage 2 accordingly
	d <- decis(thetahat1_1,thetahat2_1,
		alpha,lambda,tau,info,delta,m,dparm)

	n1 <- sum(d==1)
	n3 <- sum(d==3)

#second stage of the trial
	#set up vectors to populate for the stage 2 Z-values
	#-Inf would guarantee rejection if data are not simulated
	thetahat1_2 <- rep(-Inf,n) 
	thetahat2_2 <- rep(-Inf,n)
	#simulate the Enriched trials
	thetahat1_2[which(d==1)] <- rnorm(n1,theta1[which(d==1)],
							sqrt(1/info[2]))
	#simulate the full trials
	thetahat1_2[which(d==3)] <- rnorm(n3,theta1[which(d==3)],
							sqrt(1/(lambda*info[2])))
	thetahat2_2[which(d==3)] <- rnorm(n3,theta2[which(d==3)],
							sqrt(1/((1-lambda)*info[2])))
	#find the corresponding full pop estimates
	thetahat3_2 <- lambda*thetahat1_2 + (1-lambda)*thetahat2_2
	#find corresponding p-values
	p1_2 <- 1-pnorm(thetahat1_2,mean=delta[1],sd=sqrt(1/(lambda*info[2])))
	p3_2 <- 1-pnorm(thetahat3_2,mean=delta[2],sd=sqrt(1/(info[2])))
	#the p-value under the intersection hypothesis (given be sime's method)
	twomin <- 2*((p1_2<p3_2)*p1_2 + (1-(p1_2<p3_2))*p3_2)
	maxp <- (p1_2>p3_2)*p1_2 + (1-(p1_2>p3_2))*p3_2
	p13_2 <- (twomin<maxp)*twomin + (1-(twomin<maxp))*maxp
	#note the different treatment of p-values for enriched trial
	p1_2[which(d==1)] <- 1-pnorm(thetahat1_2[which(d==1)],
						mean=delta[1],sd=sqrt(1/info[2]))
	p3_2[which(d==1)] <- 1
	p13_2[which(d==1)] <- p1_2[which(d==1)]
	#and the Z-values

	Z1_2 <- (thetahat1_2-delta[1])*sqrt(lambda*info[2])
	Z1_2[which(d==1)]  <- (thetahat1_2[which(d==1)]-delta[1])*sqrt(info[2])
	Z3_2 <- (thetahat3_2-delta[2])*sqrt(info[2])
	Z13_2<- qnorm(1-p13_2)
	
	#weights for hypothesis testing
	w1 <- sqrt(tau)
	w2 <- sqrt(1-tau)
	#find combined Z-values
	Z1_c <- w1*Z1_1 + w2*Z1_2
	Z3_c <- w1*Z3_1 + w2*Z3_2
	Z13_c <- w1*Z13_1 + w2*Z13_2

	#now check for rejection
	#rejection boundary for Z
	k <- qnorm(1-alpha)
	
	prob1 <- mean((Z1_c > k)*(Z3_c < k)*(Z13_c > k))
	prob3 <- mean((Z1_c < k)*(Z3_c > k)*(Z13_c > k))
	probb <- mean((Z1_c > k)*(Z3_c > k)*(Z13_c > k))
	probe <- n1/n
	gain  <- mean(lambda*theta1*(Z1_c > k)*(Z3_c < k)*(Z13_c > k)+
		   (lambda*theta1+(1-lambda)*theta2)*(Z3_c > k)*(Z13_c > k))

	c(prob1,prob3,probb,probe,gain)
}

################################################################################
##dgrid
################################################################################
##find the appropriate decision from a decision grid for estimates at the 
##interim analysis
#Inputs: thetahat1 = estimate of theta1 from first stage of the trial
#   thetahat2 = estimate of theta2 from first stage of the trial
#      alpha  = FWER
#			 lambda = proportion of the population in sub-pop
#			 tau    = proportion of the sample in the first stage
#			 info   = ratio of variance and sample size
#			 n      = number of simulations, single value
#			 delta  = hypothesis test H_0:theta <= delta vector
#				    (note this is sub and full population)
#			 theta  = vector of theta values to simulate from
#				    (note supply a single value for true parameter)
#			 decis  = function to give the interim decision
#			 dparm  = list of decision function parameters
#Outputs: outcomes = 3 if trial should continue in full population
#                    1 if trial should continue in sub-population
dgrid <- function(thetahat1,thetahat2,
		alpha,lambda,tau,info,delta,m,dparm){
	#unpack the decision table
	t1vec <- dparm$rows
	t2vec <- dparm$cols
	grid  <- dparm$grid

	#find where the thetahat lie within the decision grid
	xlow <- floor(((thetahat1-min(t1vec))/(max(t1vec)-min(t1vec)))*length(t1vec))
	xup  <- floor(((thetahat1-min(t1vec))/(max(t1vec)-min(t1vec)))*length(t1vec)) + 1
	ylow <- floor(((thetahat2-min(t2vec))/(max(t2vec)-min(t2vec)))*length(t2vec))
	yup  <- floor(((thetahat2-min(t2vec))/(max(t2vec)-min(t2vec)))*length(t2vec)) + 1

	xlow[which(xlow<1)] <- 1
	ylow[which(ylow<1)] <- 1
	xup[which(xup<1)] <- 1
	yup[which(yup<1)] <- 1

	xlow[which(xlow>length(t1vec))] <- length(t1vec)
	ylow[which(ylow>length(t2vec))] <- length(t2vec)
	xup[which(xup>length(t1vec))] <- length(t1vec)
	yup[which(yup>length(t2vec))] <- length(t2vec)

	#the four closest vertecies
	x <- c(xlow,xup)
	y <- c(ylow,yup)
	v <- expand.grid(x,y)

	dist  <- sqrt((thetahat1 - v[,1])^2 + (thetahat2 - v[,2])^2)
	decis1 <- as.numeric(v[which.min(dist),][1])
	decis2 <- as.numeric(v[which.min(dist),][2])

	outcomes <- grid[decis1,decis2]
	outcomes[which(outcomes==-2)] <- 3
	outcomes[which(outcomes==2)] <- 1

	outcomes
}

#Vectorize is used to allow the arguements thetahat1 and thetahat2 to be input
#as vectors, the resulting output will be a 2 x length(thetahat1) matrix of 
#Bayes expected gain
vdgrid <- Vectorize(dgrid,vectorize.args=c("thetahat1","thetahat2"))