rm(list = ls())

# VIE Retention and Observation Experiment

# Chloe Fouilloux, Guillermo Garcia-Costoya, Bibiana Rojas 

#### BAYESIAN MODELS

# Load necessary packages 
library(tidyverse)
library(reshape2)
library(R2jags)
library(bayesplot)
library(MCMCvis)
library(coda)

### PART 1: PREPARATION OF THE DATA

# Read the data
frog <- read.csv("Elastomer_Final.csv", header = T, sep = ",")

# Keep only individuals from "Tag" treatment
frog$treatment <- frog$ï..Treatment
frog <- frog[which(frog$treatment=='Tag'), ] 

# Tagged individuals with NAs dropped
frog <- frog[!is.na(frog$Tag_Y_N), ] 

# Subset larval time 
frog <- frog[ which(frog$Week < 12), ] 

# Extract data into a simpler dataset
data <- mutate(frog,week=Week,obs=Tag_Y_N,id = ID,weight=Weight)%>%select(week,id,obs,weight)

# Add the Weight at Tagging variable 
data<-data%>%group_by(id)%>%mutate(weight_tag = rep(weight[2]))

# Modify the observation variable to have 1 & 0 instead of Y/N
data$obs <- ifelse(data$obs == "Y",1,0)

# Generate the "ret" column depedent on the observation colum. 
data <- data %>% group_by(id) %>% 
  mutate(ret=ifelse(max(week[which(obs==1)])>=week,1,NA)) %>% ungroup()


### PART 2: BAYESIAN MODELS 

## COMMON STUFF FOR ALL MODELS

# Change id from factor to unique number
data$id <- as.numeric(as.factor(data$id))

# Retention matrix
ret_matrix <- acast(data,week ~ id, value.var = "ret")

## Fill in 1 in NA gaps 
for (i in seq_len(ncol(ret_matrix))) {
  max_week <- max(which(ret_matrix[, i] == 1))
  ret_matrix[1:max_week, i] <- 1
}

# Observation matrix 
obs_matrix  <- acast(data,week ~ id, value.var="obs")

# Model data in jags format
model_data <- list(n_inds = ncol(ret_matrix),n_weeks = nrow(ret_matrix),observation = obs_matrix)

# Model data for models including weight at tag 
## Extract weight_tag vector for all ids
w_tag <- data%>%group_by(id)%>%summarise(w_tag=mean(weight_tag))
w_tag <- w_tag$w_tag

# Create the model data with weight at tag
model_data_tag <- list(n_inds = ncol(ret_matrix),n_weeks = nrow(ret_matrix),observation = obs_matrix,w_tag=w_tag)


# Options to run models
ni <- 1000000
nt <- 1000
nc <- 4


## MODEL WITH CONSTANT PROBABILITIES

# Model function
constant_prob_model <- function(){
  
  # Priors
  p_ret ~ dunif(0,1)
  p_obs ~ dunif(0,1)
  
  # Likelihood for the retention matrix (first week separated next weeks)
  for(i in 1:n_inds){
    retention[1,i]~ dbern(p_ret)
    for(w in 2:n_weeks){
      retention[w,i]~ dbern(retention[w-1,i]*p_ret)
    }
  }
  
  # Likelihood for the observation matrix 
  for(i in 1:n_inds){
    for(w in 1:n_weeks){
      observation[w,i]~ dbern(retention[w,i]*p_obs)
    }
  }
}

# Model results run 
constant_prob <- jags.parallel(data=model_data,model.file=constant_prob_model,parameters.to.save = c("p_ret","p_obs"),
                               inits = lapply(seq_len(1), function(i) {list(retention = ret_matrix)}),
                               n.chains =nc ,n.iter= 100000,n.thin= 10)

## MODEL WITH WEEK-DEPENDENT PROBABILITIES

# Model function
week_prob_model <- function(){
  
  # Prios
  alpha_r ~ dnorm(0,0.0001)
  beta_r ~ dnorm(0,0.0001)
  alpha_o ~ dnorm(0,0.0001)
  beta_o ~ dnorm(0,0.001)
  
  # Loop for the likelihood of the retention matrix
  for(i in 1:n_inds){
    
    # First week. 
    p_ret[1, i] <- 1/(1 + exp(-alpha_r))
    retention[1,i] ~ dbern(p_ret[1, i])
    
    # Second week on, there is the beta and current state is dependant on previous state 
    for(w in 2:n_weeks) {
      p_ret[w, i] <- 1/(1 + exp(-alpha_r - beta_r*(w-1)))
      retention[w,i] ~ dbern(retention[w-1,i]*p_ret[w, i]) 
    }
  }
  
  # Loop for the likelihood of the observation matrix 
  for(i in 1:n_inds){
    for(w in 1:n_weeks){
      p_obs[w, i]<- 1/(1 + exp(-alpha_o - beta_o*(w-1)))
      observation[w,i] ~ dbern(retention[w,i]*p_obs[w, i])
    }
  }
}

