rm(list=ls())
library(jagsUI) # library for running Jags
library(boot) # library for using logit function
library(hier.part) # library for hierarchical partitioning
library(vioplot) # library for making violin plots

setwd('c:/Zhao/')

#==============
# Read in data
#==============
areas <- c('Texas', 'New Mexico', 'Oklahoma')
years <- 1998:2018
nyear <- length(years)
ncoho <- 4

# 1. Population survey & distance sampling data
### 1a. Population survey of Texas
bpop_tx <- read.csv('1a. snpl_population_survey_Texas.csv')
count_tx <- bpop_tx$count
ncount_tx <- length(count_tx)
count_tx_year <- bpop_tx$year - min(years) + 1
count_tx_lake <- as.numeric(as.factor(bpop_tx$lake))
count_tx_date <- paste(count_tx_lake, bpop_tx$date, sep='-')
count_tx_date <- as.numeric(as.factor(count_tx_date))
count_tx_ndate <- length(unique(count_tx_date))

### 1b. Population survey of New Mexico
bpop_nm <- read.csv('1b. snpl_population_survey_New Mexico.csv')
count_nm <- bpop_nm$count
ncount_nm <- length(count_nm)
count_nm_year <- bpop_nm$year - min(years) + 1
count_nm_date <- as.numeric(as.factor(bpop_nm$date))
count_nm_ndate <- length(unique(count_nm_date))

### 1c. Population survey of Oklahoma
bpop_ok <- read.csv('1c. snpl_population_survey_Oklahoma.csv')
count_ok <- bpop_ok$count
ncount_ok <- length(count_ok)
count_ok_year <- bpop_ok$year - min(years) + 1
count_ok_grid <- as.numeric(as.factor(bpop_ok$grid))
count_ok_ngrid <- length(unique(count_ok_grid))

### 1d. Distance sampling of Oklahom
dist_ok <- read.csv('1d. snpl_distance_sampling_Oklahoma.csv')
dcount_ok <- dist_ok$count
ndcount_ok <- length(dcount_ok)
dcount_ok_year <- dist_ok$year - min(years) + 1
dcount_ok_dist <- dist_ok$distance / 10
dcount_ok_grid <- as.numeric(as.factor(dist_ok$grid))
dcount_ok_ngrid <- length(unique(dcount_ok_grid))

### Best guess of initial population sizes
bpop0 <- c(
  median(count_tx[which(count_tx_year == min(count_tx_year))]), 
  median(count_nm[which(count_nm_year == min(count_nm_year))]), 
  median(count_ok[which(count_ok_year == min(count_ok_year))]))

# 2. Encounter history
his <- read.csv('2. snpl_encounter_history.csv')
ench_all <- his[which(his$f != max(his$f)),]
ench_all <- ench_all[order(ench_all$region, ench_all$cohort, ench_all$f),]
ench_mat <- as.matrix(ench_all[,-c(1:4)])

### The original values are: 
### 1 not seen, 2 captured but not resighted, 3 resighted but not captured, 4 captured and resighted
### with an additional column to show the cohort
### The following code change the values to:
### 1 juvenile captured, only for the first year of capture, 
### 2 adult captured but not resighted, 3 adult resighted but not captured, 
### 4 adult captured and resighted; 5 not seen
ench_mat[which(ench_mat == 0)] <- 4
ench_mat <- ench_mat + 1
for (i in 1:dim(ench_mat)[1]) {
  ench_mat[i, ench_all$f[i]] <- ifelse(ench_all$cohort[i] %in% c('jf','jm','ju'), 1, 2)
}

ench_tx_f <- as.matrix(ench_mat[which(ench_all$cohort %in% c('af','jf') & ench_all$region=='Texas'),])
ench_tx_m <- as.matrix(ench_mat[which(ench_all$cohort %in% c('am','jm') & ench_all$region=='Texas'),])
f_tx_f <- ench_all$f[which(ench_all$cohort %in% c('af','jf') & ench_all$region=='Texas')]
f_tx_m <- ench_all$f[which(ench_all$cohort %in% c('am','jm') & ench_all$region=='Texas')]

ench_nm_f <- as.matrix(ench_mat[which(ench_all$cohort %in% c('af','jf') & ench_all$region=='New Mexico'),])
ench_nm_m <- as.matrix(ench_mat[which(ench_all$cohort %in% c('am','jm') & ench_all$region=='New Mexico'),])
f_nm_f <- ench_all$f[which(ench_all$cohort %in% c('af','jf') & ench_all$region=='New Mexico')]
f_nm_m <- ench_all$f[which(ench_all$cohort %in% c('am','jm') & ench_all$region=='New Mexico')]

ench_ok_f <- as.matrix(ench_mat[which(ench_all$cohort %in% c('af','jf') & ench_all$region=='Oklahoma'),])
ench_ok_m <- as.matrix(ench_mat[which(ench_all$cohort %in% c('am','jm') & ench_all$region=='Oklahoma'),])
f_ok_f <- ench_all$f[which(ench_all$cohort %in% c('af','jf') & ench_all$region=='Oklahoma')]
f_ok_m <- ench_all$f[which(ench_all$cohort %in% c('am','jm') & ench_all$region=='Oklahoma')]

visit_tx_temp <- colSums(ench_all[which(ench_all$region=='Texas'), -c(1:4)])
visit_nm_temp <- colSums(ench_all[which(ench_all$region=='New Mexico'), -c(1:4)])
visit_ok_temp <- colSums(ench_all[which(ench_all$region=='Oklahoma'), -c(1:4)])

visit_tx <- ifelse(visit_tx_temp == 0, 0, 1)
visit_nm <- ifelse(visit_nm_temp == 0, 0, 1)
visit_ok <- ifelse(visit_ok_temp == 0, 0, 1)

nind_tx_f <- length(f_tx_f)
nind_tx_m <- length(f_tx_m)
nind_nm_f <- length(f_nm_f)
nind_nm_m <- length(f_nm_m)
nind_ok_f <- length(f_ok_f)
nind_ok_m <- length(f_ok_m)

z_tx_f_known <- numeric(nind_tx_f)
for (i in 1:nind_tx_f) {z_tx_f_known[i] <- ench_tx_f[i,f_tx_f[i]]}
z_tx_m_known <- numeric(nind_tx_m)
for (i in 1:nind_tx_m) {z_tx_m_known[i] <- ench_tx_m[i,f_tx_m[i]]}

z_nm_f_known <- numeric(nind_nm_f)
for (i in 1:nind_nm_f) {z_nm_f_known[i] <- ench_nm_f[i,f_nm_f[i]]}
z_nm_m_known <- numeric(nind_nm_m)
for (i in 1:nind_nm_m) {z_nm_m_known[i] <- ench_nm_m[i,f_nm_m[i]]}

z_ok_f_known <- numeric(nind_ok_f)
for (i in 1:nind_ok_f) {z_ok_f_known[i] <- ench_ok_f[i,f_ok_f[i]]}
z_ok_m_known <- numeric(nind_ok_m)
for (i in 1:nind_ok_m) {z_ok_m_known[i] <- ench_ok_m[i,f_ok_m[i]]}

# 3. Clutch survey data
nest <- read.csv('3. snpl_nest_survey.csv')
clutch_year <- nest$year - min(years) + 1
clutch_region <- integer(length(nest$region)) # 1 Texas, 2 New Mexico, 3 Oklahoma
for (i in 1:length(nest$region)) {
  if (nest$region[i] == 'Texas') {
    clutch_region[i] <- 1
  } else if (nest$region[i] == 'New Mexico') {
    clutch_region[i] <- 2
  } else if (nest$region[i] == 'Oklahoma') {
    clutch_region[i] <- 3
  }
}
clutch <- as.integer(nest$clutch) # clutch size; take values of 1, 2 or 3
fate <- as.integer(nest$fate) # clutch fate; 1 succeed, 0 fail
clutch_nobs <- length(clutch)

