library(MASS)
library(emplik)


psi = function(t,tau){
	if(t>=0) return(tau) else return(tau-1)
	}

dSpikeSlab = function(x,theta,mu,sd){
	if(x==0) return(log(theta+1e-100))
	return(log(1-theta+1e-100)+dnorm(x,mean=mu,sd=sd,log=TRUE))
	}

## X should be a n by k matrix (no intercept)
## beta should be a k + length(tau) -vector 
## Y should be a n-vector
## mu and tau should be a m-vector
f_constrait = function(Y,X,beta,tau,wt=rep(1,length(tau))){
	GX = matrix(0,nrow=nrow(X),ncol=length(beta))
	beta1 = beta[1:(length(beta)-length(tau))]
	mu = beta[(length(beta1)+1):length(beta)]
	for(k in 1:length(tau)){
		t = Y - X%*%beta1 - mu[k]
		z = sapply(1:length(t),FUN=function(i){psi(t[i],tau[k])})
		GX[,1:length(beta1)] = GX[,1:length(beta1)] + wt[k]*z*X
		GX[,length(beta1)+k] = z
		}
	return(GX)
	}

### calculate the empirical likelihood
emplik = function(Y,X,beta,tau,wt=rep(1,length(tau))){
	tmp = el.cen.EM2(Y,d=rep(1,length(Y)),fun=function(YY){f_constrait(YY,X,beta,tau,wt)},mu=rep(0,ncol(X)+length(tau)))
	return(tmp$loglik)
	}

#### the MH proposing sampler
#### sigma2 is the variance of the normal prior
#### here beta include both beta and mu with its first k = length(beta)-length(tau) elements being beta
q_sample = function(i,Y,X,theta,beta,tau,sigma2,scale=1,wt=rep(1,length(tau))){
	mu_beta_i = 0.0
	sd_beta_i = 0.0
	sm.wt = sum(wt)
	beta1 = beta[1:(length(beta)-length(tau))]
	mu = beta[-c(1:length(beta1))]
	if(i<=length(beta)-length(tau)){
		for(k in 1:length(tau)){
			dta.tmp = data.frame(y=Y - X[,-i]%*%beta1[-i]-mu[k],x=X[,i])
			rq.model = rq(y~0+x,data=dta.tmp,tau=tau[k])
			rq.summary = summary(rq.model,se="boot")
			mu_beta_i = mu_beta_i + wt[k]*rq.summary$coef[1]/sm.wt
			sd_beta_i = sd_beta_i + wt[k]*rq.summary$coef[2]/sm.wt
			}
		}else{
		k = i - length(beta1)
		dta.tmp = data.frame(y=Y - X%*%beta1,x=rep(1,nrow(X)))
		rq.model = rq(y~0+x,data=dta.tmp,tau=tau[k])
		rq.summary = summary(rq.model,se="boot")
		mu_beta_i = rq.summary$coef[1]
		sd_beta_i = rq.summary$coef[2]
		}


	sigma2_new = (sigma2*scale*sd_beta_i^2)/(sigma2+scale*sd_beta_i^2)
	mu_new = sigma2_new*mu_beta_i/(scale*sd_beta_i^2)
	C0 = exp(-mu_beta_i^2/(2*scale*sd_beta_i^2))
	C1 = sqrt(sigma2_new/sigma2)*exp((sigma2_new/(scale*sd_beta_i^2)-1)*mu_beta_i^2/(2*scale*sd_beta_i^2))

	theta_new = theta[i]*C0/(theta[i]*C0+(1-theta[i])*C1)

	if(runif(1)<theta_new) {beta_i_new=0} else {beta_i_new=rnorm(1,mean=mu_new,sd=sqrt(sigma2_new))}
	if(beta_i_new==0) {density_new = log(theta_new)} else {density_new=log(1-theta_new)+dnorm(beta_i_new,mean=mu_new,sd=sqrt(sigma2_new),log=TRUE)}
	if(beta[i]==0) {density=log(theta_new)} else {density=log(1-theta_new)+dnorm(beta[i],mean=mu_new,sd=sqrt(sigma2_new),log=TRUE)}

	return(list(beta_new=beta_i_new,density_new=density_new,density=density))	
	}