# Model results run 
week_prob <- jags.parallel(data=model_data,model.file=week_prob_model,parameters.to.save = c("alpha_r","beta_r","alpha_o","beta_o"),
                           init=lapply(seq_len(1),function(i){list(retention=ret_matrix)}),
                           n.chains = nc,n.iter = 100000,n.thin = 10)

## MODEL WITH WEEK DEPENDENT PROBABILITIES AND WEIGHT AT TAGGING 

# Model function
week_prob_model_tag <- function(){
  
  # Priors
  alpha_r ~ dnorm(0,0.0001)
  beta_r ~ dnorm(0,0.0001)
  beta_r_tag ~ dnorm(0,0.0001)
  alpha_o ~ dnorm(0,0.0001)
  beta_o ~ dnorm(0,0.001)
  beta_o_tag ~ dnorm(0,0.0001)
  
  # Loop for the likelihood of the retention matrix
  for(i in 1:n_inds){
    
    # First week. 
    p_ret[1, i] <- 1/(1 + exp(-alpha_r - beta_r_tag*(w_tag[i])))
    retention[1,i] ~ dbern(p_ret[1, i])
    
    # Second week on, there is the beta and current state is dependant on previous state 
    for(w in 2:n_weeks) {
      p_ret[w, i] <- 1/(1 + exp(-alpha_r - beta_r*(w-1) - beta_r_tag*(w_tag[i])))
      retention[w,i] ~ dbern(retention[w-1,i]*p_ret[w, i]) 
    }
  }
  
  # Loop for the likelihood of the observation matrix 
  for(i in 1:n_inds){
    for(w in 1:n_weeks){
      p_obs[w, i]<- 1/(1 + exp(-alpha_o - beta_o*(w-1)-beta_o_tag*(w_tag[i])))
      observation[w,i] ~ dbern(retention[w,i]*p_obs[w, i])
    }
  }
}

# Model results run 
week_prob_tag <- jags.parallel(data=model_data_tag,model.file=week_prob_model_tag,parameters.to.save = c("alpha_r","beta_r","alpha_o","beta_o","beta_r_tag","beta_o_tag"),
                           init=lapply(seq_len(1),function(i){list(retention=ret_matrix)}),
                           n.chains = nc,n.iter = 100000,n.thin = 10)


## MODEL WITH WEEK DEPENDENT PROBABILITIES AND RANDOM EFFECT OF ID

# Model function 
random_id_model <- function(){
  
  # General Priors
  alpha_ret ~ dnorm(0,0.0001)
  beta_ret ~ dnorm(0,0.0001)
  alpha_obs ~ dnorm(0,0.0001)
  beta_obs ~ dnorm(0,0.0001)
  
  # Priors for random effects
  s_alpha_ret_id ~ dnorm(0,100); T(0,)
  tau_alpha_ret_id <- pow(s_alpha_ret_id,-2)
  s_beta_ret_id ~ dnorm(0,100); T(0,)
  tau_beta_ret_id <- pow(s_beta_ret_id,-2)
  s_alpha_obs_id ~ dnorm(0,100); T(0,)
  tau_alpha_obs_id <- pow(s_alpha_obs_id,-2)
  s_beta_obs_id ~ dnorm(0,100); T(0,)
  tau_beta_obs_id <- pow(s_beta_obs_id,-2)
  
  # Loop for the likelihood of the retention matrix
  for(i in 1:n_inds){
    
    # Parameters for random effects
    alpha_ret_id[i] ~ dnorm(0,tau_alpha_ret_id)
    beta_ret_id[i] ~ dnorm(0,tau_beta_ret_id)
    
    # First week (logit link function + likelihood of the value whithin the matrix)
    p_ret[1, i] <- 1/(1 + exp(-alpha_ret -alpha_ret_id[i]))
    retention[1,i] ~ dbern(p_ret[1, i])
    
    # Second week on (same structure)
    for(w in 2:n_weeks) {
      p_ret[w, i] <- 1/(1 + exp(-alpha_ret -alpha_ret_id[i] - (beta_ret + beta_ret_id[i])*(w-1)))
      retention[w,i] ~ dbern(retention[w-1,i]*p_ret[w, i]) 
    }
  }
  
  # Loop for the likelihood of the observation matrix 
  for(i in 1:n_inds){
    
    # Parameters for random effects
    alpha_obs_id[i] ~ dnorm(0,tau_alpha_obs_id)
    beta_obs_id[i] ~ dnorm(0,tau_beta_obs_id)
    
    # Loop for each week (same structure as before)
    for(w in 1:n_weeks){
      p_obs[w, i]<- 1/(1 + exp(-alpha_obs -alpha_obs_id[i] - (beta_obs + beta_obs_id[i])*(w-1)))
      observation[w,i] ~ dbern(retention[w,i]*p_obs[w, i])
    }
  }
}

# Model results run 
random_id <- jags.parallel(data=model_data,model.file=random_id_model,parameters.to.save = c("alpha_ret","beta_ret","alpha_obs","beta_obs",
                                                                                              "s_alpha_ret_id","s_beta_ret_id","s_alpha_obs_id","s_beta_obs_id"),
                           init = lapply(seq_len(1), function(i) {list(retention = ret_matrix)}),n.chains=nc,n.iter=100000,n.thin=10)