### 4. Covariates
cov_tx <- read.csv('4a. snpl_covariates_Texas.csv')
cov_nm <- read.csv('4b. snpl_covariates_New Mexico.csv')
cov_ok <- read.csv('4c. snpl_covariates_Oklahoma.csv')

palm_ok <- (cov_ok$palmer - mean(cov_ok$palmer)) / sd(cov_ok$palmer)
tmin_ok <- (cov_ok$tmin - mean(cov_ok$tmin)) / sd(cov_ok$tmin)
wind_ok <- (cov_ok$wind - mean(cov_ok$wind)) / sd(cov_ok$wind)

palm_nm <- (cov_nm$palmer - mean(cov_nm$palmer)) / sd(cov_nm$palmer)
tmin_nm <- (cov_nm$tmin - mean(cov_nm$tmin)) / sd(cov_nm$tmin)
wind_nm <- (cov_nm$wind - mean(cov_nm$wind)) / sd(cov_nm$wind)

palm_tx <- (cov_tx$palmer - mean(cov_tx$palmer)) / sd(cov_tx$palmer)
tmin_tx <- (cov_tx$tmin - mean(cov_tx$tmin)) / sd(cov_tx$tmin)
wind_tx <- (cov_tx$wind - mean(cov_tx$wind)) / sd(cov_tx$wind)

