# Bayesian Functional Data Clustering

rm(list=ls())
library(mvtnorm)
# Construct reproducing kernel
rk <- function(x, y) {
            k2 <- function(x) ((x - 0.5)^2 - 1/12)/2
            k4 <- function(x) ((x - 0.5)^4 - (x - 0.5)^2/2 + 7/240)/24
            k2(x) * k2(y) - k4(abs(x - y))
}
#  rk <- function(x, y) {
#   (min(x, y))^2 *( 3 * max(x, y) - min(x, y) )/6
#  }


nsub <- 50  # number of genes
nmeas <- 10  # number of time points/measurements/levels
nobs = nsub*nmeas  # nsub*nmeas
chip.time= (0:(nmeas-1))/(nmeas-1)


sigma <- 0.5
set.seed(6378)
#simulate a set of data

b <- rnorm(nsub,0,sigma)
e <- rnorm(nobs,0,0.5)
x <- rep(chip.time,nsub)
z <- matrix(0,nobs,nsub)
for(k in 1:nsub)
{
   z[(k*nmeas-nmeas+1):(k*nmeas),k] <- rep(1,nmeas)
}

z1=z[1:nmeas,1:nmeas]

z11=z1%*%t(z1)

y=NULL

y[1:(nobs/5)] <- 3*sin(2*pi*x[1:(nobs/5)])+(z%*%b)[1:(nobs/5)]+e[1:(nobs/5)]
y[(nobs/5+1):(2*nobs/5)]=3*sin(2*pi*x[(nobs/5+1):(2*nobs/5)]*3)+(z%*%b)[(nobs/5+1):(2*nobs/5)]+e[(nobs/5+1):(2*nobs/5)]
y[(2*nobs/5+1):(3*nobs/5)]=3*sin(2*pi*x[(2*nobs/5+1):(3*nobs/5)])*(1-x[(2*nobs/5+1):(3*nobs/5)]) +(z%*%b)[(2*nobs/5+1):(3*nobs/5)]+e[(2*nobs/5+1):(3*nobs/5)]
y[(3*nobs/5+1):(4*nobs/5)]=1980*x[(3*nobs/5+1):(4*nobs/5)]^7*(1-x[(3*nobs/5+1):(4*nobs/5)])^3+ 858*x[(3*nobs/5+1):(4*nobs/5)]^2*(1-x[(3*nobs/5+1):(4*nobs/5)])^10 +(z%*%b)[(3*nobs/5+1):(4*nobs/5)]+e[(3*nobs/5+1):(4*nobs/5)]
y[(4*nobs/5+1):nobs]=3*sin(2*pi*x[(4*nobs/5+1):nobs])-7 +(z%*%b)[(4*nobs/5+1):nobs]+e[(4*nobs/5+1):nobs]

y2=matrix(y,nmeas,nsub)
y2=t(y2)


# generate basis and kernel matrices

s <- matrix(1,nobs,2)
q <- matrix(0,nobs,nmeas)

s[,2] <- x-0.5

for(i in 1:nobs){for(j in 1:nmeas) q[i,j] <- rk(x[i],x[j])}

s.chip.time <- matrix(1,nmeas,2)
q.chip.time <- matrix(0,nmeas,nmeas)

s.chip.time[,2] <- chip.time-0.5

for(i in 1:nmeas){for(j in 1:nmeas) q.chip.time[i,j] <- rk(chip.time[i],chip.time[j])}


library(gss)
init <- 0.1
env <- nsub
fun <- function(zeta, env) diag(10^(-zeta), env)
sigma <- list(fun = fun, env = env)

random <- list(z=z,sigma=sigma,init=init)
class(random) <- "list"

my.est1=ssanova(y~x, random=random,  alpha=1, id.basis = 1:nmeas)

# Initialize the estimate of the parameters
library(MASS)



BICs=NULL

