#################################################
# JAGS Code for Analysis of Plethodon shermani  #
# PIT-tag detection data                        #
#                                               #
# Grant Connette                                #
# University of Missouri                        #
# 3/20/2015                                     #
#################################################

# Load Libraries
library("R2jags")

# Set working directory
setwd("C:/Users/Owner/Desktop/Data")

###################
### DATA IMPORT

# Import capture histories (1=PIT detection only, 2=PIT & Visual detection, 3=not detected)
PITdata <- read.csv("Connette_PITdata.csv", header = T)
# Add in column for capture occasion (All individuals known to be alive/present at capture)
PITdata <- cbind(PITdata[,2],rep(2,nrow(PITdata)),PITdata[,3:ncol(PITdata)])
colnames(PITdata)[1:2] <- c("PIT_ID","V1")
CH <- PITdata[,-1]  # Final capture history

# Define model parameters
Ind <- 234     # Number of Individuals in Study
Surv <- 24     # Number of Surveys (including initial capture)

# Extract survey dates
SurveyDates <- t(read.csv("Connette_PITdata.csv",sep=",",header=F)[1,-c(1,2)])
SurveyDates <- julian(as.Date(SurveyDates, format="%m/%d/%Y"))  # Convert to vector in julian days

# Import initial capture dates for each individual
CapDates <- read.csv("CapDates.csv",sep=",",header=T)
CapDates$julian <- julian(as.Date(CapDates$Capture_Date, format="%m/%d/%Y"))

# Calculate matrix giving the time between surveys for each individual
# (only capture dates differ among individuals)
AllDates <- matrix(data=rep(c(NA,SurveyDates),Ind),nrow=Ind,ncol=Surv,byrow=T)

iter=0                       # Loop over individuals to assign their first capture dates
for (i in PITdata$PIT_ID){
  iter=iter+1
  cap <- which(CapDates$PIT_ID==i)
  AllDates[iter,1] <- CapDates$julian[cap]
}

Intervals <- t(diff(t(AllDates)))   # Final matrix giving intervals between surveys for each individual

# Import Covariates
Harvest <- read.csv("HarvestCovs.csv", header = T, check.names=F)
Harvest <- as.matrix(Harvest[,3:ncol(Harvest)])

Rain <- read.csv("DaysSinceRain.csv", header = T, row.names=1, check.names=F)
RainDays <- as.numeric(Rain) # Add NA for rainfall on capture occasion since we don't estimate detection


###################
### SPECIFY MODEL (Modified from Kery & Schaub 2012)