## MODEL WITH WEEK DEPENDENT PROBABILITIES AND RANDOM EFFECT OF ID AND EFFECT OF WEIGHT AT TAG 

# Model function 
random_id_model_tag <- function(){
  
  # General Priors
  alpha_ret ~ dnorm(0,0.0001)
  beta_ret ~ dnorm(0,0.0001)
  beta_ret_tag ~ dnorm(0,0.0001)
  alpha_obs ~ dnorm(0,0.0001)
  beta_obs ~ dnorm(0,0.0001)
  beta_obs_tag ~ dnorm(0,0.0001)
  
  # Priors for random effects
  s_alpha_ret_id ~ dnorm(0,100); T(0,)
  tau_alpha_ret_id <- pow(s_alpha_ret_id,-2)
  s_beta_ret_id ~ dnorm(0,100); T(0,)
  tau_beta_ret_id <- pow(s_beta_ret_id,-2)
  s_alpha_obs_id ~ dnorm(0,100); T(0,)
  tau_alpha_obs_id <- pow(s_alpha_obs_id,-2)
  s_beta_obs_id ~ dnorm(0,100); T(0,)
  tau_beta_obs_id <- pow(s_beta_obs_id,-2)
  
  # Loop for the likelihood of the retention matrix
  for(i in 1:n_inds){
    
    # Parameters for random effects
    alpha_ret_id[i] ~ dnorm(0,tau_alpha_ret_id)
    beta_ret_id[i] ~ dnorm(0,tau_beta_ret_id)
    
    # First week (logit link function + likelihood of the value whithin the matrix)
    p_ret[1, i] <- 1/(1 + exp(-alpha_ret -alpha_ret_id[i] - beta_ret_tag*(w_tag[i])))
    retention[1,i] ~ dbern(p_ret[1, i])
    
    # Second week on (same structure)
    for(w in 2:n_weeks) {
      p_ret[w, i] <- 1/(1 + exp(-alpha_ret -alpha_ret_id[i] - (beta_ret + beta_ret_id[i])*(w-1)- beta_ret_tag*(w_tag[i])))
      retention[w,i] ~ dbern(retention[w-1,i]*p_ret[w, i]) 
    }
  }
  
  # Loop for the likelihood of the observation matrix 
  for(i in 1:n_inds){
    
    # Parameters for random effects
    alpha_obs_id[i] ~ dnorm(0,tau_alpha_obs_id)
    beta_obs_id[i] ~ dnorm(0,tau_beta_obs_id)
    
    # Loop for each week (same structure as before)
    for(w in 1:n_weeks){
      p_obs[w, i]<- 1/(1 + exp(-alpha_obs -alpha_obs_id[i] - (beta_obs + beta_obs_id[i])*(w-1)- beta_obs_tag*(w_tag[i])))
      observation[w,i] ~ dbern(retention[w,i]*p_obs[w, i])
    }
  }
}

# Model results run 
random_id_tag <- jags.parallel(data=model_data_tag,model.file=random_id_model_tag,
                               parameters.to.save = c("alpha_ret","beta_ret","alpha_obs","beta_obs","beta_ret_tag","beta_obs_tag",
                                                      "s_alpha_ret_id","s_beta_ret_id","s_alpha_obs_id","s_beta_obs_id"),
                           init = lapply(seq_len(1), function(i) {list(retention = ret_matrix)}),n.chains=nc,n.iter=100000,n.thin=10)

### PART 3 DIC SCORES COMPARISON AND MODEL SELECTION

# DIC Values for all models got into one vector
cp <- constant_prob$BUGSoutput$DIC
wp <- week_prob$BUGSoutput$DIC
ri <- random_id$BUGSoutput$DIC
wp_tag <- week_prob_tag$BUGSoutput$DIC
ri_tag <-random_id_tag$BUGSoutput$DIC

# DIC comparison dataframe
dic <- c(cp,wp,ri,wp_tag,ri_tag)
model <- c("Constant","Week ","ID","Week + Weight_Tag","ID + Weight_Tag")
delta_dic <- c(0,wp-cp,ri-cp,wp_tag-cp,ri_tag-cp)
comparison <- data.frame(model,dic,delta_dic)

### PART 4 CHAIN CONVERGENCE AND SAMPLING INDEPENDENCE

## MODEL WITH CONSTANT PROBABILITIES

# Change output into mcmc.list object
cprob <- as.mcmc.list(as.mcmc(constant_prob))

# Potential scale reduction factor (PSRF)
gelman.diag(cprob)

# Number of independent samples
effectiveSize(cprob)

## MODEL WITH WEEK-DEPENDENT PROBABILITIES

# Change output into mcmc.list object
wprob <- as.mcmc.list(as.mcmc(week_prob))

# Potential scale reduction factor (PSRF)
gelman.diag(wprob)

# Number of independent samples
effectiveSize(wprob)