#####################################################
## GKW.test is an R function that implements 
## the generalized Kruskal-Wallis (GKW) and 
## best-guess Kruskal-Wallis (BG-KW) tests.
#####################################################
# Version 1
# (c) 2011 Elif Acar and Lei Sun
#####################################################
GKW.test = function(x,p,...) UseMethod("GKW.test")
#####################################################
#####################################################
## x: numeric vector of data values
## p: matrix or data frame of group probabilities
## k: number of groups to be tested
## method: "GKW" (the Generalized KW test based on probability-weighted rank-sums) or 
## "BG-KW" (classify the individuals into their most probable group and perform the Kruskal-Wallis test) 
# see the end of the program for EXAMPLES
#####################################################

GKW.test.default= function(x,p,k, method="GKW"){

if(missing(k)){k= ncol(p)}
if(ncol(p) != k){stop("p must have k columns")}
if(any(round(rowSums(p),5) != 1) ){stop("group probabilities must add up to 1")}


N= length(x)
x.rank = rank(x,ties.method="average")
Nmethod = "Generalized Kruskal-Wallis test"
Ndata = paste(deparse(substitute(x)), "and", deparse(substitute(p)))

if(method=="BG-KW"){
Nmethod= "Best-guess Kruskal-Wallis test"
p_bg = matrix(0, nrow=N, ncol=k)
index = seq(0,(N-1)*k, by=k)+ max.col(p) 
p = t(replace(t(p_bg),index, 1))
	}

Rstar= colSums(p*x.rank)
mean.Rstar = (N+1)/2 * colSums(p)

ties = table(x)

if(any(ties > 1)){
	
T.term = sum( (ties-1)*(ties)*(ties+1) )/ (12*(N-1))
covar.Rstar = ( N*(N+1)/12 - T.term) *var(p)*(N-1)
covar.Rstar * (diag(covar.Rstar) %*% t(diag(covar.Rstar)))^(-0.5) 

print("tie-correction is used")

}else{
	covar.Rstar = N*(N+1)/12 * var(p)*(N-1)	
	cor.Rstar = covar.Rstar * (diag(covar.Rstar) %*% t(diag(covar.Rstar)))^(-0.5)   		
	}	
	

std.Rstar = (Rstar - mean.Rstar) / diag(covar.Rstar)^0.5
Hstar = t(std.Rstar)[1:(k-1)] %*% solve(cor.Rstar[1:(k-1),1:(k-1)]) %*% std.Rstar[1:(k-1)] 
df=k-1
pval = pchisq(Hstar,df,lower.tail=FALSE)	

if(any(colSums(p)<5)){print("Caution: chi-square approximation may not be reliable")}
names(Hstar) = "H*"
names(df) ="df"
Rval = list(
statistic = Hstar,
parameter = df,
p.value = pval,
method = Nmethod,
data.name = Ndata

)
class(Rval)="htest"
return(Rval)
}
#####################################################


#####################################################
## EXAMPLES
#####################################################

library(MCMCpack)


# EXAMPLE 1

x=rnorm(n=100)
p=rbind(rdirichlet(60,c(1,0,0)), rdirichlet(30,c(0,1,0)), rdirichlet(10,c(0,0,1))  )
GKW.test(x,p,method="BG-KW")
GKW.test(x,p, method="GKW")
kruskal.test(x,g=max.col(p))



# EXAMPLE 2 

x=rnorm(n=100)
p=rbind(rdirichlet(60,c(0.9,0.05,0.05)), rdirichlet(30,c(0.05,0.9,0.05)), rdirichlet(10,c(0.05,0.05,0.9))  )
GKW.test(x,p, method="GKW")
GKW.test(x,p,method="BG-KW")
# to check BG-KW with KW test
kruskal.test(x,g=max.col(p))
# KW test on actual groups
G= c(rep(1,60), rep(2,30),rep(3,10))
kruskal.test(x,G)