######################
### Open libraries ###
######################
library(rmeta) # use install.packages(rmeta) first if it hasn't been installed

#############################
### Set working directory ###
#############################
setwd("T:/uqdi/HPC/PBSHOME/nwarrington/Students/VischerSummerSchool/") # Needs to be updated to where the data sits

#################
### Load data ###
#################
data <- read.table("Data_MR-EggerWM.txt", header=T, sep="\t", na="") 
dim(data) # 249 31
head(data)
summary(data)

############################
### Scatter plot of data ###
############################
plot(data$beta_bmi[!is.na(data$BMI)], data$beta_tg[!is.na(data$BMI)], xlab="SNP-BMI association (95% CI)", 
     ylab="SNP-Triglyceride association  (95% CI)", ylim=c(-0.051, 0.051), xlim=c(0, 0.09), pch=20)
for (i in 1:nrow(data[!is.na(data$BMI),])){
  LL_bmi <- data$beta_bmi[!is.na(data$BMI)][i] - 1.96*data$se_bmi[!is.na(data$BMI)][i]
  UL_bmi <- data$beta_bmi[!is.na(data$BMI)][i] + 1.96*data$se_bmi[!is.na(data$BMI)][i]
  LL_tg <- data$beta_tg[!is.na(data$BMI)][i] - 1.96*data$se_tg[!is.na(data$BMI)][i]
  UL_tg <- data$beta_tg[!is.na(data$BMI)][i] + 1.96*data$se_tg[!is.na(data$BMI)][i]
  lines(x=c(LL_bmi,UL_bmi), y=c(data$beta_tg[!is.na(data$BMI)][i],data$beta_tg[!is.na(data$BMI)][i]))
  lines(y=c(LL_tg,UL_tg), x=c(data$beta_bmi[!is.na(data$BMI)][i],data$beta_bmi[!is.na(data$BMI)][i]))
}
grid()


################################################
### Inverse Variance Weighted (IVW) Analysis ###
################################################
ivw.m <- meta.summaries(d=(beta_tg/beta_bmi), se=se_tg/beta_bmi, data=data, subset=!is.na(data$BMI), method="fixed")
summary(ivw.m)
ivw.m$summary
ivw.m$se.summary

ivw.r  <- lm(beta_tg ~ - 1 + beta_bmi, weights = (1 / (se_tg)^2), data=data, subset=!is.na(data$BMI))  
summary(ivw.r)
DF <- length(data$beta_tg[!is.na(data$BMI)])-1
IVWBeta <- summary(ivw.r)$coef[1,1]
SE <- summary(ivw.r)$coef[1,2]/summary(ivw.r)$sigma
IVW_p <- 2*(pt(abs(IVWBeta/SE),DF, lower.tail=F))
ivw.b <- data.frame(beta = IVWBeta, se = SE, p = IVW_p)

#######################################
### Calculate Cochran's Q statistic ###
#######################################
beta_iv <- data$beta_tg[!is.na(data$BMI)] / data$beta_bmi[!is.na(data$BMI)]
weight <- (data$beta_bmi[!is.na(data$BMI)])^2 / (data$se_tg[!is.na(data$BMI)])^2 
beta_ivw <- sum(weight * beta_iv)/sum(weight)
beta_ivw == ivw.m$summary

q <- sum(weight * (beta_iv - beta_ivw)^2)
pchisq(q, df=length(weight)-1, lower.tail=F)

################
### MR-Egger ###
################
egg.r  <- lm(beta_tg ~ beta_bmi, weights = (1 / (se_tg)^2), data=data, subset=!is.na(data$BMI))             
summary(egg.r)
egg.b  <- data.frame(beta = summary(egg.r)$coefficients[2,1], se = summary(egg.r)$coefficients[2,2] , 
                     pval = summary(egg.r)$coefficients[2,4])
egg.i  <- data.frame(beta = summary(egg.r)$coefficients[1,1], se = summary(egg.r)$coefficients[1,2] , 
                     pval = summary(egg.r)$coefficients[1,4])

