##########################
#Models for American green-winged teal data
#required data in "TealData.RData"
##########################
require(jagsUI)

##load in data
load("TealData.RData")

##specify number of years and crippling rate
length<-23
cr<-0.2 

##autoregressive model
sink("autoregbrownie.jags")
cat("
model {
##############################            
##Priors and constraints
############################## 
#Priors on first-year S and H (Schaub and Kery 2021)
saf[1] ~ dbeta(alphasaf,betasaf)           # Prior of survival in year 1 - AF
logit.saf[1] <- logit(saf[1])
sam[1] ~ dbeta(alphasam,betasam)           # Prior of survival in year 1 - AM
logit.sam[1] <- logit(sam[1])
sjf[1] ~ dbeta(alphasjf,betasjf)           # Prior of survival in year 1 - JF
logit.sjf[1] <- logit(sjf[1])
sjm[1] ~ dbeta(alphasjm,betasjm)           # Prior of survival in year 1 - JM
logit.sjm[1] <- logit(sjm[1])

Haf[1] ~ dbeta(1,1)               # Prior of H in year 1
logit.Haf[1] <- logit(Haf[1])
Ham[1] ~ dbeta(1,1)               # Prior of H in year 1
logit.Ham[1] <- logit(Ham[1])
Hjf[1] ~ dbeta(1,1)               # Prior of H in year 1
logit.Hjf[1] <- logit(Hjf[1])
Hjm[1] ~ dbeta(1,1)               # Prior of H in year 1
logit.Hjm[1] <- logit(Hjm[1])

# Autoregressive models for S
for (t in 2:(n.occasions-1)){
  logit.saf[t] ~ dnorm(logit.saf[t-1], tau.saf)
  saf[t] <- ilogit(logit.saf[t])
  logit.sam[t] ~ dnorm(logit.sam[t-1], tau.sam)
  sam[t] <- ilogit(logit.sam[t])
  logit.sjf[t] ~ dnorm(logit.sjf[t-1], tau.sjf)
  sjf[t] <- ilogit(logit.sjf[t])
  logit.sjm[t] ~ dnorm(logit.sjm[t-1], tau.sjm)
  sjm[t] <- ilogit(logit.sjm[t])
}

##Sigma priors for S
sigma.saf ~ dunif(0, 5)
tau.saf <- pow(sigma.saf, -2)
sigma.sam ~ dunif(0, 5)
tau.sam <- pow(sigma.sam, -2)
sigma.sjf ~ dunif(0, 5)
tau.sjf <- pow(sigma.sjf, -2)
sigma.sjm ~ dunif(0, 5)
tau.sjm <- pow(sigma.sjm, -2)

# Autoregressive models for H
for (t in 2:n.occasions){       
  logit.Haf[t] ~ dnorm(logit.Haf[t-1], tau.Haf)
  Haf[t] <- ilogit(logit.Haf[t])
  logit.Ham[t] ~ dnorm(logit.Ham[t-1], tau.Ham)
  Ham[t] <- ilogit(logit.Ham[t])
  logit.Hjf[t] ~ dnorm(logit.Hjf[t-1], tau.Hjf)
  Hjf[t] <- ilogit(logit.Hjf[t])
  logit.Hjm[t] ~ dnorm(logit.Hjm[t-1], tau.Hjm)
  Hjm[t] <- ilogit(logit.Hjm[t])
}

##Sigma priors for H
sigma.Haf ~ dunif(0, 5)
tau.Haf <- pow(sigma.Haf, -2)
sigma.Ham ~ dunif(0, 5)
tau.Ham <- pow(sigma.Ham, -2)
sigma.Hjf ~ dunif(0, 5)
tau.Hjf <- pow(sigma.Hjf, -2)
sigma.Hjm ~ dunif(0, 5)
tau.Hjm <- pow(sigma.Hjm, -2)

##Informative priors for rr estimates (Arnold et al. 2020)
for (t in 1:(n.occasions)){
  rr[t]~dbeta(alpharr[t],betarr[t])
}
   
##MULTINOMIAL LIKELIHOODS
######################################################################################
# juvenile female m-array
######################################################################################
for (t in 1:n.occasions){
    array.jf[t,1:(n.occasions+1)] ~ dmulti(pr.jf[t,], rel.jf[t])
    pr.jf[t,t] <- fjf[t]
    pr.jf[t,n.occasions+1] <- 1-sum(pr.jf[t,1:n.occasions])
    
for (j in (t+2):n.occasions){
    pr.jf[t,j] <- sjf[t]*prod(saf[(t+1):(j-1)])*faf[j]
    }
    
for (j in 1:(t-1)){
    pr.jf[t,j] <- 0
    }}
    
for (t in 1:(n.occasions-1)){
    pr.jf[t,t+1] <- sjf[t]*faf[t+1] 
    }
    
######################################################################################
# adult female m-array
######################################################################################
for (t in 1:n.occasions){
    array.af[t,1:(n.occasions+1)] ~ dmulti(pr.af[t,], rel.af[t])
    pr.af[t,n.occasions+1] <- 1-sum(pr.af[t,1:n.occasions])
    pr.af[t,t] <- faf[t]
    
for (j in (t+1):n.occasions){
    pr.af[t,j] <- prod(saf[t:(j-1)])*faf[j]
    }
    
for (j in 1:(t-1)){
    pr.af[t,j] <- 0 
    }}    
    
######################################################################################
# juvenile male m-array
######################################################################################
for (t in 1:n.occasions){
    array.jm[t,1:(n.occasions+1)] ~ dmulti(pr.jm[t,], rel.jm[t])
    pr.jm[t,t] <- fjm[t]
    pr.jm[t,n.occasions+1] <- 1-sum(pr.jm[t,1:n.occasions])
    
for (j in (t+2):n.occasions){
    pr.jm[t,j] <- sjm[t]*prod(sam[(t+1):(j-1)])*fam[j]
    }
    
for (j in 1:(t-1)){
    pr.jm[t,j] <- 0
    }}
    
for (t in 1:(n.occasions-1)){
    pr.jm[t,t+1] <- sjm[t]*fam[t+1] 
    }
    
    
######################################################################################
# adult male m-array
######################################################################################
for (t in 1:n.occasions){
    array.am[t,1:(n.occasions+1)] ~ dmulti(pr.am[t,], rel.am[t])
    pr.am[t,n.occasions+1] <- 1-sum(pr.am[t,1:n.occasions])
    pr.am[t,t] <- fam[t]
    
    for (j in (t+1):n.occasions){
    pr.am[t,j] <- prod(sam[t:(j-1)])*fam[j]
    }
    
for (j in 1:(t-1)){
    pr.am[t,j] <- 0 
    }}
    
###############################
#Derived parameters
###############################

##recovery rates (f)
for (t in 1:n.occasions){
      faf[t] = Haf[t]*rr[t]*(1-Cr)
      fam[t] = Ham[t]*rr[t]*(1-Cr)
      fjf[t] = Hjf[t]*rr[t]*(1-Cr)
      fjm[t] = Hjm[t]*rr[t]*(1-Cr)
}

##Nonhunting mortality (N)
for (t in 1:(n.occasions-1)){
      Naf[t] = 1 - saf[t] - Haf[t]
      Nam[t] = 1 - sam[t] - Ham[t]
      Njf[t] = 1 - sjf[t] - Hjf[t]
      Njm[t] = 1 - sjm[t] - Hjm[t]
}     
}
",fill = TRUE)
sink()

##gather data
jags.data <- list(array.af=array.af, array.jf=array.jf, array.am = array.am, array.jm = array.jm,
                  rel.af = rel.af, rel.jf = rel.jf, rel.am = rel.am, rel.jm = rel.jm,
                  n.occasions = length, alpharr=rr$alpha, betarr=rr$beta, Cr=cr, 
                  alphasaf = surv.priors[1,4], betasaf = surv.priors[1,5], ##moments of beta priors
                  alphasam = surv.priors[2,4], betasam = surv.priors[2,5],
                  alphasjf = surv.priors[3,4], betasjf = surv.priors[3,5], 
                  alphasjm = surv.priors[4,4], betasjm = surv.priors[4,5])

# Supply initial values
saf <- rep(NA, (length-1))
saf[1] <- runif(1, 0.3, 0.7)
Haf <- rep(NA, length)
Haf[1] <- runif(1, 0.02, 0.1)
sam <- rep(NA, (length-1))
sam[1] <- runif(1, 0.3, 0.7)
Ham <- rep(NA, length)
Ham[1] <- runif(1, 0.02, 0.1)
sjf <- rep(NA, (length-1))
sjf[1] <- runif(1, 0.3, 0.7)
Hjf <- rep(NA, length)
Hjf[1] <- runif(1, 0.02, 0.1)
sjm <- rep(NA, (length-1))
sjm[1] <- runif(1, 0.3, 0.7)
Hjm <- rep(NA, length)
Hjm[1] <- runif(1, 0.02, 0.1)
inits <- function(){list(saf=saf, Haf=Haf, sam=sam, Ham=Ham, sjf=sjf, Hjf=Hjf, sjm=sjm, Hjm=Hjm)}

# Specify parameters monitored
parameters <- c('saf','sjf','sam','sjm', 
                'faf','fjf','fam','fjm', 
                'Haf','Hjf','Ham','Hjm',
                'Naf','Njf','Nam','Njm',
                'sigma.saf', 'sigma.sam',
                'sigma.sjf', 'sigma.sjm',
                'sigma.Haf', 'sigma.Ham',
                'sigma.Hjf', 'sigma.Hjm',
                'rr')

nc<-3
nt<-10
ni<-100000
nb<-20000

##fit model
teal.results <- jags(jags.data, inits, parameters, "autoregbrownie.jags", n.chains = nc, n.thin = nt, n.iter = ni, n.burnin = nb)


##Linear model
sink("lineartimebrownie.jags")
cat("
model {
            
##Priors and constraints
##time regression model on S
for (t in 1:(n.occasions-1)){
    logit(saf[t])<-mu_saf+b.saf*t
    logit(sjf[t])<-mu_sjf+b.sjf*t
    logit(sam[t])<-mu_sam+b.sam*t
    logit(sjm[t])<-mu_sjm+b.sjm*t
  }

##priors for intercepts  
mu_saf<-logit(mean_saf)
mean_saf~dunif(0,1) 
mu_sjf<-logit(mean_sjf)
mean_sjf~dunif(0,1) 
mu_sam<-logit(mean_sam)
mean_sam~dunif(0,1)
mu_sjm<-logit(mean_sjm)
mean_sjm~dunif(0,1)
  
##priors for coefficients
b.saf~dnorm(0,0.01)
b.sjf~dnorm(0,0.01)
b.sam~dnorm(0,0.01)
b.sjm~dnorm(0,0.01)

##time regression model on H
for (t in 1:n.occasions){
    logit(Haf[t])<-mu_Haf+b.Haf*t
    logit(Hjf[t])<-mu_Hjf+b.Hjf*t
    logit(Ham[t])<-mu_Ham+b.Ham*t
    logit(Hjm[t])<-mu_Hjm+b.Hjm*t
}

##priors for intercepts    
mu_Haf<-logit(mean_Haf)
mean_Haf~dunif(0,1)
mu_Hjf<-logit(mean_Hjf)
mean_Hjf~dunif(0,1)
mu_Ham<-logit(mean_Ham)
mean_Ham~dunif(0,1)
mu_Hjm<-logit(mean_Hjm)
mean_Hjm~dunif(0,1)

##priors for coefficients
b.Haf~dnorm(0,0.01)
b.Hjf~dnorm(0,0.01)
b.Ham~dnorm(0,0.01)
b.Hjm~dnorm(0,0.01)

##informative priors for rr estimates (Arnold et al. 2020)
for (t in 1:(n.occasions)){
  rr[t]~dbeta(alpharr[t],betarr[t])
}
    
##MULTINOMIAL LIKELIHOODS
######################################################################################
# juvenile female m-array
######################################################################################
for (t in 1:n.occasions){
    array.jf[t,1:(n.occasions+1)] ~ dmulti(pr.jf[t,], rel.jf[t])
    pr.jf[t,t] <- fjf[t]
    pr.jf[t,n.occasions+1] <- 1-sum(pr.jf[t,1:n.occasions])
    
for (j in (t+2):n.occasions){
    pr.jf[t,j] <- sjf[t]*prod(saf[(t+1):(j-1)])*faf[j]
    }
    
for (j in 1:(t-1)){
    pr.jf[t,j] <- 0
    }}
    
for (t in 1:(n.occasions-1)){
    pr.jf[t,t+1] <- sjf[t]*faf[t+1] 
    }
    
######################################################################################
# adult female m-array
######################################################################################
for (t in 1:n.occasions){
    array.af[t,1:(n.occasions+1)] ~ dmulti(pr.af[t,], rel.af[t])
    pr.af[t,n.occasions+1] <- 1-sum(pr.af[t,1:n.occasions])
    pr.af[t,t] <- faf[t]
    
    for (j in (t+1):n.occasions){
    pr.af[t,j] <- prod(saf[t:(j-1)])*faf[j]
    }
    
for (j in 1:(t-1)){
    pr.af[t,j] <- 0 
    }}    
    
######################################################################################
# juvenile male m-array
######################################################################################
for (t in 1:n.occasions){
    array.jm[t,1:(n.occasions+1)] ~ dmulti(pr.jm[t,], rel.jm[t])
    pr.jm[t,t] <- fjm[t]
    pr.jm[t,n.occasions+1] <- 1-sum(pr.jm[t,1:n.occasions])
    
for (j in (t+2):n.occasions){
    pr.jm[t,j] <- sjm[t]*prod(sam[(t+1):(j-1)])*fam[j]
    }
    
for (j in 1:(t-1)){
    pr.jm[t,j] <- 0
    }}
    
for (t in 1:(n.occasions-1)){
    pr.jm[t,t+1] <- sjm[t]*fam[t+1] 
    }
    
######################################################################################
# adult male m-array
######################################################################################
for (t in 1:n.occasions){
    array.am[t,1:(n.occasions+1)] ~ dmulti(pr.am[t,], rel.am[t])
    pr.am[t,n.occasions+1] <- 1-sum(pr.am[t,1:n.occasions])
    pr.am[t,t] <- fam[t]
    
for (j in (t+1):n.occasions){
    pr.am[t,j] <- prod(sam[t:(j-1)])*fam[j]
    }
    
for (j in 1:(t-1)){
    pr.am[t,j] <- 0 
    }}
    
###############################
#Derived parameters
###############################

##recovery rates (f)
for (t in 1:n.occasions){
      faf[t] = Haf[t]*rr[t]*(1-Cr)
      fam[t] = Ham[t]*rr[t]*(1-Cr)
      fjf[t] = Hjf[t]*rr[t]*(1-Cr)
      fjm[t] = Hjm[t]*rr[t]*(1-Cr)
}

}
",fill = TRUE)
sink()

jags.data <- list(array.af=array.af, array.jf=array.jf, array.am = array.am, array.jm = array.jm,
                  rel.af = rel.af, rel.jf = rel.jf, rel.am = rel.am, rel.jm = rel.jm,
                  n.occasions = length, alpharr=rr$alpha, betarr=rr$beta, Cr=cr)

# Parameters monitored
parameters <- c('saf','sjf','sam','sjm', 
                'faf','fjf','fam','fjm', 
                'Haf','Hjf','Ham','Hjm',
                'Naf','Njf','Nam','Njm',
                'b.saf', 'b.sam',
                'b.sjf', 'b.sjm',
                'b.Haf', 'b.Ham',
                'b.Hjm', 'b.Hjf',
                'mean_saf', 'mu_saf',
                'mean_sam', 'mu_sam',
                'mean_sjf', 'mu_sjf',
                'mean_sjm', 'mu_sjm',
                'mean_Haf', 'mu_Haf',
                'mean_Ham', 'mu_Ham',
                'mean_Hjf', 'mu_Hjf',
                'mean_Hjm', 'mu_Hjm')

nc<-3
nt<-10
ni<-100000
nb<-20000

teal.results <- jags(jags.data, inits = NULL, parameters, "lineartimebrownie.jags", n.chains = nc, 
                     n.thin = nt, n.iter = ni, n.burnin = nb, parallel = TRUE)