rm(list=ls())
library(jagsUI)
library(boot)

species <- 'ross' # snow or ross

#==============
# Read in data
#==============
# Basic
years <- 1997:2017
nyear <- length(years)
yr_st <- years[-nyear]
yr_st <- (yr_st - mean(yr_st)) / (max(yr_st) - min(yr_st)) * 4

# Band-recapture-recovery data in M-array
load(paste(c('data/marr_', species, '.RData'), collapse=''))

interlace <- function (n)
{
	interlace.idx <- rep(1:n, each = 2) + rep(c(0,n), n) # interlaced index!! :-) This is R fun :-)
	return(interlace.idx)
}

mAd <-  cbind(
  cbind(marr$A[,-nyear], marr$A.dead)[,interlace(nyear-1)], # interlace alive with dead
  marr$A[,nyear] # add last column (never captured again)
)

mJuv <- cbind(
  cbind(marr$J[,-nyear], marr$J.dead)[,interlace(nyear-1)], # interlace alive with dead
  marr$J[,nyear] # add last column (never captured again)
) 

# Bpop & Age ratio
data <- read.csv('data/pop size age ratio.csv')

if (species == 'snow') {
  bpop_obs <- data$TOTSNOW[which(data$YEAR %in% years)] / 1000
  bpop_se <- data$TOTSNOWCL[which(data$YEAR %in% years)] / 1000
} else if (species == 'ross') {
  bpop_obs <- data$TOTROSS[which(data$YEAR %in% years)] / 1000
  bpop_se <- data$TOTROSSCL[which(data$YEAR %in% years)] / 1000
} else {
  print('wrong species')
}
bpop_tau <- 1 / (bpop_se ^ 2)
lambda1 <- bpop_obs[1]
lambda2 <- bpop_obs[2]

years2 <- years[-length(years)]

if (species == 'snow') {
  cap_young <- data$SNOWFLEDGEHY[which(data$YEAR %in% years2)]
  cap_adult <- data$SNOWFLEDGEAHY[which(data$YEAR %in% years2)]
} else if (species == 'ross') {
  cap_young <- data$ROSSFLEDGEHY[which(data$YEAR %in% years2)]
  cap_adult <- data$ROSSFLEDGEAHY[which(data$YEAR %in% years2)]
} else {
  print('wrong species')
}
cap_all <- cap_young + cap_adult