#=======================
# Specify model in JAGS 
#=======================
sink('snpl_ipm_1 slope.txt')
cat("
  model {
  # Priors
  # 1.1 Survival
  ### Mean
  logit_sur_mean ~ dnorm(0, .01)
  logit_sur_mean_tau_cohort ~ dgamma(.01, .01)
  logit_sur_mean_sd_cohort <- 1 / sqrt(logit_sur_mean_tau_cohort)
  for (j in 1:ncoho) {
    logit_sur_mean_cohort[j] ~ dnorm(0, logit_sur_mean_tau_cohort)
  } # j
  logit_saf_mean <- logit_sur_mean + logit_sur_mean_cohort[1]
  logit_sam_mean <- logit_sur_mean + logit_sur_mean_cohort[2]
  logit_sjf_mean <- logit_sur_mean + logit_sur_mean_cohort[3]
  logit_sjm_mean <- logit_sur_mean + logit_sur_mean_cohort[4]

  ### Palmer drought severity index
  logit_sur_palm ~ dnorm(0, .01)
  logit_sur_palm_tau_cohort ~ dgamma(.01, .01)
  logit_sur_palm_sd_cohort <- 1 / sqrt(logit_sur_palm_tau_cohort)
  for (j in 1:ncoho) {
    logit_sur_palm_cohort[j] ~ dnorm(0, logit_sur_palm_tau_cohort)
  } # j
  logit_saf_palm <- logit_sur_palm + logit_sur_palm_cohort[1]
  logit_sam_palm <- logit_sur_palm + logit_sur_palm_cohort[2]
  logit_sjf_palm <- logit_sur_palm + logit_sur_palm_cohort[3]
  logit_sjm_palm <- logit_sur_palm + logit_sur_palm_cohort[4]

  ### Minimum temperature
  logit_sur_tmin ~ dnorm(0, .01)
  logit_sur_tmin_tau_cohort ~ dgamma(.01, .01)
  logit_sur_tmin_sd_cohort <- 1 / sqrt(logit_sur_tmin_tau_cohort)
  for (j in 1:ncoho) {
    logit_sur_tmin_cohort[j] ~ dnorm(0, logit_sur_tmin_tau_cohort)
  } # j
  logit_saf_tmin <- logit_sur_tmin + logit_sur_tmin_cohort[1]
  logit_sam_tmin <- logit_sur_tmin + logit_sur_tmin_cohort[2]
  logit_sjf_tmin <- logit_sur_tmin + logit_sur_tmin_cohort[3]
  logit_sjm_tmin <- logit_sur_tmin + logit_sur_tmin_cohort[4]

  ### Wind speed
  logit_sur_wind ~ dnorm(0, .01)
  logit_sur_wind_tau_cohort ~ dgamma(.01, .01)
  logit_sur_wind_sd_cohort <- 1 / sqrt(logit_sur_wind_tau_cohort)
  for (j in 1:ncoho) {
    logit_sur_wind_cohort[j] ~ dnorm(0, logit_sur_wind_tau_cohort)
  } # j
  logit_saf_wind <- logit_sur_wind + logit_sur_wind_cohort[1]
  logit_sam_wind <- logit_sur_wind + logit_sur_wind_cohort[2]
  logit_sjf_wind <- logit_sur_wind + logit_sur_wind_cohort[3]
  logit_sjm_wind <- logit_sur_wind + logit_sur_wind_cohort[4]

  logit_sur_tau ~ dgamma(.01, .01)
  logit_sur_sd <- 1 / sqrt(logit_sur_tau)

  # 1.1.1 Survival, Texas
  for (t in 1:nyear) {
    logit_saf_tx_mu[t] <- 
      logit_saf_mean + 
      logit_saf_palm * palm_tx[t] + 
      logit_saf_tmin * tmin_tx[t] + 
      logit_saf_wind * wind_tx[t]
    logit_sam_tx_mu[t] <- 
      logit_sam_mean + 
      logit_sam_palm * palm_tx[t] + 
      logit_sam_tmin * tmin_tx[t] + 
      logit_sam_wind * wind_tx[t]
    logit_sjf_tx_mu[t] <- 
      logit_sjf_mean + 
      logit_sjf_palm * palm_tx[t] + 
      logit_sjf_tmin * tmin_tx[t] + 
      logit_sjf_wind * wind_tx[t]
    logit_sjm_tx_mu[t] <- 
      logit_sjm_mean + 
      logit_sjm_palm * palm_tx[t] + 
      logit_sjm_tmin * tmin_tx[t] + 
      logit_sjm_wind * wind_tx[t]
    logit_saf_tx[t] ~ dnorm(logit_saf_tx_mu[t], logit_sur_tau)
    logit_sam_tx[t] ~ dnorm(logit_sam_tx_mu[t], logit_sur_tau)
    logit_sjf_tx[t] ~ dnorm(logit_sjf_tx_mu[t], logit_sur_tau)
    logit_sjm_tx[t] ~ dnorm(logit_sjm_tx_mu[t], logit_sur_tau)
    saf_tx[t] <- ilogit(logit_saf_tx[t])
    sam_tx[t] <- ilogit(logit_sam_tx[t])
    sjf_tx[t] <- ilogit(logit_sjf_tx[t])
    sjm_tx[t] <- ilogit(logit_sjm_tx[t])
  } # t

  # 1.1.2 Survival, New Mexico
  for (t in 1:nyear) {
    logit_saf_nm_mu[t] <- 
      logit_saf_mean + 
      logit_saf_palm * palm_nm[t] + 
      logit_saf_tmin * tmin_nm[t] + 
      logit_saf_wind * wind_nm[t]
    logit_sam_nm_mu[t] <- 
      logit_sam_mean + 
      logit_sam_palm * palm_nm[t] + 
      logit_sam_tmin * tmin_nm[t] + 
      logit_sam_wind * wind_nm[t]
    logit_sjf_nm_mu[t] <- 
      logit_sjf_mean + 
      logit_sjf_palm * palm_nm[t] + 
      logit_sjf_tmin * tmin_nm[t] + 
      logit_sjf_wind * wind_nm[t]
    logit_sjm_nm_mu[t] <- 
      logit_sjm_mean + 
      logit_sjm_palm * palm_nm[t] + 
      logit_sjm_tmin * tmin_nm[t] + 
      logit_sjm_wind * wind_nm[t]
    logit_saf_nm[t] ~ dnorm(logit_saf_nm_mu[t], logit_sur_tau)
    logit_sam_nm[t] ~ dnorm(logit_sam_nm_mu[t], logit_sur_tau)
    logit_sjf_nm[t] ~ dnorm(logit_sjf_nm_mu[t], logit_sur_tau)
    logit_sjm_nm[t] ~ dnorm(logit_sjm_nm_mu[t], logit_sur_tau)
    saf_nm[t] <- ilogit(logit_saf_nm[t])
    sam_nm[t] <- ilogit(logit_sam_nm[t])
    sjf_nm[t] <- ilogit(logit_sjf_nm[t])
    sjm_nm[t] <- ilogit(logit_sjm_nm[t])
  } # t

  # 1.1.3 Survival, Oklahoma
  for (t in 1:nyear) {
    logit_saf_ok_mu[t] <- 
      logit_saf_mean + 
      logit_saf_palm * palm_ok[t] + 
      logit_saf_tmin * tmin_ok[t] + 
      logit_saf_wind * wind_ok[t]
    logit_sam_ok_mu[t] <- 
      logit_sam_mean + 
      logit_sam_palm * palm_ok[t] + 
      logit_sam_tmin * tmin_ok[t] + 
      logit_sam_wind * wind_ok[t]
    logit_sjf_ok_mu[t] <- 
      logit_sjf_mean + 
      logit_sjf_palm * palm_ok[t] + 
      logit_sjf_tmin * tmin_ok[t] + 
      logit_sjf_wind * wind_ok[t]
    logit_sjm_ok_mu[t] <- 
      logit_sjm_mean + 
      logit_sjm_palm * palm_ok[t] + 
      logit_sjm_tmin * tmin_ok[t] + 
      logit_sjm_wind * wind_ok[t]
    logit_saf_ok[t] ~ dnorm(logit_saf_ok_mu[t], logit_sur_tau)
    logit_sam_ok[t] ~ dnorm(logit_sam_ok_mu[t], logit_sur_tau)
    logit_sjf_ok[t] ~ dnorm(logit_sjf_ok_mu[t], logit_sur_tau)
    logit_sjm_ok[t] ~ dnorm(logit_sjm_ok_mu[t], logit_sur_tau)
    saf_ok[t] <- ilogit(logit_saf_ok[t])
    sam_ok[t] <- ilogit(logit_sam_ok[t])
    sjf_ok[t] <- ilogit(logit_sjf_ok[t])
    sjm_ok[t] <- ilogit(logit_sjm_ok[t])
  } # t

  # 1.2 Capture/resighting probability
  pcap_af ~ dunif(0, 1) # Capture probability of adult females
  pcap_am ~ dunif(0, 1) # Capture probability of adult males
  pobs_af ~ dunif(0, 1) # Resighting probability of adult females
  pobs_am ~ dunif(0, 1) # Resighting probability of adult males

  # 2.1 Clutch size
  ### Clutch size = 2
  log_clutch2_mean ~ dnorm(0, .01)
  log_clutch2_palm ~ dnorm(0, .01)
  log_clutch2_tmin ~ dnorm(0, .01)
  log_clutch2_wind ~ dnorm(0, .01)

  ### Clutch size = 3
  log_clutch3_mean ~ dnorm(0, .01)
  log_clutch3_palm ~ dnorm(0, .01)
  log_clutch3_tmin ~ dnorm(0, .01)
  log_clutch3_wind ~ dnorm(0, .01)

  log_clutch_tau ~ dgamma(.01, .01)
  log_clutch_sd <- 1 / sqrt(log_clutch_tau)

  # 2.1.1 Clutch size, Texas
  for (t in 1:nyear) {
    log_clutch2_tx_mu[t] <- 
      log_clutch2_mean + 
      log_clutch2_palm * palm_tx[t] + 
      log_clutch2_tmin * tmin_tx[t] + 
      log_clutch2_wind * wind_tx[t]
    log_clutch3_tx_mu[t] <- 
      log_clutch3_mean + 
      log_clutch3_palm * palm_tx[t] + 
      log_clutch3_tmin * tmin_tx[t] + 
      log_clutch3_wind * wind_tx[t]
    log_clutch2_tx[t] ~ dnorm(log_clutch2_tx_mu[t], log_clutch_tau)
    log_clutch3_tx[t] ~ dnorm(log_clutch3_tx_mu[t], log_clutch_tau)
    clutch1_tx[t] <- 1 / (1 + exp(log_clutch2_tx[t]) + exp(log_clutch3_tx[t]))
    clutch2_tx[t] <- exp(log_clutch2_tx[t]) / (1 + exp(log_clutch2_tx[t]) + exp(log_clutch3_tx[t]))
    clutch3_tx[t] <- exp(log_clutch3_tx[t]) / (1 + exp(log_clutch2_tx[t]) + exp(log_clutch3_tx[t]))
    clutch_tx[t] <- 1 * clutch1_tx[t] + 2 * clutch2_tx[t] + 3 * clutch3_tx[t]
  } # t

  # 2.1.2 Clutch size, New Mexico
  for (t in 1:nyear) {
    log_clutch2_nm_mu[t] <- 
      log_clutch2_mean + 
      log_clutch2_palm * palm_nm[t] + 
      log_clutch2_tmin * tmin_nm[t] + 
      log_clutch2_wind * wind_nm[t]
    log_clutch3_nm_mu[t] <- 
      log_clutch3_mean + 
      log_clutch3_palm * palm_nm[t] + 
      log_clutch3_tmin * tmin_nm[t] + 
      log_clutch3_wind * wind_nm[t]
    log_clutch2_nm[t] ~ dnorm(log_clutch2_nm_mu[t], log_clutch_tau)
    log_clutch3_nm[t] ~ dnorm(log_clutch3_nm_mu[t], log_clutch_tau)
    clutch1_nm[t] <- 1 / (1 + exp(log_clutch2_nm[t]) + exp(log_clutch3_nm[t]))
    clutch2_nm[t] <- exp(log_clutch2_nm[t]) / (1 + exp(log_clutch2_nm[t]) + exp(log_clutch3_nm[t]))
    clutch3_nm[t] <- exp(log_clutch3_nm[t]) / (1 + exp(log_clutch2_nm[t]) + exp(log_clutch3_nm[t]))
    clutch_nm[t] <- 1 * clutch1_nm[t] + 2 * clutch2_nm[t] + 3 * clutch3_nm[t]
  } # t

  # 2.1.3 Clutch size, Oklahoma
  for (t in 1:nyear) {
    log_clutch2_ok_mu[t] <- 
      log_clutch2_mean + 
      log_clutch2_palm * palm_ok[t] + 
      log_clutch2_tmin * tmin_ok[t] + 
      log_clutch2_wind * wind_ok[t]
    log_clutch3_ok_mu[t] <- 
      log_clutch3_mean + 
      log_clutch3_palm * palm_ok[t] + 
      log_clutch3_tmin * tmin_ok[t] + 
      log_clutch3_wind * wind_ok[t]
    log_clutch2_ok[t] ~ dnorm(log_clutch2_ok_mu[t], log_clutch_tau)
    log_clutch3_ok[t] ~ dnorm(log_clutch3_ok_mu[t], log_clutch_tau)
    clutch1_ok[t] <- 1 / (1 + exp(log_clutch2_ok[t]) + exp(log_clutch3_ok[t]))
    clutch2_ok[t] <- exp(log_clutch2_ok[t]) / (1 + exp(log_clutch2_ok[t]) + exp(log_clutch3_ok[t]))
    clutch3_ok[t] <- exp(log_clutch3_ok[t]) / (1 + exp(log_clutch2_ok[t]) + exp(log_clutch3_ok[t]))
    clutch_ok[t] <- 1 * clutch1_ok[t] + 2 * clutch2_ok[t] + 3 * clutch3_ok[t]
  } # t

  clutch_prob[1,1:nyear,1] <- clutch1_tx[1:nyear]
  clutch_prob[1,1:nyear,2] <- clutch2_tx[1:nyear]
  clutch_prob[1,1:nyear,3] <- clutch3_tx[1:nyear]
  clutch_prob[2,1:nyear,1] <- clutch1_nm[1:nyear]
  clutch_prob[2,1:nyear,2] <- clutch2_nm[1:nyear]
  clutch_prob[2,1:nyear,3] <- clutch3_nm[1:nyear]
  clutch_prob[3,1:nyear,1] <- clutch1_ok[1:nyear]
  clutch_prob[3,1:nyear,2] <- clutch2_ok[1:nyear]
  clutch_prob[3,1:nyear,3] <- clutch3_ok[1:nyear]

  # 2.2 Clutch fate
  logit_fate_mean ~ dnorm(0, .01)
  logit_fate_palm ~ dnorm(0, .01)
  logit_fate_tmin ~ dnorm(0, .01)
  logit_fate_wind ~ dnorm(0, .01)

  logit_fate_tau ~ dgamma(.01, .01)
  logit_fate_sd <- 1 / sqrt(logit_fate_tau)

  # 2.2.1 Clutch fate, Texas
  for (t in 1:nyear) {
    logit_fate_tx_mu[t] <- 
      logit_fate_mean + 
      logit_fate_palm * palm_tx[t] + 
      logit_fate_tmin * tmin_tx[t] + 
      logit_fate_wind * wind_tx[t]
    logit_fate_tx[t] ~ dnorm(logit_fate_tx_mu[t], logit_fate_tau)
    fate_tx[t] <- ilogit(logit_fate_tx[t])
  } # t

  # 2.2.2 Clutch fate, New Mexico
  for (t in 1:nyear) {
    logit_fate_nm_mu[t] <- 
      logit_fate_mean + 
      logit_fate_palm * palm_nm[t] + 
      logit_fate_tmin * tmin_nm[t] + 
      logit_fate_wind * wind_nm[t]
    logit_fate_nm[t] ~ dnorm(logit_fate_nm_mu[t], logit_fate_tau)
    fate_nm[t] <- ilogit(logit_fate_nm[t])
  } # t

  # 2.2.3 Clutch fate, Oklahoma
  for (t in 1:nyear) {
    logit_fate_ok_mu[t] <- 
      logit_fate_mean + 
      logit_fate_palm * palm_ok[t] + 
      logit_fate_tmin * tmin_ok[t] + 
      logit_fate_wind * wind_ok[t]
    logit_fate_ok[t] ~ dnorm(logit_fate_ok_mu[t], logit_fate_tau)
    fate_ok[t] <- ilogit(logit_fate_ok[t])
  } # t

  fate_prob[1,1:nyear] <- fate_tx[1:nyear]
  fate_prob[2,1:nyear] <- fate_nm[1:nyear]
  fate_prob[3,1:nyear] <- fate_ok[1:nyear]

  # 3.1 Population
  log_pop_tau ~ dgamma(.01, .01)
  log_pop_sd <- 1 / sqrt(log_pop_tau)

  log_count_tau ~ dgamma(.01, .01)
  log_count_sd <- 1 / sqrt(log_count_tau)

  for (i in 1:count_tx_ndate) {
    log_count_tx_ep[i] ~ dnorm(0, log_count_tau)
  } # i

  for (i in 1:count_nm_ndate) {
    log_count_nm_ep[i] ~ dnorm(0, log_count_tau)
  } # i

  for (i in 1:count_ok_ngrid) {
    log_count_ok_ep[i] ~ dnorm(0, log_count_tau)
  } # i

  for (i in 1:dcount_ok_ngrid) {
    log_dcount_ok_ep[i] ~ dnorm(0, log_count_tau)
  } # i

  theta ~ dgamma(.01, .01)
  pcount <- exp(-1 * theta * 3)
  for (i in 1:ndcount_ok) {
    pdist[i] <- exp(-1 * theta * dcount_ok_dist[i])
  } # i

  # Process model
  # 1.1 Survival (value: 1 juvenile alive; 2 adult alive; 3 dead)
  for (t in 1:nyear) {
    # Texas, female
    ps_tx_f[1,1,t] <- 0
    ps_tx_f[1,2,t] <- sjf_tx[t]
    ps_tx_f[1,3,t] <- 1 - sjf_tx[t]
    ps_tx_f[2,1,t] <- 0
    ps_tx_f[2,2,t] <- saf_tx[t]
    ps_tx_f[2,3,t] <- 1 - saf_tx[t]
    ps_tx_f[3,1,t] <- 0
    ps_tx_f[3,2,t] <- 0
    ps_tx_f[3,3,t] <- 1

    # Texas, male
    ps_tx_m[1,1,t] <- 0
    ps_tx_m[1,2,t] <- sjm_tx[t]
    ps_tx_m[1,3,t] <- 1 - sjm_tx[t]
    ps_tx_m[2,1,t] <- 0
    ps_tx_m[2,2,t] <- sam_tx[t]
    ps_tx_m[2,3,t] <- 1 - sam_tx[t]
    ps_tx_m[3,1,t] <- 0
    ps_tx_m[3,2,t] <- 0
    ps_tx_m[3,3,t] <- 1

    # New Mexico, female
    ps_nm_f[1,1,t] <- 0
    ps_nm_f[1,2,t] <- sjf_nm[t]
    ps_nm_f[1,3,t] <- 1 - sjf_nm[t]
    ps_nm_f[2,1,t] <- 0
    ps_nm_f[2,2,t] <- saf_nm[t]
    ps_nm_f[2,3,t] <- 1 - saf_nm[t]
    ps_nm_f[3,1,t] <- 0
    ps_nm_f[3,2,t] <- 0
    ps_nm_f[3,3,t] <- 1

    # New Mexico, male
    ps_nm_m[1,1,t] <- 0
    ps_nm_m[1,2,t] <- sjm_nm[t]
    ps_nm_m[1,3,t] <- 1 - sjm_nm[t]
    ps_nm_m[2,1,t] <- 0
    ps_nm_m[2,2,t] <- sam_nm[t]
    ps_nm_m[2,3,t] <- 1 - sam_nm[t]
    ps_nm_m[3,1,t] <- 0
    ps_nm_m[3,2,t] <- 0
    ps_nm_m[3,3,t] <- 1

    # Oklahoma, female
    ps_ok_f[1,1,t] <- 0
    ps_ok_f[1,2,t] <- sjf_ok[t]
    ps_ok_f[1,3,t] <- 1 - sjf_ok[t]
    ps_ok_f[2,1,t] <- 0
    ps_ok_f[2,2,t] <- saf_ok[t]
    ps_ok_f[2,3,t] <- 1 - saf_ok[t]
    ps_ok_f[3,1,t] <- 0
    ps_ok_f[3,2,t] <- 0
    ps_ok_f[3,3,t] <- 1

    # Oklahoma, male
    ps_ok_m[1,1,t] <- 0
    ps_ok_m[1,2,t] <- sjm_ok[t]
    ps_ok_m[1,3,t] <- 1 - sjm_ok[t]
    ps_ok_m[2,1,t] <- 0
    ps_ok_m[2,2,t] <- sam_ok[t]
    ps_ok_m[2,3,t] <- 1 - sam_ok[t]
    ps_ok_m[3,1,t] <- 0
    ps_ok_m[3,2,t] <- 0
    ps_ok_m[3,3,t] <- 1
  } # t

  # Texas, female
  for (i in 1:nind_tx_f) {
    z_tx_f[i,f_tx_f[i]] <- z_tx_f_known[i]
    for (t in (f_tx_f[i]+1):nyear) {
      z_tx_f[i,t] ~ dcat(ps_tx_f[z_tx_f[i,t-1],,t-1])
    } # t
  } # i

  # Texas, male
  for (i in 1:nind_tx_m) {
    z_tx_m[i,f_tx_m[i]] <- z_tx_m_known[i]
    for (t in (f_tx_m[i]+1):nyear) {
      z_tx_m[i,t] ~ dcat(ps_tx_m[z_tx_m[i,t-1],,t-1])
    } # t
  } # i

  # New Mexico, female
  for (i in 1:nind_nm_f) {
    z_nm_f[i,f_nm_f[i]] <- z_nm_f_known[i]
    for (t in (f_nm_f[i]+1):nyear) {
      z_nm_f[i,t] ~ dcat(ps_nm_f[z_nm_f[i,t-1],,t-1])
    } # t
  } # i

  # New Mexico, male
  for (i in 1:nind_nm_m) {
    z_nm_m[i,f_nm_m[i]] <- z_nm_m_known[i]
    for (t in (f_nm_m[i]+1):nyear) {
      z_nm_m[i,t] ~ dcat(ps_nm_m[z_nm_m[i,t-1],,t-1])
    } # t
  } # i

  # Oklahoma, female
  for (i in 1:nind_ok_f) {
    z_ok_f[i,f_ok_f[i]] <- z_ok_f_known[i]
    for (t in (f_ok_f[i]+1):nyear) {
      z_ok_f[i,t] ~ dcat(ps_ok_f[z_ok_f[i,t-1],,t-1])
    } # t
  } # i

  # Oklahoma, male
  for (i in 1:nind_ok_m) {
    z_ok_m[i,f_ok_m[i]] <- z_ok_m_known[i]
    for (t in (f_ok_m[i]+1):nyear) {
      z_ok_m[i,t] ~ dcat(ps_ok_m[z_ok_m[i,t-1],,t-1])
    } # t
  } # i

  # 3.1 Population
  # 3.1.1 Texas population
  log_pop_tx[1] ~ dnorm(log(bpop0[1]), 100)
  for (t in 2:nyear) {
    log_pop_mu_tx[t-1] <- log(
      exp(log_pop_tx[t-1]) * 0.5 * sam_tx[t-1] + 
      exp(log_pop_tx[t-1]) * 0.5 * saf_tx[t-1] + 
      exp(log_pop_tx[t-1]) * 0.5 * clutch_tx[t-1] * fate_tx[t-1] * (sjm_tx[t-1] + sjf_tx[t-1]) / 2
    )
    log_pop_tx[t] ~ dnorm(log_pop_mu_tx[t-1], log_pop_tau)
  } # t

  for (t in 1:nyear) {
    pop_tx[t] <- exp(log_pop_tx[t])
  } # t

  # 3.1.2 New Mexico population
  log_pop_nm[1] ~ dnorm(log(bpop0[2]), 100)
  for (t in 2:nyear) {
    log_pop_mu_nm[t-1] <- log(
      exp(log_pop_nm[t-1]) * 0.5 * sam_nm[t-1] + 
      exp(log_pop_nm[t-1]) * 0.5 * saf_nm[t-1] + 
      exp(log_pop_nm[t-1]) * 0.5 * clutch_nm[t-1] * fate_nm[t-1] * (sjm_nm[t-1] + sjf_nm[t-1]) / 2
    )
    log_pop_nm[t] ~ dnorm(log_pop_mu_nm[t-1], log_pop_tau)
  } # t

  for (t in 1:nyear) {
    pop_nm[t] <- exp(log_pop_nm[t])
  } # t

  # 3.1.3 Oklahoma population
  log_pop_ok[1] ~ dnorm(log(bpop0[3]), 100)
  for (t in 2:nyear) {
    log_pop_mu_ok[t-1] <- log(
      exp(log_pop_ok[t-1]) * 0.5 * sam_ok[t-1] + 
      exp(log_pop_ok[t-1]) * 0.5 * saf_ok[t-1] + 
      exp(log_pop_ok[t-1]) * 0.5 * clutch_ok[t-1] * fate_ok[t-1] * (sjm_ok[t-1] + sjf_ok[t-1]) / 2
    )
    log_pop_ok[t] ~ dnorm(log_pop_mu_ok[t-1], log_pop_tau)
  } # t

  for (t in 1:nyear) {
    pop_ok[t] <- exp(log_pop_ok[t])
  } # t

  # Observation model
  # 1 Encounter data (value: 1 juvenile captured, only for the first year of capture; 
  #                          2 adult captured but not resighted; 3 adult resighted but not captured; 
  #                          4 adult captured and resighted; 5 not seen)
  for (t in 1:nyear) {
    # Texas, female
    po_tx_f[1,1,t] <- 0
    po_tx_f[1,2,t] <- 0
    po_tx_f[1,3,t] <- 0
    po_tx_f[1,4,t] <- 0
    po_tx_f[1,5,t] <- 1
    po_tx_f[2,1,t] <- 0
    po_tx_f[2,2,t] <- pcap_af * visit_tx[t] * (1 - pobs_af * visit_tx[t])
    po_tx_f[2,3,t] <- (1 - pcap_af * visit_tx[t]) * pobs_af * visit_tx[t]
    po_tx_f[2,4,t] <- pcap_af * visit_tx[t] * pobs_af * visit_tx[t]
    po_tx_f[2,5,t] <- (1 - pcap_af * visit_tx[t]) * (1 - pobs_af * visit_tx[t])
    po_tx_f[3,1,t] <- 0
    po_tx_f[3,2,t] <- 0
    po_tx_f[3,3,t] <- 0
    po_tx_f[3,4,t] <- 0
    po_tx_f[3,5,t] <- 1

    # Texas, male
    po_tx_m[1,1,t] <- 0
    po_tx_m[1,2,t] <- 0
    po_tx_m[1,3,t] <- 0
    po_tx_m[1,4,t] <- 0
    po_tx_m[1,5,t] <- 1
    po_tx_m[2,1,t] <- 0
    po_tx_m[2,2,t] <- pcap_am * visit_tx[t] * (1 - pobs_am * visit_tx[t])
    po_tx_m[2,3,t] <- (1 - pcap_am * visit_tx[t]) * pobs_am * visit_tx[t]
    po_tx_m[2,4,t] <- pcap_am * visit_tx[t] * pobs_am * visit_tx[t]
    po_tx_m[2,5,t] <- (1 - pcap_am * visit_tx[t]) * (1 - pobs_am * visit_tx[t])
    po_tx_m[3,1,t] <- 0
    po_tx_m[3,2,t] <- 0
    po_tx_m[3,3,t] <- 0
    po_tx_m[3,4,t] <- 0
    po_tx_m[3,5,t] <- 1

    # New Mexico, female
    po_nm_f[1,1,t] <- 0
    po_nm_f[1,2,t] <- 0
    po_nm_f[1,3,t] <- 0
    po_nm_f[1,4,t] <- 0
    po_nm_f[1,5,t] <- 1
    po_nm_f[2,1,t] <- 0
    po_nm_f[2,2,t] <- pcap_af * visit_nm[t] * (1 - pobs_af * visit_nm[t])
    po_nm_f[2,3,t] <- (1 - pcap_af * visit_nm[t]) * pobs_af * visit_nm[t]
    po_nm_f[2,4,t] <- pcap_af * visit_nm[t] * pobs_af * visit_nm[t]
    po_nm_f[2,5,t] <- (1 - pcap_af * visit_nm[t]) * (1 - pobs_af * visit_nm[t])
    po_nm_f[3,1,t] <- 0
    po_nm_f[3,2,t] <- 0
    po_nm_f[3,3,t] <- 0
    po_nm_f[3,4,t] <- 0
    po_nm_f[3,5,t] <- 1

    # New Mexico, male
    po_nm_m[1,1,t] <- 0
    po_nm_m[1,2,t] <- 0
    po_nm_m[1,3,t] <- 0
    po_nm_m[1,4,t] <- 0
    po_nm_m[1,5,t] <- 1
    po_nm_m[2,1,t] <- 0
    po_nm_m[2,2,t] <- pcap_am * visit_nm[t] * (1 - pobs_am * visit_nm[t])
    po_nm_m[2,3,t] <- (1 - pcap_am * visit_nm[t]) * pobs_am * visit_nm[t]
    po_nm_m[2,4,t] <- pcap_am * visit_nm[t] * pobs_am * visit_nm[t]
    po_nm_m[2,5,t] <- (1 - pcap_am * visit_nm[t]) * (1 - pobs_am * visit_nm[t])
    po_nm_m[3,1,t] <- 0
    po_nm_m[3,2,t] <- 0
    po_nm_m[3,3,t] <- 0
    po_nm_m[3,4,t] <- 0
    po_nm_m[3,5,t] <- 1

    # Oklahoma, female
    po_ok_f[1,1,t] <- 0
    po_ok_f[1,2,t] <- 0
    po_ok_f[1,3,t] <- 0
    po_ok_f[1,4,t] <- 0
    po_ok_f[1,5,t] <- 1
    po_ok_f[2,1,t] <- 0
    po_ok_f[2,2,t] <- pcap_af * visit_ok[t] * (1 - pobs_af * visit_ok[t])
    po_ok_f[2,3,t] <- (1 - pcap_af * visit_ok[t]) * pobs_af * visit_ok[t]
    po_ok_f[2,4,t] <- pcap_af * visit_ok[t] * pobs_af * visit_ok[t]
    po_ok_f[2,5,t] <- (1 - pcap_af * visit_ok[t]) * (1 - pobs_af * visit_ok[t])
    po_ok_f[3,1,t] <- 0
    po_ok_f[3,2,t] <- 0
    po_ok_f[3,3,t] <- 0
    po_ok_f[3,4,t] <- 0
    po_ok_f[3,5,t] <- 1

    # Oklahoma, male
    po_ok_m[1,1,t] <- 0
    po_ok_m[1,2,t] <- 0
    po_ok_m[1,3,t] <- 0
    po_ok_m[1,4,t] <- 0
    po_ok_m[1,5,t] <- 1
    po_ok_m[2,1,t] <- 0
    po_ok_m[2,2,t] <- pcap_am * visit_ok[t] * (1 - pobs_am * visit_ok[t])
    po_ok_m[2,3,t] <- (1 - pcap_am * visit_ok[t]) * pobs_am * visit_ok[t]
    po_ok_m[2,4,t] <- pcap_am * visit_ok[t] * pobs_am * visit_ok[t]
    po_ok_m[2,5,t] <- (1 - pcap_am * visit_ok[t]) * (1 - pobs_am * visit_ok[t])
    po_ok_m[3,1,t] <- 0
    po_ok_m[3,2,t] <- 0
    po_ok_m[3,3,t] <- 0
    po_ok_m[3,4,t] <- 0
    po_ok_m[3,5,t] <- 1
  } # t

  # Encounter data, Texas, female
  for (i in 1:nind_tx_f) {
    for (t in (f_tx_f[i]+1):nyear) {
      ench_tx_f_exp[i,t,1:5] <- po_tx_f[z_tx_f[i,t],,t]
      ench_tx_f[i,t] ~ dcat(ench_tx_f_exp[i,t,1:5])
    } # t
  } # i

  # Encounter data, Texas, male
  for (i in 1:nind_tx_m) {
    for (t in (f_tx_m[i]+1):nyear) {
      ench_tx_m_exp[i,t,1:5] <- po_tx_m[z_tx_m[i,t],,t]
      ench_tx_m[i,t] ~ dcat(ench_tx_m_exp[i,t,1:5])
    } # t
  } # i

  # Encounter data, New Mexico, female
  for (i in 1:nind_nm_f) {
    for (t in (f_nm_f[i]+1):nyear) {
      ench_nm_f_exp[i,t,1:5] <- po_nm_f[z_nm_f[i,t],,t]
      ench_nm_f[i,t] ~ dcat(ench_nm_f_exp[i,t,1:5])
    } # t
  } # i

  # Encounter data, New Mexico, male
  for (i in 1:nind_nm_m) {
    for (t in (f_nm_m[i]+1):nyear) {
      ench_nm_m_exp[i,t,1:5] <- po_nm_m[z_nm_m[i,t],,t]
      ench_nm_m[i,t] ~ dcat(ench_nm_m_exp[i,t,1:5])
    } # t
  } # i

  # Encounter data, Oklahoma
  for (i in 1:nind_ok_f) {
    for (t in (f_ok_f[i]+1):nyear) {
      ench_ok_f_exp[i,t,1:5] <- po_ok_f[z_ok_f[i,t],,t]
      ench_ok_f[i,t] ~ dcat(ench_ok_f_exp[i,t,1:5])
    } # t
  } # i

  # Encounter data, New Mexico, male
  for (i in 1:nind_ok_m) {
    for (t in (f_ok_m[i]+1):nyear) {
      ench_ok_m_exp[i,t,1:5] <- po_ok_m[z_ok_m[i,t],,t]
      ench_ok_m[i,t] ~ dcat(ench_ok_m_exp[i,t,1:5])
    } # t
  } # i

  # 2 Clutch size & fate data
  for (j in 1:clutch_nobs) {
    clutch_exp[j,1:3] <- clutch_prob[clutch_region[j],clutch_year[j],1:3]
    clutch[j] ~ dcat(clutch_exp[j,1:3])
    fate_exp[j] <- fate_prob[clutch_region[j],clutch_year[j]]
    fate[j] ~ dbin(fate_exp[j], 1)
  } # j

  # 3 Population count data
  # Texas
  for (k in 1:ncount_tx) {
    lambda_tx[k] <- exp(
      log(pop_tx[count_tx_year[k]]) + 
      log_count_tx_ep[count_tx_date[k]]) * 
      pcount
    count_tx[k] ~ dpois(lambda_tx[k])
  } # k

  # New Mexico
  for (k in 1:ncount_nm) {
    lambda_nm[k] <- exp(
      log(pop_nm[count_nm_year[k]]) + 
      log_count_nm_ep[count_nm_date[k]]) * 
      pcount
    count_nm[k] ~ dpois(lambda_nm[k])
  } # k

  # Oklahoma, point count
  for (k in 1:ncount_ok) {
    lambda_ok[k] <- exp(
      log(pop_ok[count_ok_year[k]]) + 
      log_count_ok_ep[count_ok_grid[k]])
      * pcount
    count_ok[k] ~ dpois(lambda_ok[k])
  } # k

  # Oklahoma, distance sampling
  for (k in 1:ndcount_ok) {
    lambda_dist_ok[k] <- exp(
      log(pop_ok[dcount_ok_year[k]]) + 
      log_dcount_ok_ep[dcount_ok_grid[k]])
      * pdist[k]
    dcount_ok[k] ~ dpois(lambda_dist_ok[k])
  } # k

  } # model
  ",fill = TRUE)
sink()

#==========
# Run Jags
#==========
# Texas
z_tx_f_temp <- matrix(, nind_tx_f, nyear)
for (i in 1:nind_tx_f) {
  z_tx_f_temp[i, (f_tx_f[i]+1):nyear] <- 2
}
z_tx_f_init <- z_tx_f_data <- z_tx_f_temp
z_tx_f_init[which(ench_tx_f != 5)] <- NA
z_tx_f_data[which(ench_tx_f == 5)] <- NA

z_tx_m_temp <- matrix(, nind_tx_m, nyear)
for (i in 1:nind_tx_m) {
  z_tx_m_temp[i, (f_tx_m[i]+1):nyear] <- 2
}
z_tx_m_init <- z_tx_m_data <- z_tx_m_temp
z_tx_m_init[which(ench_tx_m != 5)] <- NA
z_tx_m_data[which(ench_tx_m == 5)] <- NA

# New Mexico
z_nm_f_temp <- matrix(, nind_nm_f, nyear)
for (i in 1:nind_nm_f) {
  z_nm_f_temp[i, (f_nm_f[i]+1):nyear] <- 2
}
z_nm_f_init <- z_nm_f_data <- z_nm_f_temp
z_nm_f_init[which(ench_nm_f != 5)] <- NA
z_nm_f_data[which(ench_nm_f == 5)] <- NA

z_nm_m_temp <- matrix(, nind_nm_m, nyear)
for (i in 1:nind_nm_m) {
  z_nm_m_temp[i, (f_nm_m[i]+1):nyear] <- 2
}
z_nm_m_init <- z_nm_m_data <- z_nm_m_temp
z_nm_m_init[which(ench_nm_m != 5)] <- NA
z_nm_m_data[which(ench_nm_m == 5)] <- NA

# Oklahoma
z_ok_f_temp <- matrix(, nind_ok_f, nyear)
for (i in 1:nind_ok_f) {
  z_ok_f_temp[i, (f_ok_f[i]+1):nyear] <- 2
}
z_ok_f_init <- z_ok_f_data <- z_ok_f_temp
z_ok_f_init[which(ench_ok_f != 5)] <- NA
z_ok_f_data[which(ench_ok_f == 5)] <- NA

z_ok_m_temp <- matrix(, nind_ok_m, nyear)
for (i in 1:nind_ok_m) {
  z_ok_m_temp[i, (f_ok_m[i]+1):nyear] <- 2
}
z_ok_m_init <- z_ok_m_data <- z_ok_m_temp
z_ok_m_init[which(ench_ok_m != 5)] <- NA
z_ok_m_data[which(ench_ok_m == 5)] <- NA

# Bundle data
jags.data <- list(
  # Basic
  nyear = nyear, ncoho = ncoho, 

  # Encounter, Texas
  nind_tx_f=nind_tx_f, nind_tx_m=nind_tx_m,  
  ench_tx_f=ench_tx_f, ench_tx_m=ench_tx_m, 
  f_tx_f=f_tx_f, f_tx_m=f_tx_m, 
  z_tx_f_known=z_tx_f_known, z_tx_m_known=z_tx_m_known, 
  visit_tx=visit_tx, 
  z_tx_f=z_tx_f_data, z_tx_m=z_tx_m_data, 

  # Encounter, New Mexico
  nind_nm_f=nind_nm_f, nind_nm_m=nind_nm_m,  
  ench_nm_f=ench_nm_f, ench_nm_m=ench_nm_m, 
  f_nm_f=f_nm_f, f_nm_m=f_nm_m, 
  z_nm_f_known=z_nm_f_known, z_nm_m_known=z_nm_m_known, 
  visit_nm=visit_nm, 
  z_nm_f=z_nm_f_data, z_nm_m=z_nm_m_data, 

  # Encounter, Oklahoma
  nind_ok_f=nind_ok_f, nind_ok_m=nind_ok_m,  
  ench_ok_f=ench_ok_f, ench_ok_m=ench_ok_m, 
  f_ok_f=f_ok_f, f_ok_m=f_ok_m, 
  z_ok_f_known=z_ok_f_known, z_ok_m_known=z_ok_m_known, 
  visit_ok=visit_ok, 
  z_ok_f=z_ok_f_data, z_ok_m=z_ok_m_data, 

  # Clutch size & fate
  clutch_nobs=clutch_nobs, clutch_region=clutch_region, clutch_year=clutch_year, 
  clutch=clutch, fate=fate, 

  # Population count
  count_tx=count_tx, ncount_tx=ncount_tx, 
  count_tx_year=count_tx_year, 
  count_tx_date=count_tx_date, count_tx_ndate=count_tx_ndate, 
  count_nm=count_nm, ncount_nm=ncount_nm, 
  count_nm_year=count_nm_year, 
  count_nm_date=count_nm_date, count_nm_ndate=count_nm_ndate, 
  count_ok=count_ok, ncount_ok=ncount_ok, 
  count_ok_year=count_ok_year, 
  count_ok_grid=count_ok_grid, count_ok_ngrid=count_ok_ngrid, 
  dcount_ok=dcount_ok, ndcount_ok=ndcount_ok, 
  dcount_ok_year=dcount_ok_year, 
  dcount_ok_dist=dcount_ok_dist, 
  dcount_ok_grid=dcount_ok_grid, dcount_ok_ngrid=dcount_ok_ngrid, 
  bpop0=bpop0, 

  # Covariates
  palm_tx=palm_tx, tmin_tx=tmin_tx, wind_tx=wind_tx, 
  palm_nm=palm_nm, tmin_nm=tmin_nm, wind_nm=wind_nm, 
  palm_ok=palm_ok, tmin_ok=tmin_ok, wind_ok=wind_ok
)

# Initial values
inits <- function() {
  list(
    logit_sur_mean=0, logit_sur_mean_tau_cohort=1, 
    logit_sur_palm=0, logit_sur_palm_tau_cohort=1, 
    logit_sur_tmin=0, logit_sur_tmin_tau_cohort=1, 
    logit_sur_wind=0, logit_sur_wind_tau_cohort=1, 
    logit_sur_tau=1, 
    pcap_af=.5, pcap_am=.5, pobs_af=.5, pobs_am=.5, 
    z_tx_f=z_tx_f_init, z_tx_m=z_tx_m_init, 
    z_nm_f=z_nm_f_init, z_nm_m=z_nm_m_init, 
    z_ok_f=z_ok_f_init, z_ok_m=z_ok_m_init, 
    log_clutch2_mean=0, log_clutch2_palm=0, log_clutch2_tmin=0, log_clutch2_wind=0, 
    log_clutch3_mean=0, log_clutch3_palm=0, log_clutch3_tmin=0, log_clutch3_wind=0, 
    log_clutch_tau=1, 
    logit_fate_mean=0, logit_fate_palm=0, logit_fate_tmin=0, logit_fate_wind=0, 
    logit_fate_tau=1, 
    log_pop_tau=1, log_count_tau=1, 
    theta=1
)}

# Parameters monitored
parameters <- c(
  'saf_tx', 'sam_tx', 'sjf_tx', 'sjm_tx', 
  'saf_nm', 'sam_nm', 'sjf_nm', 'sjm_nm', 
  'saf_ok', 'sam_ok', 'sjf_ok', 'sjm_ok', 
  'clutch_tx', 'clutch_nm', 'clutch_ok', 
  'fate_tx', 'fate_nm', 'fate_ok', 
  'pop_tx', 'pop_nm', 'pop_ok', 
  'logit_sur_mean', 'logit_sur_mean_sd_cohort', 
  'logit_saf_mean', 'logit_sam_mean', 'logit_sjf_mean', 'logit_sjm_mean', 
  'logit_sur_palm', 'logit_sur_palm_sd_cohort', 
  'logit_saf_palm', 'logit_sam_palm', 'logit_sjf_palm', 'logit_sjm_palm', 
  'logit_sur_tmin', 'logit_sur_tmin_sd_cohort', 
  'logit_saf_tmin', 'logit_sam_tmin', 'logit_sjf_tmin', 'logit_sjm_tmin', 
  'logit_sur_wind', 'logit_sur_wind_sd_cohort', 
  'logit_saf_wind', 'logit_sam_wind', 'logit_sjf_wind', 'logit_sjm_wind', 
  'logit_sur_sd', 
  'pcap_af', 'pcap_am', 'pobs_af', 'pobs_am', 
  'log_clutch2_mean', 'log_clutch2_palm', 'log_clutch2_tmin', 'log_clutch2_wind', 
  'log_clutch3_mean', 'log_clutch3_palm', 'log_clutch3_tmin', 'log_clutch3_wind', 
  'log_clutch_sd', 
  'logit_fate_mean', 'logit_fate_palm', 'logit_fate_tmin', 'logit_fate_wind', 
  'logit_fate_sd', 
  'log_pop_sd', 'log_count_sd', 
  'theta', 'pcount'
)

# Call JAGS from R
fit <- jags(jags.data, inits, parameters, model.file='snpl_ipm_1 slope.txt', 
            n.chains=4, n.adapt=2000, n.burnin=50000, n.iter=300000, n.thin=100, 
            parallel=TRUE)

save(fit, file='snpl_ipm_1 slope.RData')

#===========================
# Hierarchical partitioning
#===========================
# Extracting posterior samples
### Adult female survival
saf_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='saf_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='saf_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='saf_tx[')])
saf_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='saf_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='saf_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='saf_nm[')])
saf_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='saf_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='saf_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='saf_ok[')])

