###============================================================================================
### 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.
require(dplyr)     # TO SETUP THE DATASET IN A SUITABLE FORMAT FOR rjags ANALYSIS 
require(rBeta2009) # TO GENERATE INITAL VALUES FOR LATENT CLASS PROBABILITIES
require(truncnorm) # TO GENERATE INITIAL VALUES FOR SPECIFICITY OF CUTLURE TEST BASED ON PRIOR
###============================================================================================


###============================================================================================
### TO READ AND PREPARE THE DATA IN A SUITABLE FORMAT FOR THE LC ANALYSIS 
DATA <- read.table("C_trachomatis.txt", header=TRUE)
# Joint test results for each patient (each row represents a different patient and each column represents a test result)
y <- DATA %>%  
  slice(rep(seq_len(n()), Obs_Freq)) %>% 
  select(-Obs_Freq)
# Total number of patients
N = dim(y)[1]
dataList <- list(y=y, N=N, L=3)
###============================================================================================



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

for (i in 1:N) {

  #============
  # LIKELIHOOD 
  #============
  
    for (j in 1:4) {
      y[i, j] ~ dbin(p[LC[i],i, j], 1)
    }
    
    LC[i] ~ dcat(pLC[1:L])
    r[i] ~ dnorm(0,1)
    
    #===========================================================================
    # Conditional probability of a positive observation for a given latent class
    #===========================================================================
    # LATENT CLASS L1 : viable C. trachomatis bacteria positive (D+) and true DNA status positive (l1+)
    p[1, i, 1] <- phi(a[1, 1] + b.RE[1] * r[i])
    p[1, i, 2] <- phi(a[1, 2] + b.RE[2] * r[i])
    p[1, i, 3] <- phi(a[1, 3])
    p[1, i, 4] <- phi(a[1, 4])
    
    # LATENT CLASS L2 : viable C. trachomatis bacteria negative (D-) and true DNA status positive (l1+)
    p[2, i, 1] <- phi(a[2, 1])
    p[2, i, 2] <- phi(a[2, 2])
    p[2, i, 3] <- phi(a[2, 3])
    p[2, i, 4] <- phi(a[2, 4])  
    
    # LATENT CLASS L3 : viable C. trachomatis bacteria negative (D-) and true DNA status negative (l1-)
    p[3, i, 1] <- phi(a[3, 1])
    p[3, i, 2] <- phi(a[3, 2])
    p[3, i, 3] <- phi(a[3, 3])
    p[3, i, 4] <- phi(a[3, 4])
    
    # LATENT CLASS 4 : viable C. trachomatis bacteria positive (D+) and true DNA status negative (l1-)
    # This latent class is assumed to be not possible
}

#==================================================
# Prior distributions of random effects parameters
#==================================================

a[1,1] ~ dnorm(0,1)
a[1,2] ~ dnorm(0,1)
a[1,3] ~ dnorm(0,1)
a[1,4] ~ dnorm(0,1)

a[2,1]  <- a[1,1]
a[2,2]  <- a[1,2]
a[2,3]  ~ dnorm(0,1)
a[2,4]  ~ dnorm(0,1)T(-5, -2.05)

a[3,1]  ~ dnorm(0,1) 
a[3,2]  ~ dnorm(0,1)
a[3,3]  <- a[2,3]
a[3,4]  <- a[2,4]
    
b.RE[1] ~ dunif(0, 5)
b.RE[2] <- b.RE[1]

#==================================================
# Prior distributions of latent class probabilities
#==================================================

pLC[1:L] ~ ddirch(prior[1:L])
for (i in 1:L) {
  prior[i]<-1
}

#=========================================================
# Other parameters of interest 
#=========================================================

# SENSITIVITY/SPECIFICITY WITH RESPECT TO viable C. trachomatis bacteria
se_D[1] <- phi(a[1, 1]/sqrt(1 + b.RE[1] * b.RE[1]))
se_D[2] <- phi(a[1, 2]/sqrt(1 + b.RE[2] * b.RE[2]))
se_D[3] <- phi(a[1, 3])
se_D[4] <- phi(a[1, 4])