#### now the samplers
sample_theta_i = function(i,beta,tau){
	if(i > length(beta)-length(tau)) return(0)
	if(beta[i]==0) return(rbeta(1,2,1)) else return(rbeta(1,1,2))
	}

sample_sigma2_inverse = function(beta,theta,shape0,rate0){
	shapeNew = sum(beta!=0)/2 + shape0 
	rateNew = sum(beta^2)/2 + rate0
	return(rgamma(1,shape=shapeNew,rate=rateNew))
	}

sample_beta_i = function(i,Y,X,beta,theta,tau,sigma2,scale=1,wt=rep(1,length(mu))){
	tmp = q_sample(i,Y,X,theta,beta,tau,sigma2,scale=scale,wt=wt)
	beta_new = beta
	beta_new[i] = tmp$beta_new
	alpha = emplik(Y,X,beta_new,tau,wt=wt) + tmp$density- (emplik(Y,X,beta,tau,wt=wt) + tmp$density_new)
	alpha = alpha + dSpikeSlab(beta_new[i],theta[i],0,sd=sqrt(sigma2)) - dSpikeSlab(beta[i],theta[i],0,sd=sqrt(sigma2))
	tmp = runif(1,0,1)
	if(tmp<exp(alpha)) return(beta_new[i]) else return(beta[i])
	}


#############
#### this is the major function, Y is the response, X is the design matrix, tau is the quantile vector
#### note that if tau only has one element, then  empirical likelihood quantile regression is performaed
#### if tau contains multiple elements, the bayesian multiple quantile regression is performed
QR_MCMC = function(Y,X,tau,nburn=1000,Nsample=10000,beta0=NULL,shape0=0.1,rate0=0.1,scale=1,wt=rep(1,length(tau))){
	if(Nsample-nburn<1000) Nsample=nburn+1000

	if(is.null(beta0)){
		if(length(tau)>1){
			rq.model = rq(y~.,data=data.frame(y=Y,x=X),tau=tau)	
			beta1 = rowMeans(rq.model$coef[-1,])
			mu1 = rq.model$coef[1,]
			beta0 = c(beta1,mu1)
			}else{
			rq.model = rq(y~.,data=data.frame(y=Y,x=X),tau=tau)
			beta1 = rq.model$coef[-1]
			mu1 = rq.model$coef[1]
			beta0 = c(beta1,mu1)
			}
		}

	num.pred = ncol(X)+length(tau)
	beta_chain = matrix(0,nrow=Nsample,ncol=num.pred)
	beta_chain[1,] = beta0
	sigma2_inverse_chain = rep(0,Nsample)
	theta_chain = matrix(0,nrow=Nsample,ncol=num.pred)

	for(k in c(1:Nsample)){
		### first sample theta
		for(i in 1:num.pred){
			theta_chain[k,i] = sample_theta_i(i,beta_chain[k,],tau)
			}
		### second sample sigma2_inverse
		sigma2_inverse_chain[k] = sample_sigma2_inverse(beta_chain[k,],theta_chain[k,],shape0,rate0)

		### Last sample beta
		beta_new = beta_chain[k,]
		for(i in 1:num.pred){
			beta_new[i] = sample_beta_i(i,Y,X,beta_new,theta_chain[k,],tau,1/sigma2_inverse_chain[k],scale=scale,wt=wt)
			}
		if(k<Nsample)  beta_chain[k+1,] = beta_new

		if(floor(k/100)==k/100) print(k)
		}
	

	beta.hat =  apply(beta_chain[-c(1:nburn), ], 2, median)
	return(list(beta.hat=beta.hat,betaChain=beta_chain, Sigma2Chain=sigma2_inverse_chain,thetaChain=theta_chain))
	}