#======================
# Define model in Jags
#======================
sink('ipm_trend_marr_truncated_normal.txt')
cat("
model {

#------Survival (Multistate Burnham) Model------
# Parameters:
# sj[t]: probability that juvenile survives from year t to t+1, 
# Fj[t]: probability that juvenile doesn't emigrate between years t and t+1 (we denote t+1 as the year of emigration)
# rj[t]: probability that juvenile death occuring between years t and t+1 is reported (we denote t+1 as the year of death)
# sa[t]: probability that adult survives from year t to t+1
# Fa[t]: probability that adult doesn't emigrate between years t and t+1 (we denote t+1 as the year of emigration)
# ra[t]: probability that adult death occuring between years t and t+1 is reported (we denote t+1 as the year of death)
# p[t]:  probability that bird present at site in year t is captured in year t
# -------------------------------------------------
# States (S):
# 1 juvenile alive in study area
# 2 juvenile alive outside study area (emigrated)
# 3 dead, recovered and reported as juvenile
# 4 adult alive in study area
# 5 adult alive outside study area (emigrated)
# 6 dead, recovered and reported as adult
# 7 not seen or recovered and reported
#
# Observations (O):
# 1 juvenile alive in study area
# 2 juvenile alive outside study area
# 3 dead, recovered and reported as juvenile
# 4 adult alive in study area
# 5 adult alive outside study area
# 6 dead, recovered and reported as adult
# 7 not seen or recovered and reported
# -------------------------------------------------

# Priors
mu.sj ~ dnorm(0, .01)
mean.sj <- ilogit(mu.sj)
mu.Fj ~ dnorm(0, .01)
mean.Fj <- ilogit(mu.Fj)
mu.rj ~ dnorm(0, .01)
mean.rj <- ilogit(mu.rj)
mu.sa ~ dnorm(0, .01)
mean.sa <- ilogit(mu.sa)
mu.Fa ~ dnorm(0, .01)
mean.Fa <- ilogit(mu.Fa)
mu.ra ~ dnorm(0, .01)
mean.ra <- ilogit(mu.ra)
mu.p ~ dnorm(0, .01)
mean.p <- ilogit(mu.p)

beta.sj ~ dnorm(0, .01)
beta.Fj ~ dnorm(0, .01)
beta.rj ~ dnorm(0, .01)
beta.sa ~ dnorm(0, .01)
beta.Fa ~ dnorm(0, .01)
beta.ra ~ dnorm(0, .01)
beta.p ~ dnorm(0, .01)

tau.sj ~ dgamma(.01, .01)
sd.sj <- 1 / sqrt(tau.sj)
tau.Fj ~ dgamma(.01, .01)
sd.Fj <- 1 / sqrt(tau.Fj)
tau.rj ~ dgamma(.01, .01)
sd.rj <- 1 / sqrt(tau.rj)
tau.sa ~ dgamma(.01, .01)
sd.sa <- 1 / sqrt(tau.sa)
tau.Fa ~ dgamma(.01, .01)
sd.Fa <- 1 / sqrt(tau.Fa)
tau.ra ~ dgamma(.01, .01)
sd.ra <- 1 / sqrt(tau.ra)
tau.p ~ dgamma(.01, .01)
sd.p <- 1 / sqrt(tau.p)

for (t in 1:(nyear-1)) { 
  logit.sj.mu[t] <- mu.sj + beta.sj * yr_st[t]
  logit.sj[t] ~ dnorm(logit.sj.mu[t], tau.sj)
  sj[t] <- ilogit(logit.sj[t])
  logit.Fj.mu[t] <- mu.Fj + beta.Fj * yr_st[t]
  logit.Fj[t] ~ dnorm(logit.Fj.mu[t], tau.Fj)
  Fj[t] <- ilogit(logit.Fj[t])
  logit.rj.mu[t] <- mu.rj + beta.rj * yr_st[t]
  logit.rj[t] ~ dnorm(logit.rj.mu[t], tau.rj)
  rj[t] <- ilogit(logit.rj[t])
  logit.sa.mu[t] <- mu.sa + beta.sa * yr_st[t]
  logit.sa[t] ~ dnorm(logit.sa.mu[t], tau.sa)
  sa[t] <- ilogit(logit.sa[t])
  logit.Fa.mu[t] <- mu.Fa + beta.Fa * yr_st[t]
  logit.Fa[t] ~ dnorm(logit.Fa.mu[t], tau.Fa)
  Fa[t] <- ilogit(logit.Fa[t])
  logit.ra.mu[t] <- mu.ra + beta.ra * yr_st[t]
  logit.ra[t] ~ dnorm(logit.ra.mu[t], tau.ra)
  ra[t] <- ilogit(logit.ra[t])
  logit.p.mu[t] <- mu.p  + beta.p  * yr_st[t]
  logit.p[t] ~ dnorm(logit.p.mu[t], tau.p)
  p[t+1] <- ilogit(logit.p[t])
}

# Likelihood for M-array
# M-array implementation of the Multistate Burnham model (joint live and dead encounter model)
#
# Since states 2, 5, and 7 are not observable, and states 3 and 6 (dead ad or juv) have only one possible continuation, 
# we need m-arrays only for states 1 (juv, alive, inside) and 4 (adult, alive, inside). In these m-arrays, we need two 
# numbers for two different output states - alive or dead.
#
# Input data in the m-array format are defined like this:
#
# mJuv[i,2*(j-1) - 1] .....	for i=1..nyear-1, j=2..nyear, number of birds banded as juveniles in year i, and next seen captured 
#							in year j as adults (next observable state = 4)
# mJuv[i,2*(j-1)    ] ..... for i=1..nyear-1, j=2..nyear, number of birds banded as juveniles in year i, and next seen 
#							in year j as dead (next observable state = dead juv or ad, 3 or 6 - I don't think you need two states for this BTW... 
#							3 and 6 could be merged?)
# mJuv[i,2*nyear - 1] ..... for i=1..nyear-1,  number of birds banded as juveniles in year i and never seen again
#
# mAd[i,2*(j-1) - 1] ......	for i=1..nyear-1, j=2..nyear, number of birds captured as adults in year i, and next seen captured 
#							in year j
# mAd[i,2*(j-1)    ] ...... for i=1..nyear-1, j=2..nyear, number of birds captured as adults in year i, and next seen 
#							in year j as dead (next observable state = dead ad = 6)
# mAd[i,2*nyear - 1] ...... for i=1..nyear-1,  number of birds captured as adults in year i and never seen again
#
# Row sums of these matrices contain total captured ad/juv in a given year.
#
# For each cell of these m-arrays, probabilities are defined in matrices pAd[] and pJuv[]. The rows of these matrices sum up to 1.

for (i in 1:(nyear-1)) { # row index - number of birds captured in year i, and next ....
  ### below the diagonal - zeros
  for (j in 2:i) {
    pAd[i, 2*(j-1)-1] <- 0
    pAd[i, 2*(j-1)] <- 0
    pJuv[i, 2*(j-1)-1] <- 0
    pJuv[i, 2*(j-1)] <- 0		
  } # j

  ### the diagonal: next year (j = i+1)
  pAd [i, 2*((i+1)-1) - 1] <- Fa[i]*sa[i]*p[i+1] # captured alive next year 
  pAd [i, 2*((i+1)-1)    ] <- (1 - sa[i])*ra[i]  # reported dead next year (couldn't have emigrated in this case)
  pJuv[i, 2*((i+1)-1) - 1] <- Fj[i]*sj[i]*p[i+1] # captured alive next year 
  pJuv[i, 2*((i+1)-1)    ] <- (1 - sj[i])*rj[i]  # reported dead next year (couldn't have emigrated in this case)

  ### j = i+2, if it is <= nyear
  for (j in (i+2):min(i+2,nyear)) { # next captured/recovered in year j = i+2, if i+2 <= nyear
	
    # probability that bird captured as adult/juvenile in year i is next captured as adult in year j
    pAd [i, 2*(j-1)-1] <-         prod(Fa[ i   :(j-1)]) * prod(1 - p[(i+1):(j-1)])         * prod(sa[ i   :(j-1)]) * p[j]
    pJuv[i, 2*(j-1)-1] <- Fj[i] * prod(Fa[(i+1):(j-1)]) * prod(1 - p[(i+1):(j-1)]) * sj[i] * prod(sa[(i+1):(j-1)]) * p[j]
		
    # probability that bird captured as adult/juvenile in year i is next seen as dead in year j:
    # sum up over all posibilities of emigration (latent state). 
    # Emigration possible in years i+1 to j-1 (since the bird cannot emigrate in the year of death = j)
    # Recommendation: watch the code with fixed font, all the expressions are aligned with the most complicated case below (emigration in year >= i + 3)
		
    # emigration in year i+1
    for (e in (i+1)) {
      pAd.emig [i, 2*(j-1), e] <- (1-Fa[e-1]) * prod(sa[i:(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
      pJuv.emig[i, 2*(j-1), e] <- (1-Fj[e-1]) * sj[i]             * (1 - sa[j-1]) * ra[j-1] 				
    } # e
				
    # no emigration at all (faithful till year j-1, in year j faithfulness/emigration isn't defined since it's year of death)
    # for simplicity, we index it as emigration in year j
    for (e in j) {
      pAd.emig [i, 2*(j-1), e] <- prod(Fa[i:(e-2)]) * prod(1 - p[(i+1):(e-1)]) * prod(sa[i:(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
      pJuv.emig[i, 2*(j-1), e] <- Fj[i]             * prod(1 - p[(i+1):(e-1)]) * sj[i]             * (1 - sa[j-1]) * ra[j-1] 
    } # e
				
    # now, sum up over all posibilities of emigration:
    # probability that bird captured as adult in year i is next seen as dead in year j:
    pAd [i, 2*(j-1)] <- sum(pAd.emig [i, 2*(j-1), (i+1):j]) 
    pJuv[i, 2*(j-1)] <- sum(pJuv.emig[i, 2*(j-1), (i+1):j]) 
  } # j

  ### j = i+3 up to nyear
  for (j in (i+3):nyear) { # next captured/recovered in year j
	
    # probability that bird captured as adult/juvenile in year i is next captured as adult in year j
    pAd [i, 2*(j-1)-1] <-         prod(Fa[ i   :(j-1)]) * prod(1 - p[(i+1):(j-1)])         * prod(sa[ i   :(j-1)]) * p[j]
    pJuv[i, 2*(j-1)-1] <- Fj[i] * prod(Fa[(i+1):(j-1)]) * prod(1 - p[(i+1):(j-1)]) * sj[i] * prod(sa[(i+1):(j-1)]) * p[j]

    # probability that bird captured as adult/juvenile in year i is next seen as dead in year j:
    # sum up over all posibilities of emigration (latent state). 
    # Emigration possible in years i+1 to j-1 (since the bird cannot emigrate in the year of death = j)
    # Recommendation: watch the code with fixed font, all the expressions are aligned with the most complicated case below (emigration in year >= i + 3)
		
    # emigration in year i+1
    for (e in (i+1)) {
      pAd.emig [i, 2*(j-1), e] <- (1-Fa[e-1])         * prod(sa[ i   :(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
      pJuv.emig[i, 2*(j-1), e] <- (1-Fj[e-1]) * sj[i] * prod(sa[(i+1):(j-2)]) * (1 - sa[j-1]) * ra[j-1] 				
    } # e

    # emigration in year i+2, if i+2 <= j-1
    for (e in (i+2):min(i+2,j-1)) {
      pAd.emig [i, 2*(j-1), e] <- prod(Fa[i:(e-2)]) * (1-Fa[e-1]) * prod(1 - p[(i+1):(e-1)])         * prod(sa[ i   :(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
      pJuv.emig[i, 2*(j-1), e] <- Fj[i]             * (1-Fa[e-1]) * prod(1 - p[(i+1):(e-1)]) * sj[i] * prod(sa[(i+1):(j-2)]) * (1 - sa[j-1]) * ra[j-1] 			
    } # e

    # emigration in years i+3 to j-1
    for (e in (i+3):(j-1)) {
      pAd.emig [i, 2*(j-1), e] <-         prod(Fa[ i     :(e-2)]) * (1-Fa[e-1]) * prod(1 - p[(i+1):(e-1)])         * prod(sa[ i   :(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
      pJuv.emig[i, 2*(j-1), e] <- Fj[i] * prod(Fa[(i + 1):(e-2)]) * (1-Fa[e-1]) * prod(1 - p[(i+1):(e-1)]) * sj[i] * prod(sa[(i+1):(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
    } # e, i + 1 <= e - 2 => e >= i + 3		

    # no emigration at all (faithful till year j-1, in year j faithfulness/emigration isn't defined since it's year of death)
    # for simplicity, we index it as emigration in year j
    for (e in j) {
      pAd.emig [i, 2*(j-1), e] <-         prod(Fa[ i     :(e-2)]) * prod(1 - p[(i+1):(e-1)])         * prod(sa[ i   :(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
      pJuv.emig[i, 2*(j-1), e] <- Fj[i] * prod(Fa[(i + 1):(e-2)]) * prod(1 - p[(i+1):(e-1)]) * sj[i] * prod(sa[(i+1):(j-2)]) * (1 - sa[j-1]) * ra[j-1] 
    } # e
		
    # now, sum up over all posibilities of emigration:
    # probability that bird captured as adult in year i is next seen as dead in year j:
    pAd [i, 2*(j-1)] <- sum(pAd.emig [i, 2*(j-1), (i+1):j]) 
    pJuv[i, 2*(j-1)] <- sum(pJuv.emig[i, 2*(j-1), (i+1):j]) 
  } # j

  # probability that bird captured as adult/juvenile in year i is never seen again
  pAd [i, 2*nyear - 1] <- 1 - sum(pAd [i, 1:(2*nyear - 2)])
  pJuv[i, 2*nyear - 1] <- 1 - sum(pJuv[i, 1:(2*nyear - 2)])

  # now, use multinomial distribution to match probabilities with data!
  mAd [i,] ~ dmulti(pAd [i,], mAd.sums [i])
  mJuv[i,] ~ dmulti(pJuv[i,], mJuv.sums[i])
} # i

#--------------Productivity Model------------------
# Parameters:
# ar: age ratio
# q: proportion of juvenile females
# -------------------------------------------------

# Priors
mu.ar ~ dnorm(0, .01)
mean.ar <- exp(mu.ar)
beta.ar ~ dnorm(0, .01)
tau.ar ~ dgamma(.01, .01)
sd.ar <- 1 / sqrt(tau.ar)

for (t in 1:(nyear-1)){
  log.ar.mu[t] <- mu.ar + beta.ar * yr_st[t]
  log.ar[t] ~ dnorm(log.ar.mu[t], tau.ar)
  ar[t] <- exp(log.ar[t])
  q[t] <- ar[t] / (1 + ar[t])
}

# Likelihood for age-ratio data
for (t in 1:(nyear-1)) {
  cap_young[t] ~ dbinom(q[t], cap_all[t])
}

#--------------Population Model------------------
# Parameters:
# bpop: true breeding population size
# sd.bpop: SD for process error
# -----------------------------------------------

# Priors
lambda.i ~ dgamma(.01, .01)

# First-year population size
Nj[1] ~ dnorm(lambda1 / 6 * 1, 1 / (lambda1 / 6 * 1))T(0,)
Na[1] ~ dnorm(lambda1 / 6 * 4, 1 / (lambda1 / 6 * 4))T(0,)
Ni[1] ~ dnorm(lambda1 / 6 * 1, 1 / (lambda1 / 6 * 1))T(0,)
N[1] <- Nj[1] + Na[1] + Ni[1]

# Second-year population size
Nj[2] ~ dnorm(lambda2 / 6 * 1, 1 / (lambda2 / 6 * 1))T(0,)
Na[2] ~ dnorm(lambda2 / 6 * 4, 1 / (lambda2 / 6 * 4))T(0,)
Ni[2] ~ dnorm(lambda2 / 6 * 1, 1 / (lambda2 / 6 * 1))T(0,)
N[2] <- Nj[2] + Na[2] + Ni[2]

# Population size from the third year
for (t in 3:nyear) {
  Nj[t] ~ dnorm(N[t-2] * ar[t-2] * sj[t-2] * Fj[t-2] * sa[t-1] * Fa[t-1], 1 / (N[t-2] * ar[t-2] * sj[t-2] * Fj[t-2] * sa[t-1] * Fa[t-1]))T(0,)
  Na[t] ~ dnorm(N[t-1] * sa[t-1] * Fa[t-1], 1 / (N[t-1] * sa[t-1] * Fa[t-1] * (1 - sa[t-1] * Fa[t-1])))T(0,)
  Ni[t] ~ dnorm(lambda.i, 1 / lambda.i)T(0,)
  N[t] <- Nj[t] + Na[t] + Ni[t]
}

# likelihood for population survey data
for (t in 1:nyear) {
  bpop_obs[t] ~ dnorm(N[t], bpop_tau[t])
}

} # model

",fill = TRUE)
sink()

#==========
# Run Jags
#==========
# Bundle data
data <- list(mAd = mAd, mJuv = mJuv, 
	       mAd.sums = rowSums(mAd), mJuv.sums = rowSums(mJuv),
	       nyear = nyear, yr_st = yr_st, 
             cap_young = cap_young, cap_all = cap_all, 
             bpop_obs = bpop_obs, bpop_tau = bpop_tau, 
             lambda1 = lambda1, lambda2 = lambda2)

# Initial values 
inits <- function() {
  list(
    mu.sj=logit(0.5), mu.Fj=logit(0.5), mu.rj=logit(0.1), 
    mu.sa=logit(0.9), mu.Fa=logit(0.9), mu.ra=logit(0.1), 
    mu.p=logit(0.02), 
    beta.sj=0, beta.Fj=0, beta.rj=0, 
    beta.sa=0, beta.Fa=0, beta.ra=0, 
    beta.p=0, 
    tau.sj=1, tau.Fj=1, tau.rj=1, 
    tau.sa=1, tau.Fa=1, tau.ra=1, 
    tau.p=1, 
    logit.sj=rep(.5,nyear-1), logit.Fj=rep(.5,nyear-1), logit.rj=rep(.1,nyear-1), 
    logit.sa=rep(.9,nyear-1), logit.Fa=rep(.9,nyear-1), logit.ra=rep(.1,nyear-1), 
    logit.p=rep(.02,nyear-1), 
    mu.ar=log(.3), beta.ar=0, tau.ar=1, 
    log.ar=rep(1,nyear-1), 
    lambda.i=10)
}

# Parameters monitored
parms <- c(
  'mean.sa', 'mean.Fa', 'mean.ra', 
  'mean.sj', 'mean.Fj', 'mean.rj', 
  'mean.p', 'mean.ar', 
  'beta.sa', 'beta.Fa', 'beta.ra', 
  'beta.sj', 'beta.Fj', 'beta.rj', 
  'beta.p', 'beta.ar', 
  'sd.sa', 'sd.Fa', 'sd.ra', 
  'sd.sj', 'sd.Fj', 'sd.rj', 
  'sd.p', 'sd.ar', 
  'sa', 'Fa', 'ra', 
  'sj', 'Fj', 'rj', 
  'p', 'ar', 
  'N', 'Nj', 'Na', 'Ni', 
  'lambda.i'
)

# MCMC settings
fit <- jags(data, inits, parms, 'ipm_trend_marr_truncated_normal.txt', 
            n.chains=3, n.adapt=10000, n.burnin=250000, n.iter=500000, n.thin=50, 
#            n.chains=1, n.adapt=100, n.burnin=100, n.iter=200, n.thin=1, 
            parallel=TRUE)

save(fit, file=paste(c('ipm_trend_marr_truncated_normal_', species, '.RData'), collapse=''))

print(fit)
par(mfrow=c(5,5))
par(mar=c(2,2,2,1))
traceplot(fit, 
  parameters=c('mean.sa', 'mean.Fa', 'mean.ra', 
               'mean.sj', 'mean.Fj', 'mean.rj', 
               'mean.p', 'mean.ar', 
               'beta.sa', 'beta.Fa', 'beta.ra', 
               'beta.sj', 'beta.Fj', 'beta.rj', 
               'beta.p', 'beta.ar', 
               'sd.sa', 'sd.Fa', 'sd.ra', 
               'sd.sj', 'sd.Fj', 'sd.rj', 
               'sd.p', 'sd.ar', 
               'lambda.i'))