sp_D[1] <- ( phi(-a[2, 1])*pLC[2] + phi(-a[3, 1])*pLC[3] )/(pLC[2]+pLC[3])
sp_D[2] <- ( phi(-a[2, 2])*pLC[2] + phi(-a[3, 2])*pLC[3] )/(pLC[2]+pLC[3])
sp_D[3] <- ( phi(-a[2, 3])*pLC[2] + phi(-a[3, 3])*pLC[3] )/(pLC[2]+pLC[3])
sp_D[4] <- ( phi(-a[2, 4])*pLC[2] + phi(-a[3, 4])*pLC[3] )/(pLC[2]+pLC[3])


# SENSITIVITY/SPECIFICITY WITH RESPECT TO true DNA status
se_l1[1] <- ( phi(a[1, 1]/sqrt(1 + b.RE[1] * b.RE[1]))*pLC[1] + phi(a[2, 1]/sqrt(1 + b.RE[1] * b.RE[1]))*pLC[2] )/(pLC[1]+pLC[2])
se_l1[2] <- ( phi(a[1, 2]/sqrt(1 + b.RE[2] * b.RE[2]))*pLC[1] + phi(a[2, 2]/sqrt(1 + b.RE[2] * b.RE[2]))*pLC[2] )/(pLC[1]+pLC[2])
se_l1[3] <- ( phi(a[1, 3])*pLC[1] + phi(a[2, 3])*pLC[2] )/(pLC[1]+pLC[2])
se_l1[4] <- ( phi(a[1, 4])*pLC[1] + phi(a[2, 4])*pLC[2] )/(pLC[1]+pLC[2])

sp_l1[1] <- phi(-a[3, 1])
sp_l1[2] <- phi(-a[3, 2])
sp_l1[3] <- phi(-a[3, 3])
sp_l1[4] <- phi(-a[3, 4])

# PREVALENCE OF viable C. trachomatis bacteria
prev <- pLC[1]

# PREVALENCE OF true DNA status
prev_DNA <- pLC[1]+pLC[2]

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



###============================================================================================
### WRITE A HOME MADE FUNCTION TO GENERATE INITIAL VALUES ACCORDING TO THE PRIOR DISTRIBUTIONS
GenInits  = function(){  
  a11 <- rnorm(1,0,1)
  a12 <- rnorm(1,0,1)
  a13 <- rnorm(1,0,1)
  a14 <- rnorm(1,0,1)
  a21 <- a11
  a22 <- a12
  a23 <- rnorm(1,0,1)
  a24 <- rtruncnorm(1,-5,-2.05,0,1)
  a31 <- rnorm(1,0,1)
  a32 <- rnorm(1,0,1)
  a33 <- a23
  a34 <- a24
  b1 <- runif(1,0,5)
  b2 <- b1
  pLC <- rdirichlet(1,c(1,2,3))
  
  b.RE <- c(b1, b2)
  a <- matrix(c(a11, a12, a13, a14,
                a21, a22, a23, a24,
                a31, a32, a33, a34), byrow=TRUE, ncol=4)
  pLC <- as.vector(pLC)
  
  
  list(
    a = a,
    pLC = pLC,
    b.RE = b.RE,
    .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_D","sp_D", "se_l1", "sp_l1", "a", "b.RE", "prev", "prev_DNA", "pLC")
###============================================================================================

###============================================================================================
### 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 WITH RESPECT TO viable *C. trachomatis* bacteria (D)
for(k in 1:4) {
  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()
  }
}

### SENSITIVITY AND SPECIFICITY WITH RESPECT TO *C. trachomatis* DNA (l1)
for(k in 1:4) {
  for(i in 3: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()
  }
}


### PREVALENCE
for(i in 7) {
  #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()
}

### PREVALENCE DNA
for(i in 8) {
  #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()
}

### LATENT CLASS PROBABILITIES
for(k in 1:3) {
  for(i in 9) {
    #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()
  }
}


### RANDOM EFFECTS PARAMETERS

### a-PARAMETERS
for(i in 5) { #POSITION OF THE PARAMETER IN THE "parameters" OBJECT
  for(k in 1:3){
    for(j in 1:4) { 
      #jpeg(paste(result_folder,"/",parameters[i],"[",k,",",j,"].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(i in 6) {
  #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()
}
###============================================================================================