###============================================================================================
### PACKAGES USED IN THIS SCRIPT
require(rjags)     # PACKAGE TO RUN THE jags MODEL. MANDATORY
require(MCMCvis)   # THIS PACKAGE CONTAINS THE MCMCsummary FUNCTION USED IN THIS SCRIPT
require(mcmcplots) # USED FOR THE CREATION OF THE CONVERGENCE PLOTS
require(DT)        # THIS LIBRARY ALLOWS THE NICE OUTPUT DISPLAY WITH THE SEARCH BAR OPTION.
###============================================================================================


###============================================================================================
### TO READ AND PREPARE THE DATA IN A SUITABLE FORMAT FOR THE LC ANALYSIS 
DATA <- read.table("Strongyloides.txt", header=TRUE)
t1 <- c(rep(1, DATA$n11), rep(1, DATA$n10), rep(0, DATA$n01), rep(0, DATA$n00) )
t2 <- c(rep(1, DATA$n11), rep(0, DATA$n10), rep(1, DATA$n01), rep(0, DATA$n00) )
# Joint test results for each patient (each row represents a different patient and each column represents a test result)
y <- cbind(t1, t2)
# Total number of patients
N = dim(y)[1]
dataList <- list(y=y, N=N)
###============================================================================================



###============================================================================================
### WRITE THE 2-LC MODEL AND SAVE IT AS A TEXT FILE
modelString =  
  "model {

for (i in 1:N) {

  #============
  # LIKELIHOOD 
  #============
  
	D1[i]~dbern(prev)
	D[i]<-D1[i]+1
	y[i,1]~dbern(p1[i,D[i]])
	y[i,2]~dbern(p2[i,D[i]])	
	
	# Defining the sensitivity and specificity of each subject (with inclusion of random effects)
	r[i]~dnorm(0,1)
	s1[i]<-phi(a[1,1]+b.RE[1]*r[i])
	c1[i]<-phi(a[1,2]+b.RE[2]*r[i])
	s2[i]<-phi(a[2,1]+b.RE[1]*r[i])
	c2[i]<-phi(a[2,2]+b.RE[2]*r[i])
	
	# Conditional probability of a positive observation
	p1[i,2]<-s1[i]
	p1[i,1]<-1-c1[i]
	p2[i,2]<-s2[i]
	p2[i,1]<-1-c2[i]

}

#==================================================
# Prior distributions of prevalence, sensitivities
# and specificities
#==================================================

prev~dbeta(1,1)
se[1]~dbeta(21.96,5.49)
sp[1]~dbeta(4.1,1.76) 
se[2]~dbeta(4.44,13.31)
sp[2]~dbeta(71.25,3.75)

#=========================================================
# Random effects parameters 
#=========================================================

a[1,1]<-probit(se[1])*sqrt(1+b.RE[1]*b.RE[1])
a[1,2]<-probit(sp[1])*sqrt(1+b.RE[1]*b.RE[1])
a[2,1]<-probit(se[2])*sqrt(1+b.RE[2]*b.RE[2])
a[2,2]<-probit(sp[2])*sqrt(1+b.RE[2]*b.RE[2])

b.RE[1]~dnorm(0,1)I(0,)
b.RE[2]~dnorm(0,1)I(0,)

}"
writeLines(modelString,con="model.txt")
###============================================================================================



###============================================================================================
### WRITE A HOME MADE FUNCTION TO GENERATE INITIAL VALUES ACCORDING TO THE PRIOR DISTRIBUTIONS
GenInits  = function(){  
  se1 <- rbeta(1,21.96,5.49)
  sp1 <- rbeta(1,4.1,1.76) 
  se2 <- rbeta(1,4.44,13.31)
  sp2 <- rbeta(1,71.25,3.75)
  b1 <- abs(rnorm(1,0,1))
  b2 <- abs(rnorm(1,0,1))
  prev <- rbeta(1,1,1)
  
  se <- c(se1, se2)
  sp <- c(sp1, sp2)
  b.RE <- c(b1, b2)
  
  list(
    se = se,
    sp = sp,
    b.RE = b.RE,
    prev = prev,
    .RNG.name="base::Wichmann-Hill",
    .RNG.seed=321
  )  
}
###============================================================================================


###============================================================================================
### GENERATE INITIAL VALUES FROM THE HOME MADE FUNCTION.  A SEED IS PROVIDED FOR REPRODUCIBILITY
set.seed(123)
initsList = vector('list',3)
for(i in 1:3){
  initsList[[i]] = GenInits()
}
###============================================================================================


###============================================================================================
### COMPILE THE MODEL
jagsModel = jags.model("model.txt",data=dataList,n.chains=3,n.adapt=0, inits=initsList)
###============================================================================================

###============================================================================================
### BURN-IN STEP 
update(jagsModel,n.iter=5000)
###============================================================================================

###============================================================================================
### PARAMETERS TO BE MONITORED
parameters = c( "se","sp", "a", "b.RE", "prev")
###============================================================================================

###============================================================================================
### POSTERIOR SAMPLING
posterior_results = coda.samples(jagsModel,variable.names=parameters,n.iter=5000)
output = posterior_results
###============================================================================================