#######################################
### Calculate Cochran's Q statistic ###
#######################################
qdash <- sum(weight *(beta_iv - (egg.i[1]/data$beta_bmi[!is.na(data$BMI)]) - egg.b[1])^2)
pchisq(qdash, df=length(weight)-2, lower.tail=F)

#################################
### Weighted median functions ###
#################################
weighted.median <- function(betaIV.in, weights.in) {
  betaIV.order = betaIV.in[order(betaIV.in)]
  weights.order = weights.in[order(betaIV.in)]
  weights.sum = cumsum(weights.order)-0.5*weights.order
  weights.sum = weights.sum/sum(weights.order)
  below = max(which(weights.sum<0.5))
  weighted.est = betaIV.order[below] + (betaIV.order[below+1]-betaIV.order[below])*
    (0.5-weights.sum[below])/(weights.sum[below+1]-weights.sum[below])
  return(weighted.est) }

weighted.median.boot = function(betaXG.in, betaYG.in, sebetaXG.in, sebetaYG.in, weights.in){
  med = NULL
  for(wmb in 1:1000){
    betaXG.boot = rnorm(length(betaXG.in), mean=betaXG.in, sd=sebetaXG.in)
    betaYG.boot = rnorm(length(betaYG.in), mean=betaYG.in, sd=sebetaYG.in)
    betaIV.boot = betaYG.boot/betaXG.boot
    med[wmb] = weighted.median(betaIV.boot, weights.in)
  }
  return(sd(med)) }

###################################
### Fit Weighted median to data ###
###################################
# Simple
b_exp <- data$beta_bmi[!is.na(data$BMI)]
se_exp <- data$se_bmi[!is.na(data$BMI)]
b_out <- data$beta_tg[!is.na(data$BMI)]
se_out <- data$se_tg[!is.na(data$BMI)]
b_iv <- c(b_out/b_exp)
weight_iv <- rep(1/sum(!is.na(data$BMI)), sum(!is.na(data$BMI))) # equal weights for all SNPs
b_sm = weighted.median(b_iv, weight_iv) 
se_sm = weighted.median.boot(b_exp, b_out, se_exp, se_out, weight_iv)
zscore_sm <- b_sm/se_sm
pval_sm <- 2*(1-pnorm(abs(zscore_sm),0,1))
sm.b <- data.frame(beta = b_sm, se = se_sm, p = pval_sm)

# Weighted
weight_iv <- weight/sum(weight) # weight each SNP by it's variance
b_wm = weighted.median(b_iv, weight_iv) 
se_wm = weighted.median.boot(b_exp, b_out, se_exp, se_out, weight_iv)
zscore_wm <- b_wm/se_wm
pval_wm <- 2*(1-pnorm(abs(zscore_wm),0,1))
wm.b <- data.frame(beta = b_wm, se = se_wm, p = pval_wm)

# Penalized
penalty = pchisq(weight_iv * (b_iv - beta_ivw)^2, df=1, lower.tail=FALSE)
pen.weights = weight_iv*pmin(1, penalty*20)
b_pwm = weighted.median(b_iv, pen.weights) 
se_pwm = weighted.median.boot(b_exp, b_out, se_exp, se_out, pen.weights)
zscore_pwm <- b_pwm/se_pwm
pval_pwm <- 2*(1-pnorm(abs(zscore_pwm),0,1))
pwm.b <- data.frame(beta = b_pwm, se = se_pwm, p = pval_pwm)

#################################
### Add lines to scatter plot ###
#################################
abline(ivw.r, col="red", lwd=2)
abline(egg.r, col="blue", lwd=2)
abline(a=0, b=b_wm, col="green", lwd=2)

legend("bottomright", c("IVW", "MR-Egger", "Weighted-median"), col=c("red", "blue", "green"), lty=1)

###################
### Funnel plot ###
###################
plot(data$beta_tg[!is.na(data$BMI)]/data$beta_bmi[!is.na(data$BMI)], sqrt(weight), xlab="Causal estimate", 
     ylab="1/s.e(causal estimate)", pch=16)
abline(v=ivw.b[1], col="red")
abline(v=egg.b[1], col="blue")
abline(v=b_wm, col="green")