######################################################################################### 
# Hardy Weinberg Equilibrium Test in structured populations (HWES). x (a matrix)	# 
# is the genotypes at the testing SNPs. Each row of x represents the genotypes 		#
# of one individual. G (a matrix) is the genotypes at a set of unlinked markers 	#
# used to control population stratification. Each row of G represents the genot-	#
# ypes of one individual at a set of unlinked markers. K represents the number of	#
# PCs used to control population stratification.					#
#########################################################################################

HWES=function(x,G,K)
{
	N=nrow(x)
	M=ncol(x)
	L=ncol(G)
	GG=G
	for(i in 1:L)	GG[,i]=G[,i]-mean(G[,i])
	if(N>=L)
		{
		A=t(GG)%*%GG
		B=eigen(A)
		E=B$vectors
		T=GG%*%E[,1:K]
		}
	else
		{
		GGG=GG%*%t(GG)
		BB=eigen(GGG)
		EE=BB$vectors
		T=EE[,1:K]	
		}
	pvalue=c(1:M)
	for(i in 1:M)
		{
		yy1=(x[,i]==1)
		yy2=(x[,i]==2)
		aa1=glm(yy1~T,family=binomial)
		aa2=glm(yy2~T,family=binomial)
		pp=aa1$fitted.values/2+aa2$fitted.values	
		pvalue[i]=testhwnew(x[,i],pp)
		}
pvalue
}



##############################

testhwnew=function(x,x1)
{
	N=length(x)
	n=rep(0,3)
	for(i in 1:2) n[i]=sum(x==(i-1))
	n[3]=N-n[1]-n[2]
	m=rep(0,3)
	m[3]=sum(x1*x1)
	m[2]=2*sum(x1*(1-x1))
	m[1]=sum((1-x1)*(1-x1))
	t1=sum((n-m)^2/m)

pvalue=1-pchisq(t1,1)
}