for(h in 4:6){
	b.est <- rep(0,nsub)
	c.est <- 10^my.est1$theta* my.est1$c[my.est1$c!=0]
	c.est <- c.est[1:nmeas]
	d.est <- my.est1$d

	sigmabsq.est <- 0.5
	tausq.est <- 0.01
	sigmasq.est <- 0.5

	delta0=10^4
	delta1=10^4


	# set the prior for the parameters.


	# The prior for tausq is assumed to be IG(tausq.alpha,tausq.beta). When shape
	#and scale parameters are small, we essentially put an noninformtive prior on
	#the variable

	tausq.alpha=0.01
	tausq.beta=0.01

	# The prior for sigmasq is assumed to be IG(sigmasq.alpha,sigmasq.beta)
	sigmasq.alpha=0.01
	sigmasq.beta=0.01

	# The prior for sigmabsq is assumed to be IG(sigmasqb.alpha,sigmasqb.beta)
	sigmabsq.alpha=0.01
	sigmabsq.beta=0.01

	# Initial values of probabilities of each cluster, and indice 


	index= kmeans(y2,h,iter.max = 20, nstart=5)$cluster
	p=c(rep(1/h,h))
	
	p1=c(rep(1/h,h))
	ns=c(rep(0,h))
	finalp<-matrix(0,h,nsub)

	# initial values of mus and Bs
	mus=matrix(0,nmeas,h)
	taus=rep(tausq.est,h)
	Bs=rep(sigmabsq.est,h)
	yhat=y

	# Initial values of b,c,d for different clusters
	cb.est=matrix(0,nsub,h)
	cc.est=matrix(0,nmeas,h)
	cd.est=matrix(0,2,h)
	for(iclust in 1:h){
		yclust=as.numeric(t(y2[which(index == iclust),]))
		xclust=rep(chip.time, length(which(index == iclust)))
		id.clust=as.factor(sort(rep( 1:length(which(index == iclust)), nmeas )))
#		xyclustpar <- list(nphi=1, mkphi=mkphi.cubic, mkrk=mkrk.cubic, env=c(min(xclust), max(xclust)) )
#		my.ss1=ssanova(yclust~xclust, random=~1|id.clust, alpha=1.0, id.basis=1:nmeas, type=list(xclust=list("custom",xyclustpar)) )
 		my.ss1=ssanova(yclust~xclust, random=~1|id.clust, alpha=1.0, id.basis=1:nmeas)
                predict(my.ss1, newdata=data.frame(xclust=chip.time))
		cc.est[,iclust] = 10^my.ss1$theta *my.ss1$c
		cd.est[,iclust] =  my.ss1$d
		mus[,iclust]= s.chip.time%*%my.ss1$d +  10^my.ss1$theta *q.chip.time%*%my.ss1$c
	}


	# Gibbs Sampler
	for (iteration in 1:5000){
		print(paste("iteration", iteration, "is running", sep=" ")) 
		for(i in 1:h){
			# doing the Gibbs sampler for the i-th cluster
			indexi=which(index==i)
			csub=length(indexi)
			if(csub>0){
				cobs=csub*nmeas
				index1=rep(0,cobs)

				for(jclust in 1:csub){
					index1[((jclust-1)*nmeas+1):(jclust*nmeas)]=5*rep(indexi[jclust]-1,nmeas)+1:nmeas
				}
				# select subjects from the i-th cluster
				si=s[index1,]
				yi=y[index1]
				qi=q[index1,]
				zi=z[index1,]

				# Sample fxied effects of the function
				d.var <- chol2inv(chol(diag(c(delta0^(-1),delta1^(-1)),2) +t(si)%*%si/sigmasq.est))
				d.mean <- d.var%*%t(si)%*%(yi-zi%*%cb.est[,i]-qi%*%cc.est[,i])/sigmasq.est
				d.est <- mvrnorm(1, mu=d.mean, Sigma=d.var)
	
				# Sample random effects coefficients
				b.var <- chol2inv(chol(diag(1/sigmabsq.est,nsub) + t(zi)%*%zi/sigmasq.est))
				b.mean <- b.var%*%t(zi)%*%(yi-si%*%d.est-qi%*%cc.est[,i])/sigmasq.est
				b.est <- mvrnorm(1, mu=b.mean, Sigma=b.var)

				# Sample random effects of the function
				c.var <- chol2inv(chol(diag(1/taus[i], nmeas) + t(qi)%*%qi/sigmasq.est))
				c.mean <- c.var%*%t(qi)%*%(yi-si%*%d.est-zi%*%b.est)/sigmasq.est
				c.est <- mvrnorm(1, mu=c.mean, Sigma=c.var)

				# Sample Wishart(v, S), simulate x1, ... xv from N(0,S), theta=sum(xixi^T)

				# Sample random effect variance sigmab^2

				Bs[i] <- 1/rgamma(1, shape=sigmabsq.alpha+nsub/2, rate=sigmabsq.beta+sum(b.est*b.est)/2)

				#Sample tau^2 (inverse of smoothing par)

				tausq.est <- 1/rgamma(1, shape=tausq.alpha+nmeas/2,
				rate=tausq.beta+sum(c.est*c.est)/2)

				yhat[index1]=si%*%d.est+qi%*%c.est+zi%*%b.est

				# update cb.est ,cc.est ,ed.est, taus, mus and Bs
				cb.est[,i]=b.est
				cc.est[,i]=c.est
				cd.est[,i]=d.est
				taus[i]=tausq.est
				mus[,i]=s[1:nmeas,]%*%d.est+q[1:nmeas,]%*%c.est
			}
		}


		# Sample random error variance sigma^2
		sigmabsq.est=min(Bs)
		rss <- t(y-yhat)%*%(y-yhat)
		sigmasq.est <- 1/rgamma(1, shape=sigmasq.alpha+nobs/2, rate=sigmasq.beta+rss/2)

		# Update probabilities of each subject belonging to different clusters
		# and then generate new cluster index
		sigma1=sigmabsq.est*z11+diag(sigmasq.est,nmeas)
		invsigma1=chol2inv(chol(sigma1))
		for(j in 1:nsub){

			for(icl in 1:h){
			        sigmasq.clust.est = Bs[icl]*z11+diag(sigmasq.est,nmeas)
			        p1[icl]= dmvnorm(y[(j*nmeas-nmeas+1):(j*nmeas)], mean=mus[,icl],sigma=sigmasq.clust.est, log=T) +log(p[icl])
				#p1[icl]=-t(y[(j*nmeas-nmeas+1):(j*nmeas)]-mus[,icl])%*%invsigma1%*%(y[(j*nmeas-nmeas+1):(j*nmeas)]-mus[,icl])/2+log(p[icl])
			}
		
			maxp=max(p1)
			p1=p1-maxp
			p1=exp(p1)

			p1=p1/sum(p1)
			finalp[,j]=p1
			index[j]=sample(seq(1,h,1),1,prob=p1)
		}

		# update probabilities of each cluster
		for(icm in 1:h){
			ns[icm]=length(index[index==icm])
			p[icm]=rgamma(1,ns[icm]+0.01,1)
		}

		p=p/sum(p)
	}
	BIC=(h*(nmeas+2+1)+nsub+2)*log(nsub)
	for(jgene in 1:nsub){
		temp0=0
		temp.x=as.numeric(y[(jgene*nmeas-nmeas+1):(jgene*nmeas)])

		for(kclust in 1:h){
			temp0=temp0+ p[kclust]*dmvnorm(temp.x, mean=mus[,kclust], sigma=sigma1)
		}
		BIC = BIC - 2*log(temp0)
	}


	BICs=c(BICs,BIC)
	print(index)
	print(finalp)
	print(mus)  

}

which.max=function(x)(which(x==max(x)))
apply(finalp,2, which.max)


BICs