#Chapter 11
#Example 11.1
set.seed(1)
N=100
x=rnorm(N,0,1)
e1=rnorm(N,0,1)
e2=rnorm(N,0,1)
alpha=0.98
beta=1.96
c=1.28
M=alpha*x+e1
y=c*x+beta*M+e2

install.packages("R2WinBUGS")
library(R2WinBUGS)

#11.2.1 Method 1
data1<-list(x=x,M=m,y=y,N=N)
inits<- function(){
  list(alpha=0,c=0,beta=0,prec1=0.5,prec2=0.5)
}
med1.1<-bugs(data1, inits, model.file =
               "C:/folder2/mediation_c_c_c_1.txt",
             parameters = c("alpha","c","beta",
                            "var1","var2","IE","DE"),
             n.chains=1,n.iter=11000,n.burnin=1000,
             n.thin=1,bugs.directory =
               "C:/folder1/winbugs14_full_patched/WinBUGS14/",
             debug = F)

med1.1

#WinBUGS model for method 1 (noted as model ccc1):
  model {
    #The model structure
    for(i in 1:N){
      mu_M[i] <- alpha*x[i]
      M[i] ~ dnorm(mu_M[i],prec1)
      mu_y[i] <-c*x[i]+beta*M[i]
      y[i] ~ dnorm(mu_y[i],prec2)
    }
    #Set up prior distributions
    alpha ~ dnorm(0.0,0.000001)
    beta ~ dnorm(0.0,0.000001)
    c ~ dnorm(0.0,0.000001)
    var1 ~ dgamma(1,0.1)
    var2 ~ dgamma(1,0.1)
    prec1 <-1/var1
    prec2 <-1/var2
    #Calculation of third-variable effects
    IE<-alpha*beta
    DE<-c
  }

#11.2.1.2 Method 2

#WinBUGS model for method 2 (noted as model ccc2):
  model {
    for(i in 1:N){
      mu_M[i] <- alpha*x[i]
      M[i] ~ dnorm(mu_M[i],prec1)
      mu_y[i]<- c*x[i]+beta*M[i]
      y[i] ~ dnorm(mu_y[i],prec2)
      mu_y1[i]<- c*(x[i]+deltax)+beta*M[i]
      de[i]<-(mu_y1[i]-mu_y[i])/deltax
      mu_M1[i] <- alpha*(x[i]+deltax)
      mu_y2[i]<- c*x[i]+beta*(M[i]+deltam)
      ie[i]<-(mu_M1[i]-mu_M[i])/deltax*(mu_y2[i]-mu_y[i])/deltam
      te[i]<-ie[i]+de[i]
    }
    alpha ~ dnorm(0.0,0.000001)
    beta ~ dnorm(0.0,0.000001)
    c ~ dnorm(0.0,0.000001)
    var1 ~ dgamma(1,0.1)
    var2 ~ dgamma(1,0.1)
    prec1 <-1/var1
    prec2 <-1/var2
  }

#11.2.1.3 Method 3
model {
  for(i in 1:N){
    mu_M[i] <- alpha*x[i]
    M[i] ~ dnorm(mu_M[i],prec1)
    mu_y[i]<- c*x[i]+beta*M[i]
    y[i] ~ dnorm(mu_y[i],prec2)
    mu_y0[i]<-c*x[i]+beta*mu_M[i]
    mu_M1[i] <- alpha*(x[i]+deltax)
    mu_y1[i]<- c*(x[i]+deltax)+beta*mu_M1[i]
    te[i]<-(mu_y1[i]-mu_y0[i])/deltax
    #draw random numbers between 1 and N
    j1[i]~dunif(0.5,100.5)
    j2[i]~dunif(0.5,100.5)
    j3[i]<- round(j1[i])
    j4[i]<- round(j2[i])
    mu_M2[i] <- alpha*(x[j3[i]])
    mu_y2[i]<- mu_y0[i]-beta*mu_M[i]+beta*mu_M2[i]
    mu_M3[i] <- alpha*(x[j4[i]])
    mu_y3[i]<- mu_y1[i]-beta*mu_M1[i]+beta*mu_M3[i]
    de[i]<-(mu_y3[i]-mu_y2[i])/deltax
    ie[i]<-te[i]-de[i]
  }
  alpha ~ dnorm(0.0,0.000001)
  beta ~ dnorm(0.0,0.000001)
  c ~ dnorm(0.0,0.000001)
  var1 ~ dgamma(1,0.1)
  var2 ~ dgamma(1,0.1)
  prec1 <-1/var1
  prec2 <-1/var2
}

# Example 11.4
mu_m=alpha*x
M3=rbinom(100,1,exp(mu_m)/(1+exp(mu_m)))
y2.1=c*x+beta*M3+e2

