#Program for arbitrary n multistage designs for disease model parameters 

#It is necessary to install "mvtnorm" package before you run this program.

#Arbitrary n-stage design with pisamples are equally divided into each stage and pimarkers are set to equal proportions to adjust the product of them to a ratio (here, the ratio is 0.0001)

     #Input parameters
        Dim <- 5                             # number of stages
	MDL <- 1                             # genotype model 1:multiplicative model, 2:additive model, 3:dominant model, 4:recessive model
	alphagenome <- 0.05                  # genome-wide false positive rate
	N <- 1000                            # number of samples in cases and control
	pisamples1 <- 1/(Dim)                # proportion of the sample size at 1st stage. Here the following stages' pisamples are equally divided into each stage. When each pisamplesk is variously different, don't use this command.  
        pimarkers1 <- (0.0001)^(1/(Dim-1))   # proportion of loci to be selected at 1st stage. Here the following stages' pimarkers are set to equal proportions to adjust the product of them to a ratio (here, the ratio is 0.0001). When each pimarkersk is variously different, don't use this command. 
	#pisamples<-matrix(1:(Dim-1),1,(Dim-1))  # When each pisamplesk is variously different, use this command.  
	#pimarkers<-matrix(1:(Dim-1),1,(Dim-1))  # When each pimarkersk is variously different, use this command. 
	#pisamples[,1] <- 1/(Dim)                # proportion of the sample size at 1st stage. When each pisamplesk is variously different, use this command. 
	#pisamples[,2] <- 1/(Dim) 	         # proportion of the sample size at 2nd stage. When each pisamplesk is variously different, use this command. 
	#pisamples[,3] <- 1/(Dim)                # proportion of the sample size at 3rd stage. When each pisamplesk is variously different, use this command. 
	#pisamples[,4] <- 1/(Dim) 	         # proportion of the sample size at 4th stage. When each pisamplesk is variously different, use this command. 
	#to be continued as well as the above-mentioned "Dim-1" times in all, in the following lines. 

	#pimarkers[,1] <- (0.0001)^(1/(Dim-1))   # proportion of loci to be selected at 1st stage. When each pimarkersk is variously different, use this command. 
	#pimarkers[,2] <- (0.0001)^(1/(Dim-1))   # proportion of loci to be selected at 2nd stage. When each pimarkersk is variously different, use this command. 
	#pimarkers[,3] <- (0.0001)^(1/(Dim-1))   # proportion of loci to be selected at 3rd stage. When each pimarkersk is variously different, use this command. 
	#pimarkers[,4] <- (0.0001)^(1/(Dim-1))   # proportion of loci to be selected at 4th stage. When each pimarkersk is variously different, use this command. 
	#to be continued as well as the above-mentioned "Dim-1" times in all, in the following lines.      
	
	M <- 500000                          # number of candidate loci
	prev <-	0.1                          # population prevalence of disease
	freq <-	0.3                          # population disease allele rate
	grr <- 1.5                           # population genotype relative risk(GRR)
	tm <- 5                              # number of true loci in M candidate loci 
        
                #m
                invisible(switch(MDL,
                (m <- 2),
                (m <- 1),
                (m <- 1),
                (m <- 1)))

		#penetration
                invisible(switch(MDL,
                (f0<-prev/(freq^2*grr^m+2*freq*(1-freq)*grr+(1-freq)^2)) & (f1<-grr*f0) & (f2<-grr^2*f0),
                (f0<-prev/(freq^2*(2*grr-1)+2*freq*(1-freq)*grr^m+(1-freq)^2)) & (f1<-grr*f0) & (f2<-(f1-f0)+f1), 
                (f0<-prev/(freq^2*grr^m+2*freq*(1-freq)*grr+(1-freq)^2)) & (f1<-grr*f0) & (f2<-grr^1*f0),
                (f0<-prev/(freq^2*grr^m+2*freq*(1-freq)+(1-freq)^2)) & (f1<-prev/(freq^2*grr^1+2*freq*(1-freq)+(1-freq)^2)) & (f2<-grr^m*f0)))
	     
		#Hardy-Weinberg's equilibrium
		HW0<-(1-freq)^2
		HW1<-2*freq*(1-freq)
		HW2<-freq^2

		#Frequency by each genotype
		pcase2<-HW2*f2/prev
		pcont2<-HW2*(1-f2)/(1-prev)
		pcase1<-HW1*f1/prev
		pcont1<-HW1*(1-f1)/(1-prev)
		pcase0<-HW0*f0/prev
		pcont0<-HW0*(1-f0)/(1-prev)

		invisible(switch(MDL,
		(pcase<-pcase2+pcase1/2) & (pcont<-pcont2+pcont1/2),
		(pcase<-pcase2+pcase1/2) & (pcont<-pcont2+pcont1/2),
		(pcase<-pcase2+pcase1/2) & (pcont<-pcont2+pcont1/2),
		(pcase<-pcase1/2+pcase0) & (pcont<-pcont1/2+pcont0)))

	D<- Dim-1

	if(D==0){
                
		alphamarker<-alphagenome/M
		T<-2*N*M
		library(mvtnorm)
		c<-(qnorm(1-alphamarker/2,0,1))
		mu0<-abs((pcase-pcont)/((pcase*(1-pcase)+pcont*(1-pcont))/(2*N))^(1/2))
		p<-(1-pnorm(c,mu0,1)+pnorm(-c,mu0,1))
		p0<-(1-pnorm(c,0,1)+pnorm(-c,0,1))

		tp<-tm*p
		fp<-(M-tm)*p0
		ppv0<-tp/(tp+fp)
 	
		OneStageDesign <- c("Co:", "Power:", "PPV:", "Typings:" )
		Output <- c(sprintf("%1.3f", c),sprintf("%1.3f", p),sprintf("%1.3f", ppv0),sprintf("%15.0f", T))
        	(x <- data.frame(OneStageDesign, Output))            
        	        	
	}else{            
	
		pisamples<-matrix(1:D,1,D)                     # When each pisamplesk is variously different, don't use this command.  
		pimarkers<-matrix(1:D,1,D)                     # When each pimarkersk is variously different, don't use this command.
		for( i in 1:D){                                # When each pisamplesk and pimarkersk are variously different, don't use this command. 
			pisamples[,i] <- (1-pisamples1)*1/(D)  # When each pisamplesk is variously different, don't use this command. 
			pimarkers[,i] <- pimarkers1            # When each pimarkersk is variously different, don't use this command.
		}                                              # When each pisamplesk and pimarkersk are variously different, don't use this command. 


		alphamarker<-alphagenome/M
		T<-0
		for( i in 1:(D+1)){
			ink<- 1
			if (i==D+1){
				ix <- 1
				for( j in 1:D){
					ix <- ix-pisamples[,j]
				}
				for( j in 1:(i-1)){
					ink <- ink*pimarkers[,j]
				} 
			}else{
				ix <- pisamples[,i]
				if(i!=1){
					for( j in 1:(i-1)){
						ink <- ink*pimarkers[,j]
					}
				}
			}
			T <- T+2*N*ix*M*ink
		}
		library(mvtnorm)

		cc<-matrix(1:(D+2),1,(D+2))
		muu<-matrix(1:(D+2),1,(D+2))
		for( i in 1:(D+2)){
			if(i==1){
				cc[,1] <- (qnorm(1-alphamarker/2,0,1))
				muu[,1]<- abs((pcase-pcont)/((pcase*(1-pcase)+pcont*(1-pcont))/(2*N))^(1/2))	
			}else if(i==D+2){
				ix <- 1
				ink<- 1
				for( j in 1:D){
					ix <- ix-pisamples[,j]
					ink <- ink*pimarkers[,j]
				}
				cc[,i] <- qnorm(1-(alphamarker/ink),0,1)
				muu[,i]<-abs((pcase-pcont)/((pcase*(1-pcase)+pcont*(1-pcont))/(2*N*ix))^(1/2))
				crep<-qnorm(1-alphamarker/ink,0,1)
			}else{
				cc[,i] <- qnorm(1-(pimarkers[,(i-1)]/2),0,1)
				muu[,i]<-abs((pcase-pcont)/((pcase*(1-pcase)+pcont*(1-pcont))/(2*N*pisamples[,(i-1)]))^(1/2))
			}
		}

		pp<-matrix(1:(D+2),1,(D+2))
		pp0<-matrix(1:(D+2),1,(D+2))
		for( i in 1:(D+2)){
			if(i<=2){
				pp[,i] <- (1-pnorm(cc[,i],muu[,i],1)+pnorm(-cc[,i],muu[,i],1))
				pp0[,i] <- (1-pnorm(cc[,i],0,1)+pnorm(-(cc[,i]),0,1))
			}else{
				pink<-1
				ppink<-1
				pinkk<-1
				ppinkk<-1
				for( j in 2:i){
					pink <- pink*(1-pnorm(cc[,j],muu[,j],1))
					pinkk <- pinkk*(pnorm(-cc[,j],muu[,j],1))
					ppink <- ppink*(1-pnorm(cc[,j],0,1))
					ppinkk <- ppinkk*(pnorm(-cc[,j],0,1))
				} 
				pp[,i]<-pink/pp[,2]+pinkk/pp[,2]
				pp0[,i] <- ppink/pp0[,2]+ppinkk/pp0[,2]
			}
		}
		prep<-pp[,2]*pp[,D+2]


		#Joint
		me<-numeric(D+1)
		cor<-matrix(0,nrow=D+1,ncol=D+1)
		diag(cor)<-1
		for( i in 1:D){
			cor[D+1,i]=sqrt(pisamples[,i])
			cor[i,D+1]=sqrt(pisamples[,i])
		}

		#f standby
		f<-function(x){
			d1<-matrix(-Inf,,D+1)
			d1[,D+1]<- -Inf
			d2<-matrix(0,,D+1)
			for( i in 1:D){
				d2[,i]<- -cc[,(i+1)]
			}
			d2[,D+1]<- -x
			d3<-matrix(0,,D+1)
			for( i in 1:D){
				d3[,i]<- (cc[,(i+1)])
			}
			d3[,D+1]<- x
			d4<-matrix(Inf,,D+1)
			d4[,D+1]<- Inf
			d5<-matrix(0,,D+1)
			for( i in 1:D){
				d5[,i]<- cc[,(i+1)]
			}
			d5[,D+1]<- -Inf
			d6<-matrix(Inf,,D+1)
			d6[,D+1]<- -x
			d7<-matrix(-Inf,,D+1)
			d7[,D+1]<- x
			d8<-matrix(0,,D+1)
			for( i in 1:D){
				d8[,i]<- -cc[,(i+1)]
			}
			d8[,D+1]<- Inf
		return((pmvnorm(lower=c(d1),upper=c(d2),mean=me,corr=cor)+
			pmvnorm(lower=c(d3),upper=c(d4),mean=me,corr=cor)+
			pmvnorm(lower=c(d5),upper=c(d6),mean=me,corr=cor)+
			pmvnorm(lower=c(d7),upper=c(d8),mean=me,corr=cor))-(alphamarker))}

		cjj<-uniroot(f,c(crep,cc[,1]),tol=1e-10)$root

		mujj<-matrix(0,,D+1)
		for( i in 0:D+1){
			if(i==D+1){
				ix <- 1
				for( j in 1:D){
					ix <- ix-pisamples[,j]
				}
				mujj[,i] <- muu[,i+1]/sqrt(ix)
			}else{
				mujj[,i] <- muu[,i+1]
			}
		}

		#pjj standby
			d1<-matrix(-Inf,,D+1)
			d1[,D+1]<- -Inf
                        d2<-matrix(0,,D+1)
			for( i in 1:D){
				d2[,i]<- -cc[,(i+1)]
			}
			d2[,D+1]<- -cjj
			d3<-matrix(0,,D+1)
			for( i in 1:D){
				d3[,i]<- (cc[,(i+1)])
			}
			d3[,D+1]<- cjj
			d4<-matrix(Inf,,D+1)
			d4[,D+1]<- Inf
			d5<-matrix(0,,D+1)
			for( i in 1:D){
				d5[,i]<- cc[,(i+1)]
			}
			d5[,D+1]<- -Inf
			d6<-matrix(Inf,,D+1)
			d6[,D+1]<- -cjj
			d7<-matrix(-Inf,,D+1)
			d7[,D+1]<- cjj
			d8<-matrix(0,,D+1)
			for( i in 1:D){
				d8[,i]<- -cc[,(i+1)]
			}
			d8[,D+1]<- Inf
		pjj<-(pmvnorm(lower=c(d1),upper=c(d2),mean=c(mujj),corr=cor)+
			pmvnorm(lower=c(d3),upper=c(d4),mean=c(mujj),corr=cor)+
			pmvnorm(lower=c(d5),upper=c(d6),mean=c(mujj),corr=cor)+
			pmvnorm(lower=c(d7),upper=c(d8),mean=c(mujj),corr=cor))

		tp<-tm*pp[,1]
		fp<-(M-tm)*pp0[,1]
		tprn<-tm*prep
		fprn<-(M-tm)*pp0[,2]*pp0[,D+2]
		tpjn<-tm*pjj
		fpjn<-(M-tm)*alphamarker
		ppv0<-tp/(tp+fp)
		ppvrep<-tprn/(tprn+fprn)
		ppvj<-tpjn/(tpjn+fpjn)

		nStageDesign <- c("C of RBA:","C of JA:", "Power of RBA:", "Power of JA:", "PPV of RBA:", "PPV of JA:", "Typings:" )
		Output <- c(sprintf("%1.3f", crep), sprintf("%1.3f", cjj), sprintf("%1.3f", prep), sprintf("%1.3f", pjj), sprintf("%1.3f", ppvrep), sprintf("%1.3f", ppvj), sprintf("%15.0f", T))
		(x <- data.frame(nStageDesign, Output))	

		}

		remove(list=ls())