#################################################################################################
# This program conducts a robust mediation analysis for single mediation model based on median regression
# To use the program, you first need to download and install the package "quantreg" from 
# http://www.r-project.org, and import MMediation into R using the command
#		source("C:/mmediation.r")
# where C:/ should be replaced with the address of the mmediation.r file which will vary by computer.  
# Then run the program use the command 
#		robustmed(y, m, x, n, nboot)
# where
#        x: indepdent variable; 
#        y: dependent variable;
#        m: mediating variable; 
#        boostrap: indicating whether the bootrap 95% confidence interval (CI) should be calculated. 
#                  By default, the bootrap CI will be calculated, which may take a couple of minutes. 
#                  Set bootstrap=FALSE to skip the calculation;           
#        nboot: number of bootstrap samples
#
#################################################################################################

# load quantile regression library. 
library(quantreg);


# The function to conduct the standard and robust mediation analyses
robustmed<- function(x, y, m, bootstrap=TRUE, nboot=1000)
{
	
	# subroutine to calculate 95% CI using the bootstrap method
	bootstrapCI<-function(x, y, m, nboot)
	{
		set.seed(12);
		Theta.hat = numeric(nboot);
		Theta.hat.q = numeric(nboot);
		for(i in 1:nboot)
		{
			resamp.ind<-sample(1:length(y), replace=T);
			y.b = y[resamp.ind];
			m.b = m[resamp.ind];
			x.b = x[resamp.ind];
			alpha = lm(m.b~x.b)$coeff[2];
			beta  = lm(y.b~m.b+x.b)$coeff[2];
			alpha2 = rq(m.b~x.b, tau=0.5)$coeff[2];
			beta2  = rq(y.b~m.b+x.b, tau=0.5)$coeff[2];
			Theta.hat[i] = alpha*beta;
			Theta.hat.q[i] = alpha2*beta2;
		}
		return(c(quantile(Theta.hat, c(0.025, 0.975)), quantile(Theta.hat.q, c(0.025, 0.975))));	
	}	
	
	##### the standard mediation analysis based on the OLS mean regression 
	# fit model m = beta02+alpha*x+e2
	mx.reg = lm(m~x);
	alpha = mx.reg$coeff[2];
	alpha.var = vcov(mx.reg)["x", "x"];
	# robust variance estimate
	N = length(y);
	xx.inv = matrix(ncol=2, nrow=2);
	X = cbind(rep(1, N), x);
	xbar = mean(x);
	ssx = sum((x-xbar)^2);
	xx.inv[1,1] = sum(x^2)/(N*ssx);
	xx.inv[1,2] = -xbar/ssx;
	xx.inv[2,1] = xx.inv[1,2]
	xx.inv[2,2] = 1/ssx;
	HC3 = xx.inv%*%t(X)%*%diag((resid(mx.reg)/(1-lm.influence(mx.reg)$hat))^2)%*%X%*%xx.inv;
	alpha.var.hc3 = HC3[2, 2];

	
	# fit model y = beta03+beta*m+tau.prime*x+e3
	ymx.reg  = lm(y~m+x);
	beta = ymx.reg$coeff[2];
	taup = ymx.reg$coeff[3];
	beta.var = vcov(ymx.reg)["m", "m"];
	# robust variance estimate
	X = cbind(rep(1,N), m, x)
	XX.inv = chol2inv(chol(t(X)%*%X));
	HC3 = XX.inv%*%t(X)%*%diag((resid(ymx.reg)/(1-lm.influence(ymx.reg)$hat))^2)%*%X%*%XX.inv;
	beta.var.hc3 = HC3[2, 2];	
	
	# inference	
	theta.hat = alpha*beta;
	ci.product = quantile(rnorm(100000, alpha, sqrt(alpha.var))*rnorm(100000, beta, sqrt(beta.var)),  c(0.025, 0.975));
	ci.product.hc3 = quantile(rnorm(100000, alpha, sqrt(alpha.var.hc3))*rnorm(100000, beta, sqrt(beta.var.hc3)),  c(0.025, 0.975));
	
	
	##### a robust mediation analysis based on median regression
	# fit model m = beta02+alpha*x+e2
	mx.qreg = rq(m~x, tau=0.5);
	alpha = mx.qreg$coeff[2];
	alpha.se = summary(mx.qreg, se="ker")$coeff[2,2];
	
	# fit model y = beta03+beta*m+tau.prime*x+e3 
	ymx.qreg  = rq(y~m+x, tau=0.5);
	beta = ymx.qreg$coeff[2];
	taup = ymx.qreg$coeff[3];
	beta.se = summary(ymx.qreg, se="ker")$coeff[2,2];
	
	# inference
	theta.hat.q = alpha*beta;
	ci.product.q = quantile(rnorm(100000, alpha, alpha.se)*rnorm(100000, beta, beta.se),  c(0.025, 0.975));
	
	
	#### calculate 95% CI based on bootstrap
	if(bootstrap==TRUE)
	{
		bs = bootstrapCI(x, y, m, nboot)
		ci.boot = bs[1:2];
		ci.boot.q = bs[3:4];
	}
	cat("************************************************************************\n");
	cat("* Standard mediation analysis based on the ordinary least square (OLS) *\n");
	cat("************************************************************************\n");
	cat(" estimate of mediation effect = ", theta.hat);
	cat("\n 95% CI using OLS product of coefficients method: (", ci.product[1], ", ", ci.product[2],")");
	cat("\n 95% CI using HC3 product of coefficients method: (", ci.product.hc3[1], ", ", ci.product.hc3[2],")");
	if(bootstrap)
	{cat("\n 95% CI using bootstrap method: (", ci.boot[1], ", ", ci.boot[2],")");}
	cat("\n\n\n");
	
	cat("********************************************************\n");
	cat("* Robust mediation analysis based on median regression *\n");
	cat("********************************************************\n");
	cat(" estimate of mediation effect = ", theta.hat.q);
	cat("\n 95% CI using product of coefficients method: (", ci.product.q[1], ", ", ci.product.q[2],")");
	if(bootstrap)
	{cat("\n 95% CI using bootstrap method: (", ci.boot.q[1], ", ", ci.boot.q[2],")");}
	cat("\n");
}