#modeld_cbc2
model {
  for(i in 1:N){
    logit(mu_M[i]) <- alpha*x[i]
    M[i] ~ dbern(mu_M[i])
    mu_y[i]<- c*x[i]+beta*M[i]
    y[i] ~ dnorm(mu_y[i],prec2)
    logit(mu_M1[i]) <- alpha*(x[i]+deltax)
    M1[i] ~ dbern(mu_M1[i])
    mu_y1[i]<- c*(x[i]+deltax)+beta*M[i]
    mu_y2[i]<- c*x[i]
    mu_y3[i]<- c*x[i]+beta
    ie[i]<-(mu_M1[i]-mu_M[i])/deltax*(mu_y3[i]-mu_y2[i])/deltam
    de[i]<-(mu_y1[i]-mu_y[i])/deltax
    te[i]<-ie[i]+de[i]
  }
  alpha ~ dnorm(0.0,0.01)
  beta ~ dnorm(0.0,0.01)
  c ~ dnorm(0.0,0.001)
  var2 ~ dgamma(1,0.1)
  prec2 <-1/var2
}

deltax=diff(range(x))/10
deltam=1
data1<-list(x=x,M=M3,y=y2.1,N=N,deltax=deltax,deltam=deltam)
inits<-function(){list(alpha=0,c=0,beta=0,prec2=0.5)}
med3.2<- bugs(data1, inits, parameters = c("alpha","c","beta",
                                           "var2","ie","de","te"),
              model.file = "C:/folder2/mediation_c_b_c_2.txt",
              n.chains = 1, n.iter = 11000,n.burnin=1000,
              bugs.directory = "C:/folder1/WinBUGS14/",
              debug = F)


#Example 11.5
set.seed(2)
alpha1=0.5
alpha2=0.8
beta1=1
beta2=1.2
mu_m1=alpha1*x
mu_m2=alpha2*x
p=cbind(1,exp(mu_m1),exp(mu_m2))
M4=t(apply(p,1,rmultinom,n=1,size=1))
y4=c*x+beta1*M4[,2]+beta2*M4[,3]+e2
M4.1=as.vector(M4%*%(1:3))

# Model mediation_c_cat_c_2
model {
  for(i in 1:N){
    mu_M1[i,1] <- 1 #baseline is the 1st category
    mu_M1[i,2] <- exp(alpha1*x[i])
    mu_M1[i,3] <- exp(alpha2*x[i])
    sum_M[i] <-mu_M1[i,1]+mu_M1[i,2]+mu_M1[i,3]
    for (k in 1:3)
    {mu_M[i,k] <- mu_M1[i,k]/sum_M[i]}
    M[i]~dcat(mu_M[i,])
    mu_y[i] <-c*x[i]+beta1*equals(M[i],2)+beta2*equals(M[i],3)
    y[i] ~ dnorm(mu_y[i],prec2)
    mu_y1[i]<- c*(x[i]+deltax)+beta1*equals(M[i],2)+beta2*equals(M[i],3)
    de[i]<-(mu_y1[i]-mu_y[i])/deltax
    mu_M1.2[i,1] <- 1 #baseline is the 1st category
    mu_M1.2[i,2] <- exp(alpha1*(x[i]+deltax))
    mu_M1.2[i,3] <- exp(alpha2*(x[i]+deltax))
    sum_M.2[i] <-mu_M1.2[i,1]+mu_M1.2[i,2]+mu_M1.2[i,3]
    for (k in 1:3)
    {mu_M.2[i,k] <- mu_M1.2[i,k]/sum_M.2[i]}
    mu_y2[i]<- c*x[i]
    mu_y3[i]<- c*x[i]+beta1
    mu_y4[i]<- c*x[i]+beta2
    ie1[i]<-(mu_M.2[i,2]-mu_M[i,2])/deltax*(mu_y3[i]-mu_y2[i])/deltam1
    ie2[i]<-(mu_M.2[i,3]-mu_M[i,3])/deltax*(mu_y4[i]-mu_y2[i])/deltam2
    ie[i]<-ie1[i]+ie2[i]
    te[i]<-ie[i]+de[i]
  }
  alpha1 ~ dnorm(0.0,0.01)
  alpha2 ~ dnorm(0.0,0.01)
  beta1 ~ dnorm(0.0,0.000001)
  beta2 ~ dnorm(0.0,0.000001)
  c ~ dnorm(0.0,0.001)
  var2 ~ dgamma(1,0.1)
  prec2 <-1/var2
}

