#----------------------------------------------------------------------------------------
#  Appendix S3 of Specht et al. "Occupancy surveys with conditional replicates: An alternative sampling design for rare species"
#  Data simulation and analysis code for standard, removal and conditional occupancy designs
#
# code generated by T Arnold, Dept of Fisheries & Wildlife, Univ of Minnesota, arnol065@umn.edu
#----------------------------------------------------------------------------------------

rm(list=ls(all=TRUE)) # clear memory
library(R2jags)
#R2jags requires library "rjags", which may need to be downloaded from source (example below for mac): 
#install.packages("https://cran.r-project.org/bin/macosx/mavericks/contrib/3.3/rjags_4-6.tgz", repos=NULL, type="source")
#R2jags Also requires JAGS to be installed on machine (https://sourceforge.net/projects/mcmc-jags/files/)

# change this to location where you want to store simulation files
#setwd("C:/Users/") 

#Lines 88, 108, 109, 113 and 114 need input values. 

# Simple removal model, sample k times or until detected
sink("removal.txt")
cat("
    model {
    # Uniform priors bounded 0 to 1
    psi.rem ~ dunif(0, 1)
    p.rem ~ dunif(0, 1)
    
    # Estimate conditional detection probability for each of k surveys
    for (i in 1:k) {
    p.eff.rem[i] <- (p.rem*psi.rem*(1-p.rem)**(i-1))/(psi.rem*(1-p.rem)**(i-1)+(1-psi.rem))
    }
    
    for (i in 1:R) {
    for (j in 1:rem[i]) {
    y[i,j] ~ dbern(p.eff.rem[j])
    }
    } #close i
    } # end bugs model
    ",fill = TRUE)
sink()

# conditional detection model, sample once unless spp detected in first survey, then sample k times 
sink("conditional.txt")
cat("
    model {
    # Uniform priors bounded 0 to 1
    psi.cond ~ dunif(0, 1)
    p.cond ~ dunif(0, 1)
    
    # Estimate conditional detection probability for each of k surveys
    p.eff.cond[1] <- psi.cond*p.cond
    for (i in 2:k) {
    p.eff.cond[i] <- p.cond
    }
    
    # Likelihood
    for (i in 1:R) {
    for (j in 1:cond[i]) {
    y[i,j] ~ dbern(p.eff.cond[j])
    } # close j 
    } #close i
    }
    ",fill = TRUE)
sink()

# standard model, sample k times regardless
# code for this model adapted from M Kery's book
sink("standard.txt")
cat("
    model {
    # Uniform priors bounded 0 to 1
    psi.stand ~ dunif(0, 1)
    p.stand ~ dunif(0, 1)
    # Likelihood
    for (i in 1:R) {
    z[i] ~ dbern(psi.stand)
    p.eff.stand[i] <- z[i]*p.stand
    for (j in 1:k) {
    y[i,j] ~ dbern(p.eff.stand[i])
    } # close j 
    } #close i
    }
    ",fill = TRUE)
sink()

# MCMC settings (use 12000 ni, 2000 nb for real to get R-hat < 1.01 on nearly all simulations)
ni <- 12000 # iterations
nt <- 2     # thinning rate
nb <- 2000  # burn-in 
nc <- 3     # chains

Reps =  # number of simulations for each combination 

# Specify matrix dimensions for simulated data that will be sampled by each method
R = 1500     # R indicates potential # of sample plots simulated (with enough excess for all scenarios)
C = 24       # C indicates each plot can be surveyed up to 24 times (enough excess for all scenarios)

jags.out <- matrix(NA, nrow = Reps , ncol = 37) # create matrix to store output for simulations

# S is sites, E is total effort, Sd is sites where spp was detected, d is total detections, occ is true # occupied sites
colnames(jags.out) <- c("true.psi","true.p","max.reps","target.E", 
                        "rem.S","rem.E","rem.Sd","rem.d","rem.occ","psi.rem","psi.rem.sd","p.rem","p.rem.sd","rem.corr","rem.cov", 
                        "cond.S", "cond.E","cond.Sd", "cond.d","cond.occ","psi.cond","psi.cond.sd","p.cond","p.cond.sd","cond.corr","cond.cov",
                        "stand.S","stand.E","stand.Sd","stand.d","stand.occ","psi.stand","psi.stand.sd","p.stand","p.stand.sd","stand.corr","stand.cov")


#######
# Set unique parameters for psi, p, and removal sites
# only these next 11 lines and the final line (file name) need to be modified.
#######

psi =      # True occupancy probability (simulations include 0.1-0.9 at 0.1 intervals)
  p =       # True detection probability, given site is occupied (simulations include 0.1-0.9 at 0.1 intervals)
  
  ##### change values of k and E for each simulation
  # Specify true parameter values and number of replicates (k) and total effort (E = sites*replicates)
  k =      # max number of visits per site (we used 3, 5, 15)
  E =      # TOTAL sample effort (sites * reps), (we used 120, 420, 1500)
  
  
  R1= trunc(E/((k*(1-psi))+((psi*(1-(1-p)^k))/p))) # Expected value of removal sites at given psi, p, k, E
  R2 = trunc(E/(psi*p*k-psi*p+1)) # Expected value of conditional sites at given psi, p, k, E 
  R3 = trunc(E/k) # Expected value of standard sites at given E, k 

#### Simulate data across enough sites (s) and surveys (k) that they can be sampled by any design

for (s in 1:Reps){
  
  #### simulate data
  # Create structure to contain counts
  y <- matrix(NA, nrow = R, ncol = C) # create matrix with sites as rows, visits as columns (more sites and rows than needed by any single design)
  # Ecological process: Sample true occurrence (z, yes/no) from a Bernoulli (occurrence probability = psi)
  z <- rbinom(n = R, size = 1, prob = psi)
  # Observation process: Sample detection/nondetection observations from a Bernoulli(with p) if z=1
  for (j in 1:C){
    y[,j] <- rbinom(n = R, size = 1, prob = z * p) # if z = 0, prob will also be zero, if z 1, prob = p
  }
  #store data generating parameters
  jags.out[s,1]<-psi # true occupancy
  jags.out[s,2]<-p  # true detection
  jags.out[s,3]<-k  # number of visits per site
  jags.out[s,4]<-E  # target number of total surveys (sites*k)
  
  #REMOVAL SAMPLING: 
  
  rem<-c(rep(1,R1))  
  seen.rem<-c(rep(1,R1)) 
  #Loop creates vectors of lenght R1 (total sites) representing sites where >0 detections and # visits at each site until positive detection
  for (i in 1:R1){
    seen.rem[i] <- max(y[i,1:k])
    rem[i] <- min(which(y[i,]==1)) #Can ignore associated warnings
    rem[i] <- min(rem[i],k)}
  jags.out[s,5]<-R1  # removal site reps
  jags.out[s,6]<-sum(rem)  # removal total reps (sites * surveys)
  jags.out[s,7]<-sum(seen.rem)  # total occupied sites detected by removal method
  jags.out[s,8]<-sum(seen.rem)  # total detections by removal method (same as sites due to protocol)
  jags.out[s,9]<-sum(z[1:R1])  # total occupied sites for removal method 
  # save 10-15 for psi, SD(psi), p, SD(p), corr(Psi,p), cov(Psi,p)
  
  # create vector of total visits for conditional sampling (k if seen in survey 1, else 1)
  seen1<-c(rep(1,R2)) 
  cond<-c(rep(1,R2))
  sum.row<-c(rep(1,R2))
  for (i in 1:R2){
    seen1[i] <- y[i,1]
    cond[i] <- seen1[i]*(k-1) + 1  # takes on value of 1 (not seen) or k (seen)
    sum.row[i] <- seen1[i]*sum(y[i,1:k])}
  jags.out[s,16]<-R2  # conditional site reps
  jags.out[s,17]<-sum(cond)  # conditional total reps (sites * surveys)
  jags.out[s,18]<-sum(seen1)  # total occupied sites detected by conditional method
  jags.out[s,19]<-sum(sum.row)  # total organisms detected by conditional method
  jags.out[s,20]<-sum(z[1:R2])  # total occupied sites for conditional method 
  # save 21-26 for psi, SD(psi), p, SD(p), corr(Psi,p), cov(Psi,p)
  
  # summarize total detections for standard method (k visits at all sites)
  stand<-c(rep(1,R3))
  seen<-c(rep(1,R3))
  for (i in 1:R3){
    seen[i] <- max(y[i,1:k])
    stand[i] <- sum(y[i,1:k])}
  jags.out[s,27]<-R3  # conditional site reps
  jags.out[s,28]<-R3*k  # conditional total reps (sites * surveys)
  jags.out[s,29]<-sum(seen)  # total occupied sites detected by standard method
  jags.out[s,30]<-sum(stand)  # total organisms detected by standard method
  jags.out[s,31]<-sum(z[1:R3])  # total occupied sites for standard method 
  # save 32-37 for psi, SD(psi), p, SD(p), corr(Psi,p), cov(Psi,p)
  
  # Sample sizes (R#) to get 1000 reps for each estimator, on average
  win.data1 <- list(y=y, rem=rem, k=k, R=R1)  #Removal
  win.data2 <- list(y=y, cond=cond, k=k, R=R2)  #Conditional
  win.data3 <- list(y=y, k=k, R=R3)  #Standard
  
  #set initial values as observed occurrence
  inits1 <- function() list(psi.rem=runif(1,0,1), p.rem=runif(1,0,1))
  inits2 <- function() list(psi.cond=runif(1,0,1), p.cond=runif(1,0,1))
  zst <- apply(y[1:R3,1:k], 1, max)    # Observed occurrence as starting values for z
  inits3 <- function() list(z=zst, psi.stand=runif(1,0,1), p.stand=runif(1,0,1))
  
  params1 <- c("psi.rem", "p.rem")
  params2 <- c("psi.cond", "p.cond")
  params3 <- c("psi.stand", "p.stand")
  
  #### Run jags models. Output all saved to jags.out table.
  
  # Call jags from R 
  out.jags1 <- jags(win.data1, inits1, params1, "removal.txt", n.chains = nc, 
                    n.thin = nt, n.iter = ni, n.burnin = nb, working.directory = getwd())
  jags.out[s,10]=round(out.jags1$BUGSoutput$mean$psi.rem,4) # psi estimate
  jags.out[s,11]=round(out.jags1$BUGSoutput$sd$psi.rem,4) # sd psi
  jags.out[s,12]=round(out.jags1$BUGSoutput$mean$p.rem,4) # p estimate
  jags.out[s,13]=round(out.jags1$BUGSoutput$sd$p.rem,4) # sd p 
  jags.out[s,14]=round(cor(out.jags1$BUGSoutput$sims.matrix[,2],out.jags1$BUGSoutput$sims.matrix[,3]),6)
  jags.out[s,15]=round(cov(out.jags1$BUGSoutput$sims.matrix[,2],out.jags1$BUGSoutput$sims.matrix[,3]),6)
  
  # Call jags from R 
  out.jags2 <- jags(win.data2, inits2, params2, "conditional.txt", n.chains = nc, 
                    n.thin = nt, n.iter = ni, n.burnin = nb, working.directory = getwd())
  jags.out[s,21]=round(out.jags2$BUGSoutput$mean$psi.cond,4) # psi estimate
  jags.out[s,22]=round(out.jags2$BUGSoutput$sd$psi.cond,4) # sd psi
  jags.out[s,23]=round(out.jags2$BUGSoutput$mean$p.cond,4) # p estimate
  jags.out[s,24]=round(out.jags2$BUGSoutput$sd$p.cond,4) # sd p 
  jags.out[s,25]=round(cor(out.jags2$BUGSoutput$sims.matrix[,2],out.jags2$BUGSoutput$sims.matrix[,3]),6)
  jags.out[s,26]=round(cov(out.jags2$BUGSoutput$sims.matrix[,2],out.jags2$BUGSoutput$sims.matrix[,3]),6)
  
  # Call jags from R 
  out.jags3 <- jags(win.data3, inits3, params3, "standard.txt", n.chains = nc, 
                    n.thin = nt, n.iter = ni, n.burnin = nb, working.directory = getwd())
  jags.out[s,32]=round(out.jags3$BUGSoutput$mean$psi.stand,4) # psi estimate
  jags.out[s,33]=round(out.jags3$BUGSoutput$sd$psi.stand,4) # sd psi
  jags.out[s,34]=round(out.jags3$BUGSoutput$mean$p.stand,4) # p estimate
  jags.out[s,35]=round(out.jags3$BUGSoutput$sd$p.stand,4) # sd p 
  jags.out[s,36]=round(cor(out.jags3$BUGSoutput$sims.matrix[,2],out.jags3$BUGSoutput$sims.matrix[,3]),6)
  jags.out[s,37]=round(cov(out.jags3$BUGSoutput$sims.matrix[,2],out.jags3$BUGSoutput$sims.matrix[,3]),6)
  
  print(s)
} # end simulation
jags.out