### Adult male survival
sam_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sam_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sam_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sam_tx[')])
sam_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sam_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sam_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sam_nm[')])
sam_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sam_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sam_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sam_ok[')])

### Juvenile female survival
sjf_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sjf_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sjf_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sjf_tx[')])
sjf_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sjf_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sjf_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sjf_nm[')])
sjf_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sjf_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sjf_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sjf_ok[')])

### Juvenile male survival
sjm_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sjm_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sjm_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sjm_tx[')])
sjm_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sjm_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sjm_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sjm_nm[')])
sjm_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='sjm_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='sjm_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='sjm_ok[')])

### Clutch size
clu_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,10)=='clutch_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,10)=='clutch_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,10)=='clutch_tx[')])
clu_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,10)=='clutch_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,10)=='clutch_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,10)=='clutch_nm[')])
clu_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,10)=='clutch_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,10)=='clutch_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,10)=='clutch_ok[')])

### Clutch fate
fat_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,8)=='fate_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,8)=='fate_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,8)=='fate_tx[')])
fat_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,8)=='fate_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,8)=='fate_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,8)=='fate_nm[')])
fat_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,8)=='fate_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,8)=='fate_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,8)=='fate_ok[')])

### Population size
pop_tx <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='pop_tx[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='pop_tx[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='pop_tx[')])
pop_nm <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='pop_nm[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='pop_nm[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='pop_nm[')])
pop_ok <- rbind(
  fit$samples[[1]][,which(substr(row.names(fit$summary),1,7)=='pop_ok[')], 
  fit$samples[[2]][,which(substr(row.names(fit$summary),1,7)=='pop_ok[')], 
  fit$samples[[3]][,which(substr(row.names(fit$summary),1,7)=='pop_ok[')])