# model mediation_c_cat_c_3
model{
   for(i in 1:N){
    mu_M1[i,1] <- 1 #baseline is the 1st category
    mu_M1[i,2] <- exp(alpha1*x[i])
    mu_M1[i,3] <- exp(alpha2*x[i])
    sum_M[i] <-mu_M1[i,1]+mu_M1[i,2]+mu_M1[i,3]
    for (k in 1:3)
    {mu_M[i,k] <- mu_M1[i,k]/sum_M[i]}
    M[i]~dcat(mu_M[i,])
    mu_y[i] <-c*x[i]+beta1*equals(M[i],2)+beta2*equals(M[i],3)
    y[i] ~ dnorm(mu_y[i],prec2)
    j1[i]~dunif(0.5,100.5)
    j2[i]~dunif(0.5,100.5)
    j3[i]<- round(j1[i])
    j4[i]<- round(j2[i])
    mu_M1.2[i,1] <- 1 #baseline is the 1st category
    mu_M1.2[i,2] <- exp(alpha1*(x[i]+deltax))
    mu_M1.2[i,3] <- exp(alpha2*(x[i]+deltax))
    sum_M.2[i] <-mu_M1.2[i,1]+mu_M1.2[i,2]+mu_M1.2[i,3]
    for (k in 1:3)
    {mu_M.2[i,k] <- mu_M1.2[i,k]/sum_M.2[i]}
    mu_y1[i]<- c*(x[i]+deltax)+beta1*mu_M.2[i,2]+beta2*mu_M.2[i,3]
    mu_y0[i] <-c*x[i]+beta1*mu_M[i,2]+beta2*mu_M[i,3]
    te[i]<-(mu_y1[i]-mu_y0[i])/deltax
    mu_M1.3[i,1] <- 1 #baseline is the 1st category
    mu_M1.3[i,2] <- exp(alpha1*x[j3[i]])
    mu_M1.3[i,3] <- exp(alpha2*x[j3[i]])
    sum_M.3[i] <- mu_M1.3[i,1]+mu_M1.3[i,2]+mu_M1.3[i,3]
    for (k in 1:3)
    {mu_M.3[i,k] <- mu_M1.3[i,k]/sum_M.3[i]}
    mu_y2[i]<- c*x[i]+beta1*mu_M.3[i,2]+beta2*mu_M.3[i,3]
    mu_M1.4[i,1] <- 1 #baseline is the 1st category
    #changed j4 to j3 to reduce variance
    mu_M1.4[i,2] <- exp(alpha1*x[j4[i]])
    mu_M1.4[i,3] <- exp(alpha2*x[j4[i]])
    sum_M.4[i] <- mu_M1.4[i,1]+mu_M1.4[i,2]+mu_M1.4[i,3]
    for (k in 1:3)
    {mu_M.4[i,k] <- mu_M1.4[i,k]/sum_M.4[i]}
    mu_y3[i]<- c*(x[i]+deltax)+beta1*mu_M.4[i,2]+beta2*mu_M.4[i,3]
    de[i]<-(mu_y3[i]-mu_y2[i])/deltax
    ie[i]<-te[i]-de[i]
   }
  alpha1 ~ dnorm(0.0,0.01)
  alpha2 ~ dnorm(0.0,0.01)
  beta1 ~ dnorm(0.0,0.000001)
  beta2 ~ dnorm(0.0,0.000001)
  c ~ dnorm(0.0,0.001)
  var2 ~ dgamma(1,0.1)
  prec2 <-1/var2
}

#11.3
#model_ccb2
model {
  for(i in 1:N){
    mu_M[i] <- alpha*x[i]
    M[i] ~ dnorm(mu_M[i],prec1)
    mu_y[i]<- c*x[i]+beta*M[i]
    y[i] ~ dnorm(mu_y[i],prec2)
    mu_y3[i]<-beta*M[i] #when x=0
    mu_y4[i]<-c+beta*M[i] #when x=1
    de[i]<-(mu_y4[i]-mu_y3[i])/deltax
    mu_y2[i]<-c*x[i]+beta*(M[i]+deltam)
    ie[i]<-alpha/deltax*(mu_y2[i]-mu_y[i])/deltam #alpha is alpha(1-0)
    te[i]<-ie[i]+de[i]
  }
  alpha ~ dnorm(0.0,0.000001)
  beta ~ dnorm(0.0,0.000001)
  c ~ dnorm(0.0,0.000001)
  var1 ~ dgamma(1,0.1)
  var2 ~ dgamma(1,0.1)
  prec1 <-1/var1
  prec2 <-1/var2
}

#model_ccb3
model {
  for(i in 1:N){
    mu_M[i] <- alpha*x[i]
    M[i] ~ dnorm(mu_M[i],prec1)
    mu_y[i]<- c*x[i]+beta*M[i]
    y[i] ~ dnorm(mu_y[i],prec2)
    j1[i]~dunif(0.5,100.5)
    j2[i]~dunif(0.5,100.5)
    j3[i]<- round(j1[i])
    j4[i]<- round(j2[i])
    M1[i] ~ dnorm(alpha,prec1)
    M0[i] ~ dnorm(0,prec1)
    mu_y1[i]<- c+beta*M1[i]
    mu_y0[i]<- beta*M0[i]
    te[i]<-(mu_y1[i]-mu_y0[i])/deltax
    mu_M2[i] <- alpha*(x[j3[i]])
    M2[i] ~ dnorm(mu_M2[i],prec1)
    mu_y2[i]<- beta*M2[i]
    mu_M3[i] <- alpha*(x[j4[i]])
    M3[i] ~ dnorm(mu_M3[i],prec1)
    mu_y3[i]<- c+beta*M3[i]
    de[i]<-(mu_y3[i]-mu_y2[i])/deltax
    ie[i]<-te[i]-de[i]
  }
  alpha ~ dnorm(0.0,0.000001)
  beta ~ dnorm(0.0,0.000001)
  c ~ dnorm(0.0,0.000001)
  var1 ~ dgamma(1,0.1)
  var2 ~ dgamma(1,0.1)
  prec1 <-1/var1
  prec2 <-1/var2
}