# ===========================================================================
# Refereces
# An Adaptive Association Test for Multiple Phenotypes with GWAS Summary Statistics
# Genetic Epidemiology Volume 39, Issue 8, pages 651--663
# Junghi Kim, Yun Bai, and Wei Pan
# ===========================================================================
library(matrixStats)
library(mvtnorm)
library(MASS)



estcov <- function(allZ){
  # -- Authors: Junghi Kim, Yun Bai, Wei Pan
  # -- allZ: matrix of summary Z-scores for all SNP
  # -- each row for SNP; each column for single trait  
  # -- return: estimated correlation matrix 
  n.snp <- dim(allZ)[1]
  n.traits <- dim(allZ)[2]
  minZ = rowMins(2*pnorm(-abs(allZ)))
  nullSNPs = which(minZ>(0.05/(n.snp/n.traits)))
  return( cor(allZ[nullSNPs,]))
}



minP <- function(Zi, v){
	# -- Zi: a vector of summary Z-scores for single SNP
	# -- v: estimated covariance matrix based on the summary Z-scores
	# -- return exact minP test
    n <- dim(v)[1]
	x <- as.numeric(max(abs(Zi)))
	return(as.numeric(1 - pmvnorm(lower=c(rep(-x,n)), upper=c(rep(x,n)), mean=c(rep(0,n)), sigma=v)))
}







MTaSPUs <- function(Z, v, B, pow, transform = FALSE){
   # -- Z: matrix of summary Z-scores, SNPs in rows and traits in columns  
   # -- Or a vector of summary Z-scores for a single snp
   # -- v: output of estcov
   # --    estimated estimated correlation matrix  based on the summary Z-scores
   # -- tranform: if TRUE, the inference is made on transformed Z
   # -- B: number of Monte Carlo samples simulated to compute p-values max(B)=1e8
   # -- results: compute p-values for SPU(gamma) i.e. pow=1:8, and infinity
   # --          aSPU, based on the minimum p-values over SPU(power)
   # --          each row for single SNP    
	if (B < 1e8) p <- MTaSPUsmallB(Z, v, B, pow, transform) else  p <- MTaSPUsB1e8(Z, v, pow, transform)
	return(p)
}


	
MTaSPUsmallB <- function(Z, v, B, pow, transform = FALSE){
  # -- Z: matrix of summary Z-scores, SNPs in rows and traits in columns  
  # -- Or a vector of summary Z-scores for a single snp
  # -- v: output of estcov
  # --    estimated estimated correlation matrix  based on the summary Z-scores
  # -- tranform: if TRUE, the inference is made on transformed Z
  # -- B: number of Monte Carlo samples simulated to compute p-values 
  # -- results: compute p-values for SPU(gamma) i.e. pow=1:8, and infinity
  # --          aSPU, based on the minimum p-values over SPU(power)
  # --          each row for single SNP    
  
    
    if (transform){
        v <- ginv(v)
        Z <- tcrossprod(Z, v)
    }
    
    if (is.vector(Z)) {
        N <- 1; K <- length(Z)
    }	else {
        N <- dim(Z)[1]; K <- dim(Z)[2]
    }
 	 if (N ==1) Z <- matrix(Z, nrow=1)
	 dim(Z) = c(N, K)
    pval <- matrix(0, N, length(pow))
	 dim(pval) = c(N, length(pow))
    aSPU <- length(N)
    ponum <- pow[pow < Inf]
	 
    set.seed(1000)
    Z0 <- rmvnorm(B, mean = rep(0, nrow(v)), sigma = v)
	 ## SPU for integer power
    for(k0 in 1:length(ponum)){
        k <- ponum[k0]  
        z1 <- abs(rowSums(Z^k))
        z0b <- abs(rowSums(Z0^k))
        for(i in 1:N){
            pval[i,k0] <- (1+sum(z0b>z1[i]))/(B+1)
        }
    }
  
    ## SPU(max)
    if (Inf %in% pow){
        z1 <- rowMaxs(abs(Z))
        z0 <- rowMaxs(abs(Z0))
        for(i in 1:N){
			  pval[i,length(pow)] = (1+sum(z0>z1[i]))/(B+1)
		  }
    }
    
    ## aSPU
    p1m <- rowMins(pval)
    p0 <- matrix(NA, B, length(pow))
    for(k0 in 1:length(ponum)){
        k <- ponum[k0]
        zb <- abs(rowSums(Z0^k))
        p0[,k0] <- (1+B-rank(abs(zb)))/B
    }

    
    if (Inf %in% pow){
      zb <- rowMaxs(abs(Z0))
      p0[,length(pow)] = (1+B-rank(zb))/B
    }
    p0m = rowMins(p0)
    for(i in 1:N){
        aSPU[i] = (1+sum(p0m<p1m[i]))/(B+1)
    }
	 pval <- cbind(pval, aSPU)
    if(Inf %in% pow) s <- c(paste("SPU", ponum, sep=""),"SPUInf","MTaSPUs") else {
      s <- c(paste("SPU", ponum, sep=""),"MTaSPUs")}
    colnames(pval) <- s
    rownames(pval) <- rownames(Z)
    return(pval)
}