### Population growth rate
grw_tx <- pop_tx[,-1] / pop_tx[,-nyear]
grw_nm <- pop_nm[,-1] / pop_nm[,-nyear]
grw_ok <- pop_ok[,-1] / pop_ok[,-nyear]

# Conducting hierarchical partitioning
nmcmc <- dim(pop_tx)[1]
clu_tx_con <- fat_tx_con <- saf_tx_con <- sam_tx_con <- sjf_tx_con <- sjm_tx_con <- 
clu_nm_con <- fat_nm_con <- saf_nm_con <- sam_nm_con <- sjf_nm_con <- sjm_nm_con <- 
clu_ok_con <- fat_ok_con <- saf_ok_con <- sam_ok_con <- sjf_ok_con <- sjm_ok_con <- 
numeric(nmcmc)

for (k in 1:nmcmc) {
  y_tx <- log(grw_tx[k,])
  x1_tx <- log(clu_tx[k,-nyear])
  x2_tx <- logit(fat_tx[k,-nyear])
  x3_tx <- logit(saf_tx[k,-nyear])
  x4_tx <- logit(sam_tx[k,-nyear])
  x5_tx <- logit(sjf_tx[k,-nyear])
  x6_tx <- logit(sjm_tx[k,-nyear])
  x_tx <- cbind(x1_tx, x2_tx, x3_tx, x4_tx, x5_tx, x6_tx)
  hp_tx <- hier.part(y_tx, x_tx, family='gaussian', gof='RMSPE', barplot=T)
  clu_tx_con[k] <- hp_tx$I.perc[1,]
  fat_tx_con[k] <- hp_tx$I.perc[2,]
  saf_tx_con[k] <- hp_tx$I.perc[3,]
  sam_tx_con[k] <- hp_tx$I.perc[4,]
  sjf_tx_con[k] <- hp_tx$I.perc[5,]
  sjm_tx_con[k] <- hp_tx$I.perc[6,]
  
  y_nm <- log(grw_nm[k,])
  x1_nm <- log(clu_nm[k,-nyear])
  x2_nm <- logit(fat_nm[k,-nyear])
  x3_nm <- logit(saf_nm[k,-nyear])
  x4_nm <- logit(sam_nm[k,-nyear])
  x5_nm <- logit(sjf_nm[k,-nyear])
  x6_nm <- logit(sjm_nm[k,-nyear])
  x_nm <- cbind(x1_nm, x2_nm, x3_nm, x4_nm, x5_nm, x6_nm)
  hp_nm <- hier.part(y_nm, x_nm, family='gaussian', gof='RMSPE', barplot=T)
  clu_nm_con[k] <- hp_nm$I.perc[1,]
  fat_nm_con[k] <- hp_nm$I.perc[2,]
  saf_nm_con[k] <- hp_nm$I.perc[3,]
  sam_nm_con[k] <- hp_nm$I.perc[4,]
  sjf_nm_con[k] <- hp_nm$I.perc[5,]
  sjm_nm_con[k] <- hp_nm$I.perc[6,]
  
  y_ok <- log(grw_ok[k,])
  x1_ok <- log(clu_ok[k,-nyear])
  x2_ok <- logit(fat_ok[k,-nyear])
  x3_ok <- logit(saf_ok[k,-nyear])
  x4_ok <- logit(sam_ok[k,-nyear])
  x5_ok <- logit(sjf_ok[k,-nyear])
  x6_ok <- logit(sjm_ok[k,-nyear])
  x_ok <- cbind(x1_ok, x2_ok, x3_ok, x4_ok, x5_ok, x6_ok)
  hp_ok <- hier.part(y_ok, x_ok, family='gaussian', gof='RMSPE', barplot=T)
  clu_ok_con[k] <- hp_ok$I.perc[1,]
  fat_ok_con[k] <- hp_ok$I.perc[2,]
  saf_ok_con[k] <- hp_ok$I.perc[3,]
  sam_ok_con[k] <- hp_ok$I.perc[4,]
  sjf_ok_con[k] <- hp_ok$I.perc[5,]
  sjm_ok_con[k] <- hp_ok$I.perc[6,]

  print(k)
} # k