###============================================================================================
### TO DISPLAY THE POSTERIOR ESTIMATES OF THE PARAMETERS WE MONITORED
res = MCMCsummary(output, digits=4)
datatable(res, extensions = 'AutoFill')
###============================================================================================




###============================================================================================
### CONVERGENCE DIAGNOSTIC TOOLS.
### FOR EACH PARAMETER MONITORED, WE CONSTRUCT DENSITY, HISTORY AND RUNNING MEAN PLOTS TO HELP
### ASSESS CONVERGENCE

### SENSITIVITY AND SPECIFICITY
for(k in 1:2) {
  for(i in 1:2) {
    #jpeg(paste(result_folder,"/",parameters[i],"[",k,"].jpeg",sep=""),width = 23, height = 23, units = "cm", res=200)
    par(oma=c(0,0,3,0))
    layout(matrix(c(1,2,3,3), 2, 2, byrow = TRUE))
    denplot(output, parms=c(paste(parameters[i],"[",k,"]",sep="")), auto.layout=FALSE, main="(a)", xlab=paste(parameters[i],"",sep=""), ylab="Density")
    rmeanplot(output, parms=c(paste(parameters[i],"[",k,"]",sep="")), auto.layout=FALSE, main="(b)")
    title(xlab="Iteration", ylab="Running mean")
    traplot(output, parms=c(paste(parameters[i],"[",k,"]",sep="")), auto.layout=FALSE, main="(c)")
    title(xlab="Iteration", ylab=paste(parameters[i],"[",k,"]",sep=""))
    mtext(paste("Diagnostics for ", parameters[i],"[",k,"]","",sep=""), side=3, line=1, outer=TRUE, cex=2)
    #dev.off()
  }
}

### PREVALENCE
for(i in 5) {
  #jpeg(paste(result_folder,"/",parameters[i],".jpeg",sep=""),width = 23, height = 23, units = "cm", res=200)
  par(oma=c(0,0,3,0))
  layout(matrix(c(1,2,3,3), 2, 2, byrow = TRUE))
  denplot(output, parms=c(paste(parameters[i], sep="")), auto.layout=FALSE, main="(a)", xlab=paste(parameters[i],"",sep=""), ylab="Density")
  rmeanplot(output, parms=c(paste(parameters[i], sep="")), auto.layout=FALSE, main="(b)")
  title(xlab="Iteration", ylab="Running mean")
  traplot(output, parms=c(paste(parameters[i], sep="")), auto.layout=FALSE, main="(c)")
  title(xlab="Iteration", ylab=paste(parameters[i], sep=""))
  mtext(paste("Diagnostics for ", parameters[i], sep=""), side=3, line=1, outer=TRUE, cex=2)
  #dev.off()
}


### RANDOM EFFECTS PARAMETERS

### a-PARAMETERS
for(i in c(3)) { #POSITION OF THE PARAMETER IN THE "parameters" OBJECT
  for(k in 1:2){
    for(j in 1:2) { 
      #jpeg(paste(result_folder,"/",parameters[i],"[",k,"].jpeg",sep=""),width = 23, height = 23, units = "cm", res=200)
      par(oma=c(0,0,3,0))
      layout(matrix(c(1,2,3,3), 2, 2, byrow = TRUE))
      denplot(output, parms=c(paste(parameters[i],"[",k,",",j,"]",sep="")), auto.layout=FALSE, main="(a)", xlab=paste(parameters[i],"[",k,"]",sep=""), ylab="Density")
      rmeanplot(output, parms=c(paste(parameters[i],"[",k,",",j,"]",sep="")), auto.layout=FALSE, main="(b)")
      title(xlab="Iteration", ylab="Running mean")
      traplot(output, parms=c(paste(parameters[i],"[",k,",",j,"]",sep="")), auto.layout=FALSE, main="(c)")
      title(xlab="Iteration", ylab=paste(parameters[i],"[",k,",",j,"]",sep=""))
      mtext(paste("Diagnostics for ", parameters[i],"[",k,",",j,"]",sep=""), side=3, line=1, outer=TRUE, cex=2)
      #dev.off()
    }
  }
}


### b.RE-PARAMETERS
for(k in 1:2) {
  for(i in 4) {
    #jpeg(paste(result_folder,"/",parameters[i],"[",k,"].jpeg",sep=""),width = 23, height = 23, units = "cm", res=200)
    par(oma=c(0,0,3,0))
    layout(matrix(c(1,2,3,3), 2, 2, byrow = TRUE))
    denplot(output, parms=c(paste(parameters[i],"[",k,"]",sep="")), auto.layout=FALSE, main="(a)", xlab=paste(parameters[i],"",sep=""), ylab="Density")
    rmeanplot(output, parms=c(paste(parameters[i],"[",k,"]",sep="")), auto.layout=FALSE, main="(b)")
    title(xlab="Iteration", ylab="Running mean")
    traplot(output, parms=c(paste(parameters[i],"[",k,"]",sep="")), auto.layout=FALSE, main="(c)")
    title(xlab="Iteration", ylab=paste(parameters[i],"[",k,"]",sep=""))
    mtext(paste("Diagnostics for ", parameters[i],"[",k,"]","",sep=""), side=3, line=1, outer=TRUE, cex=2)
    #dev.off()
  }
}
###============================================================================================