sink("PIT.txt")
cat("
model {
    
# PRIOR SPECIFICATION #
for (c in 1:3){
  b.pV[c] ~ dunif(-5, 5)   # Visual detection coefficients (Intercept,Rain,Harvest)
}
    
for (T in 1:2){            # Priors for survival, fidelity, temp. emigration, PIT-detection
  sT[T] ~ dunif(0,1)       # Survival
  fT[T] ~ dunif(0,1)       # Site Fidelity (1-P(Emigration))
  psiIOT[T] ~ dunif(0,1)   # Transition In-Out
  psiOIT[T] ~ dunif(0,1)   # Transition Out-In
  pT[T] ~ dunif(0,1)       # PIT-tag detection probability
}
    
# RANDOM EFFECTS AND CONSTRAINTS #
for (i in 1:Ind){  
  for (t in 1:(Surv-1)){
    s[i,t] <- pow(sT[Harvest[i,t]+1],Intervals[i,t]/30)   # scaling survival to the number of days between surveys
    f[i,t] <- pow(fT[Harvest[i,t]+1],Intervals[i,t]/30)   # same for site fidelity
    psiIO[i,t] <- psiIOT[Harvest[i,t]+1]
    psiOI[i,t] <- psiOIT[Harvest[i,t]+1]
    p1[i,t] <- pT[Harvest[i,t]+1]
    logit(p2[i,t]) <- b.pV[1] + b.pV[2]*RainDays[t] + b.pV[3]*Harvest[i,t]
    
    # Define transition matrix
    ps[1,i,t,1] <- f[i,t]*s[i,t]*(1-psiIO[i,t])
    ps[1,i,t,2] <- f[i,t]*s[i,t]*psiIO[i,t]
    ps[1,i,t,3] <- (1-f[i,t])*s[i,t]
    ps[1,i,t,4] <- (1-s[i,t])
    ps[1,i,t,5] <- 0
    ps[2,i,t,1] <- f[i,t]*s[i,t]*psiOI[i,t]
    ps[2,i,t,2] <- f[i,t]*s[i,t]*(1-psiOI[i,t])
    ps[2,i,t,3] <- (1-f[i,t])*s[i,t]
    ps[2,i,t,4] <- 0
    ps[2,i,t,5] <- (1-s[i,t])
    ps[3,i,t,1] <- 0
    ps[3,i,t,2] <- 0
    ps[3,i,t,3] <- 1
    ps[3,i,t,4] <- 0
    ps[3,i,t,5] <- 0
    ps[4,i,t,1] <- 0
    ps[4,i,t,2] <- 0
    ps[4,i,t,3] <- 0
    ps[4,i,t,4] <- 1
    ps[4,i,t,5] <- 0
    ps[5,i,t,1] <- 0
    ps[5,i,t,2] <- 0
    ps[5,i,t,3] <- 0
    ps[5,i,t,4] <- 0
    ps[5,i,t,5] <- 1
    
    # Define observation matrix
    po[1,i,t,1] <- (1-p2[i,t])*p1[i,t]
    po[1,i,t,2] <- p2[i,t]*p1[i,t]
    po[1,i,t,3] <- 1-p1[i,t]
    po[2,i,t,1] <- 0
    po[2,i,t,2] <- 0
    po[2,i,t,3] <- 1
    po[3,i,t,1] <- 0
    po[3,i,t,2] <- 0
    po[3,i,t,3] <- 1
    po[4,i,t,1] <- p1[i,t]
    po[4,i,t,2] <- 0
    po[4,i,t,3] <- 1-p1[i,t]
    po[5,i,t,1] <- 0
    po[5,i,t,2] <- 0
    po[5,i,t,3] <- 1
  } #t
} #i
    
# Likelihood 
for (i in 1:Ind){
  # Define latent state at first capture
  z[i,1] <- 1
    
  for (t in 2:Surv){
    # State process (state at time t given state at time t-1)
    z[i,t] ~ dcat(ps[z[i,t-1],i,t-1,])
    
    # Observation process (observation at time t given state at time t)
    y[i,t] ~ dcat(po[z[i,t],i,t-1,])
  } #t
} #i
    
# Derived Parameters
for (T in 1:2){
  eT[T] <- 1-fT[T]  # Emigration rate is the reciprocal of site fidelity
}
    
s.contrast <- sT[1]-sT[2]   # Harvest effect on survival
e.contrast <- eT[1]-eT[2]   # Harvest effect on emigration
    
}
",fill = TRUE)
sink()

# Bundle data
y=CH
jags.data <- list("y", "Surv", "Ind", "Intervals", "Harvest", "RainDays")

# Initial values
inits <- function(){list(fT = runif(2, 0.8, 1),
                         sT = runif(2, 0.8, 1),
                         psiIOT = runif(2, 0, 0.3),
                         psiOIT = runif(2, 0, 0.3),
                         pT = runif(2, 0.5, 1),
                         b.pV = rnorm(3,mean=c(-1,0,0.5),sd=0.2))}  

# Parameters monitored
parameters <- c("b.pV","pT","sT","fT","psiIOT","psiOIT","s.contrast","e.contrast")

library(R2jags)
# ~ 44 hrs. runtime
system.time(m1 <- jags.parallel(data=jags.data, inits=inits, parameters.to.save = parameters, model.file="PIT.txt", 
                                   n.chains = 3, 
                                   n.iter = 500000,
                                   n.burnin = 250000,
                                   n.thin = 100))

# Below runs analysis on a single processor
#system.time(m1 <- jags(data=jags.data, inits=inits, parameters.to.save = parameters, model.file="PIT.txt", 
#                                n.chains = 3, 
#                                n.iter = 500000,
#                                n.burnin = 250000,
#                                n.thin = 100))


# Check Parameter Identifiability
###############################################################################################

Priorcheck=function(y) {  # Using trapezoidal integration to calculate prior-posterior overlap
  Y<-rep(0,length(y$y))
  Y1<-rep(0,length(y$y))
  
  for (i in 1:length(y$y)){
    Y[i] <- min(1,y$y[i])
  }
  
  id <- order(y$x)
  AUCFull <- sum(diff(y$x[id])*rollmean(y$y[id],2))
  (AUCOverlap <- sum(diff(y$x[id])*rollmean(Y[id],2)))
}

sT1 <- m1$BUGSoutput$sims.list$sT[,1]
sT2 <- m1$BUGSoutput$sims.list$sT[,2]
fT1 <- m1$BUGSoutput$sims.list$fT[,1]
fT2 <- m1$BUGSoutput$sims.list$fT[,2]
pT1 <- m1$BUGSoutput$sims.list$pT[,1]
pT2 <- m1$BUGSoutput$sims.list$pT[,2]
eT1 <- 1-fT1
eT2 <- 1-fT2
psiIO1 <- m1$BUGSoutput$sims.list$psiIOT[,1]
psiIO2 <- m1$BUGSoutput$sims.list$psiIOT[,2]
psiOI1 <- m1$BUGSoutput$sims.list$psiOIT[,1]
psiOI2 <- m1$BUGSoutput$sims.list$psiOIT[,2]
b1.V <- m1$BUGSoutput$sims.list$b.pV[,1]
b2.V <- m1$BUGSoutput$sims.list$b.pV[,2]
b3.V <- m1$BUGSoutput$sims.list$b.pV[,3]

# Assess identifiability of parameters (Gimenez et al. 2009 - values <0.35 indicate identifiability)
y1 <- density(sT1)
round(Priorcheck(y1),digits=2)
y2 <- density(sT2)
round(Priorcheck(y2),digits=2)
y3 <- density(eT1)
round(Priorcheck(y3),digits=2)
y4 <- density(eT2)
round(Priorcheck(y4),digits=2)
y5 <- density(psiIO1)
round(Priorcheck(y5),digits=2)
y6 <- density(psiIO2)
round(Priorcheck(y6),digits=2)
y7 <- density(psiIO1)
round(Priorcheck(y7),digits=2)
y8 <- density(psiIO2)
round(Priorcheck(y8),digits=2)
y9 <- density(pT1)
round(Priorcheck(y9),digits=2)
y10 <- density(pT2)
round(Priorcheck(y10),digits=2)
y11 <- density(b1.V)
round(Priorcheck(y11),digits=2)
y12 <- density(b2.V)
round(Priorcheck(y12),digits=2)
y13 <- density(b3.V)
round(Priorcheck(y13),digits=2)