# Plotting the results
con_tx <- cbind(clu_tx_con, fat_tx_con, saf_tx_con, sam_tx_con, sjf_tx_con, sjm_tx_con)
con_nm <- cbind(clu_nm_con, fat_nm_con, saf_nm_con, sam_nm_con, sjf_nm_con, sjm_nm_con)
con_ok <- cbind(clu_ok_con, fat_ok_con, saf_ok_con, sam_ok_con, sjf_ok_con, sjm_ok_con)
cols <- c('tomato','lightcoral','royalblue','lightskyblue','orchid','plum')

pdf(file='4. contribution.pdf', width=7, height=7)

par(mfrow=c(1,3))
par(mar=c(0,0,4,0))
par(oma=c(1,5,0,1))

plot(1, xlim=c(.5,6.5), ylim=c(0,100), type='n', axes=F, xlab='', ylab='')
vioplot(con_tx, col=cols, border=NA, rectCol='grey16', lineCol='grey16', add=T)
axis(2, at=seq(0,100,20), las=2)
title(main='Texas', cex.main=2.4, line=.6)
legend('topleft', pch=15, col=cols, bty='n', cex=1.5, 
       legend=c('clutch size', 'clutch fate', 
                'adult female survival', 'adult male survival', 
                'juvenile female survival', 'juvenile male survival'))

plot(1, xlim=c(.5,6.5), ylim=c(0,100), type='n', axes=F, xlab='', ylab='')
vioplot(con_nm, col=cols, border=NA, rectCol='grey16', lineCol='grey16', add=T)
title(main='New Mexico', cex.main=2.4, line=.6)

plot(1, xlim=c(.5,6.5), ylim=c(0,100), type='n', axes=F, xlab='', ylab='')
vioplot(con_ok, col=cols, border=NA, rectCol='grey16', lineCol='grey16', add=T)
title(main='Oklahoma', cex.main=2.4, line=.6)

title(ylab='Independent Contribution', line=2.6, cex.lab=2.6, outer=T)

dev.off()