MTaSPUsB1e8 <- function(Z, v, pow, transform = FALSE){
   # -- B: number of Monte Carlo simulation is fixed at B=1e8
   # -- Z: matrix of summary Z-scores, SNPs in rows and traits in columns  
   # -- Or a vector of summary Z-scores for a single snp
   # -- v: output of estcov
   # --    estimated estimated correlation matrix  based on the summary Z-scores
   # -- tranform: if TRUE, the inference is made on transformed Z
   # -- results: compute p-values for SPU(gamma) i.e. pow=1:8, and infinity
   # --          aSPU, based on the minimum p-values over SPU(power)
   # --          each row for single SNP    
  
    
   B<- 1e7; B2<- 10
	if (transform){
	    v <- ginv(v)
	    Z <- tcrossprod(Z, v)
	}
    
	if (is.vector(Z)) {
		N <- 1; K <- length(Z)
	} else {
		N <- dim(Z)[1]; K <- dim(Z)[2]
	}
	if (N ==1) Z <- matrix(Z, nrow=1)
	ponum <- pow[pow < Inf]
	Zpu = apply(Z, 1, function(zj){
	  tmp = rep(NA,length(pow))
	  for(k0 in 1:length(ponum)){
	    tmp[k0]<- abs(sum(zj^ponum[k0]))
	  }
	  if (Inf %in% pow){tmp[length(pow)] = max(abs(zj))}
	  return(tmp)
	})

	P1perm <- matrix(0, N, length(pow))
	Qq <- matrix(0, 50, length(pow))
	for(b in 1:B2){
	  zb <- rmvnorm(B, sigma = v)
	  for(k0 in 1:length(ponum)){
		 k <- ponum[k0]  
	    res = abs(rowSums(zb^k))
	    for(j in 1:N){
	      P1perm[j,k0] = P1perm[j,k0]+mean(res>Zpu[k0,j])/B2
	    }
	    res = c(res, Qq[,k0])
	    Qq[,k0] = -sort.int(-res,partial=50)[1:50]
	  }
	  if (Inf %in% pow){
		 res = rowMaxs(abs(zb))
		 for(i in 1:N){
		 	P1perm[i,length(pow)] = P1perm[i,length(pow)]+mean(res>Zpu[length(pow),i])/B2
		 }
		 res = c(res, Qq[,length(pow)])
		 Qq[,length(pow)] = -sort.int(-res,partial=50)[1:50]
	  }
	}


	## aSPU
	Qq = apply(Qq, 2, sort)
	Q0 = rbind(0,Qq)
	pr0 = c(1,50:1/(B*B2))

	pvi = matrix(NA, B,length(pow))
	Pmin = rowMins(P1perm)
	aSPU = numeric(N)
	for(b in 1:B2){
	  zb = rmvnorm(B, sigma = v)
	  for(k0 in 1:length(ponum)){
		 k <- ponum[k0]  
	    pvi[,k0] = approx(Q0[,k0], pr0, abs(rowSums(zb^k)), method='linear',rule=2)$y
	  }
	  if (Inf %in% pow){
	  	pvi[,length(pow)] = approx(Q0[,length(pow)], pr0, rowMaxs(abs(zb)), method='linear',rule=2)$y
	  }
  	  tmp = rowMins(pvi)
	  for(i in 1:N){
	    aSPU[i] = aSPU[i] + mean(tmp < Pmin[i])/B2
	  }
	}
	P1perm = cbind(P1perm, aSPU)
	P1perm[P1perm==0] = 1/(B*B2)
   if(Inf %in% pow) s <- c(paste("SPU", ponum, sep=""),"SPUInf","MTaSPUs") else {
      s <- c(paste("SPU", ponum, sep=""),"MTaSPUs")}
	colnames(P1perm) <- s
   rownames(P1perm) <- rownames(Z)
	return(P1perm)
}