#--------------------------------------------------------------------------------------------------------------
# American woodcock harvest model, part of AMWO IPM (Saunders et al, 2019, Ecology)
# Code written by Saunders et al. January 2017 - December 2017
# Estimating annual harvest of AMWO during duck stamp surveys, harvest information program, and parts-collection survey
# Summed harvest estimates (Eastern, Central management units) were used for Lincoln estimates of population size
#--------------------------------------------------------------------------------------------------------------

library(jagsUI)

# Load data file (included on DRUM site)
load(file="Harvest_model_data.Rda")
states <- c("WI","NH-VT","NJ-PA","CT-MA-RI","NC-SC","FL-GA","DE-MD-VA-WV","MN-IA","IL-IN-OH","LA","ME","MI","NY","S-Miss")

## data structure of Harvest_model_data.Rda
 ## all data series are in reverse time (2013 - 1964) **EXCEPT** parts data (wings.total, wings.class), which are forward time (1964-2013)
 # DSS: state*year matrix of annual harvest estimates from the Duck Stamp Survey (NA = no survey that year)
 # HIP: state*year matrix of annual harvest estimates from the Harvest Information Program (NA = no survey)
 # stamps: state*year matrix of annual duck stamp sales
 # tau.DSS: state*year matrix of precision of DSS survey (precision = 1/SD^2)
 # tau.HIP: state*year matrix of precision of HIP survey (precision = 1/SD^2)
 # wings.total: year*management unit matrix of total wings from the Parts Collection Survey (used to assign age and sex to harvest data)
   # note wings.total is 1964-2013
 # always.hunt: vector of state specific minimum duck stamp sales
 # max.stamps: vector of state-specific maximum duck stamp sales
 # mean5.HIP: vector of mean harvest per state during 2009-2013 (used as prior for starting value in reverse-time analysis)
 # n.HIP: number of years of HIP data used in analysis (1999-2013 = 15)
 # n.states: number of states (see states vector)
 # n.yrs: number of years in combined analysis (1964-2013 = 50)
 # rev: vector of seq(50,1,-1) used to reverse time series
 # sd.DSS: vector of state-specific SD of DSS surveys (process + observation error)
 # sd.HIP: vector of state-specific SD of HIP surveys (process + observation error)
 # wings.class: array of year (1964-2013) by cohort (juv, adult M, adult F) by region (Eastern, Central) wing counts from Parts Collection Survey

#############################################################
# specify jags model to predict total annual harvest
#############################################################
sink("AMWO.state.harvest.pi2.jags")
cat("
    model {
    
    # priors for state specific 2013 harvest (max of each component is 5-year mean total harvest)
    for (i in 1:n.states){
    Harvest.a[i,1] ~ dunif(0,mean5.HIP[i]) # Harvest by hunters who don't purchase duck stamps
    Harvest.b[i,1] ~ dunif(0,mean5.HIP[i]) # Harvest by hunters who always purchase duck stamps
    Harvest.c[i,1] ~ dunif(0,mean5.HIP[i]) # Harvest by hunters who sometimes purchase duck stamps
    Harvest[i,1] <- Harvest.a[i,1] + Harvest.b[i,1] + Harvest.c[i,1] # Total harvest = sum of these 3 components
    
    # priors for latent process variation in components of true harvest
    sigma.Ha[i] ~ dunif(1,20000)  
    tau.Ha[i] <- pow(sigma.Ha[i],-2)
    sigma.Hb[i] ~ dunif(1,20000)
    tau.Hb[i] <- pow(sigma.Hb[i],-2)
    sigma.Hc[i] ~ dunif(1,20000)
    tau.Hc[i] <- pow(sigma.Hc[i],-2)
    }
    
    # State process (estimate true Harvest by each segment of hunters for each region in years 2-51)
    for (i in 1:n.states){
    for (j in 2:n.yrs){
    eps.a[i,j] ~ dnorm(0,tau.Ha[i]) 
    eps.b[i,j] ~ dnorm(0,tau.Hb[i])
    eps.c[i,j] ~ dnorm(0,tau.Hc[i])
    Harvest.a[i,j] <- max(0,Harvest.a[i,j-1] + eps.a[i,j]) # autoregressive model, don't allow negative harvest
    Harvest.b[i,j] <- max(0,Harvest.b[i,j-1] + eps.b[i,j]) 
    Harvest.c[i,j] <- max(0,Harvest.c[i,j-1] + eps.c[i,j]) 
    Harvest[i,j] <- Harvest.a[i,j] + Harvest.b[i,j] + Harvest.c[i,j] # total harvest, derived parm
    }
    }
    
    # observation process, recoveries and harvest data
    for (i in 1:n.states){
    for (j in 1:n.HIP){
      HIP[i,j] ~ dnorm(Harvest[i,j],tau.HIP[i,j]) # HIP estimates total combined harvest
    } # end HIP years
    for (j in 13:n.yrs){
      frac[i,j] <- (stamps[i,j] - always.hunt[i])/(max.stamps[i] - always.hunt[i]) # estimate proportion of sometimes hunters
      DSS[i,j] ~ dnorm(Harvest.b[i,j]+frac[i,j]*Harvest.c[i,j],tau.DSS[i,j]) # Dss estimates component b and part of c
    } # end DSS yrs
    } # end i loop
    
    # calculating derived parameters of E and C harvest
    for (j in 1:n.yrs){
      Eastern.H[j] <- Harvest[2,j] + Harvest[3,j] + Harvest[4,j] + Harvest[5,j] + Harvest[6,j] + Harvest[7,j] + Harvest[11,j] + Harvest[13,j]
      Central.H[j] <- Harvest[1,j] + Harvest[8,j] + Harvest[9,j] + Harvest[10,j] + Harvest[12,j] + Harvest[14,j]
    
      region.H[j,1] <- Eastern.H[j]
      region.H[j,2] <- Central.H[j]
    }
    
    # Determine age-sex class proportions (pi's) of harvest
    for (p in 1:2){
    for (j in 1:n.yrs){
    pi[j,1,p] ~ dunif(0.1, 0.8)           # juveniles
    pi[j,2,p] ~ dunif(0.1, 0.8)           # adult males
    pi[j,3,p] <- 1-(pi[j,1,p]+pi[j,2,p])  # adult females
    
    # Summarize known wing samples by age and sex
    # populations (1 eastern, 2 central)
    wings.class[j,,p] ~ dmulti(pi[j,,p], wings.total[j,p])          # pi is 3 proportions: 1=juvs, 2=males, 3=females    
    H[j,1,p] <- pi[j,1,p]*region.H[rev[j],p]
    H[j,2,p] <- pi[j,2,p]*region.H[rev[j],p]
    H[j,3,p] <- pi[j,3,p]*region.H[rev[j],p]
    } #j
    } #p
    
    } # end jags model
    ",fill = TRUE)
sink()


# Bundle data
bugs.data <- list(n.yrs=n.yrs, n.states=n.states, n.HIP=n.HIP, mean5.HIP=mean5.HIP,
                  HIP=HIP, tau.HIP=tau.HIP, DSS=DSS, tau.DSS=tau.DSS,
                  stamps=stamps, max.stamps=max.stamps, always.hunt=always.hunt, rev=rev, wings.total=wings.total, wings.class=wings.class)

# Initial values (not all priors listed)
pi.inits <- array(dim=c(n.yrs,3,2))
for (c in 1:3){
  for (p in 1:2){
    for (j in 1:n.yrs){
      pi.inits[j,1,p] <- 0.46
      pi.inits[j,2,p] <- 0.21
      pi.inits[j,3,p] <- NA
    }}}

inits <- function(){list(pi=pi.inits)}

# Parameters monitored
parameters <- c("Harvest","Harvest.a","Harvest.b","Harvest.c", "Eastern.H","Central.H","pi","H")

# MCMC settings (about 2 hours at these settings)
na <- 1000
ni <- 150000
nt <- 10        
nb <- 50000
nc <- 3

# call jags
AMWO.state.harvest.pi <- jagsUI(bugs.data, inits, parameters, "AMWO.state.harvest.pi2.jags", n.adapt = na, n.chains = nc, n.thin = nt, n.iter = ni, n.burnin = nb, parallel=TRUE)
save(AMWO.state.harvest.pi, file="AMWO_harvest_pi.Rda")

max(max.rhat<-c(sapply(AMWO.state.harvest.pi$Rhat,max,na.rm=TRUE))) # report max rhat of all posteriors

# generate plot of total estimated harvest in each management unit
plot(1964:2013-0.2, rev(AMWO.state.harvest.pi$mean$Central.H),ylim=c(0,1200000),pch=19,type="b",
     col="blue",ylab="Annual harvest",xlab="")
segments(1964:2013-0.2, rev(AMWO.state.harvest.pi$mean$Central.H-1.96*AMWO.state.harvest.pi$sd$Central.H),
         1964:2013-0.2, rev(AMWO.state.harvest.pi$mean$Central.H+1.96*AMWO.state.harvest.pi$sd$Central.H),col="blue",lty=2)
points(1964:2013+0.2, rev(AMWO.state.harvest.pi$mean$Eastern.H), pch=19, col="red", type="b")
segments(1964:2013+0.2, rev(AMWO.state.harvest.pi$mean$Eastern.H-1.96*AMWO.state.harvest.pi$sd$Eastern.H),
         1964:2013+0.2, rev(AMWO.state.harvest.pi$mean$Eastern.H+1.96*AMWO.state.harvest.pi$sd$Eastern.H),col="red",lty=2)
legend(1995,1000000,legend=c("Central Population","Eastern Population"),text.col = c("blue","red"),pch=c(19,19),col=c("blue","red"),bty="n")

# generate plot for any given state by setting j to desired value
# these are the 14 state groups used for analysis in Saunders et al. 2019
states <- c("WI","NH-VT","NJ-PA","CT-MA-RI","NC-SC","FL-GA","DE-MD-VA-WV","MN-IA","IL-IN-OH","LA","ME","MI","NY","S-Miss")
j = 1 # Wisconsin
plot(1964:2013-0.2, rev(AMWO.state.harvest.pi$mean$Harvest[j,]),ylim=c(0,2*max(AMWO.state.harvest.pi$mean$Harvest[j,])),pch=19,type="b",
     col="blue",ylab="Annual harvest",xlab="", main=states[j])
segments(1964:2013-0.2, rev(AMWO.state.harvest.pi$mean$Harvest[j,]-1.96*AMWO.state.harvest.pi$sd$Harvest[j,]),
         1964:2013-0.2, rev(AMWO.state.harvest.pi$mean$Harvest[j,]+1.96*AMWO.state.harvest.pi$sd$Harvest[j,]),col="blue",lty=2)

# plot proportion of juveniles [,1,1] = Eastern, [,1,2] = Central
 # Note: these are not yet corrected for vulnerability (see Zimmerman et al. 2010, Saunders et al. 2019)
plot(1964:2013-0.2, AMWO.state.harvest.pi$mean$pi[,1,1],ylim=c(0.35,0.65), pch=19, type="b",
     col="blue",ylab="Proportion of juveniles",xlab="",main="Harvest age ratios")
segments(1964:2013-0.2, AMWO.state.harvest.pi$mean$pi[,1,1]-1.96*AMWO.state.harvest.pi$sd$pi[,1,1],
         1964:2013-0.2, AMWO.state.harvest.pi$mean$pi[,1,1]+1.96*AMWO.state.harvest.pi$sd$pi[,1,1],col="blue",lty=1)
points(1964:2013+0.2, AMWO.state.harvest.pi$mean$pi[,1,2],ylim=c(0,1), pch=19, type="b", col="red")
segments(1964:2013+0.2, AMWO.state.harvest.pi$mean$pi[,1,2]-1.96*AMWO.state.harvest.pi$sd$pi[,1,2],
         1964:2013+0.2, AMWO.state.harvest.pi$mean$pi[,1,2]+1.96*AMWO.state.harvest.pi$sd$pi[,1,2],col="red",lty=1)
legend(1990, 0.65, legend=c("Eastern","Central"), text.col = c("blue","red"), bty="n", pch=c(19,19), col=c("blue","red"),
       lty=c(1,1))

# write model-based estimates and SD (equivalent to frequentist SE) to files
mean.harvest2 <- sd.harvest2 <- matrix(NA, nrow=n.yrs, ncol=18)

for (j in 1:n.yrs){
  for (i in 1:14){
    mean.harvest2[j,i+1] <- round(AMWO.state.harvest.pi$mean$Harvest[i,j],1)
    sd.harvest2[j,i+1] <- round(AMWO.state.harvest.pi$sd$Harvest[i,j],1)
  }
  mean.harvest2[j,15] <- mean.harvest2[j,1] + sum(mean.harvest2[j,8:10]) + mean.harvest2[j,12] + mean.harvest2[j,14] # central
  mean.harvest2[j,16] <- sum(mean.harvest2[j,2:7]) + mean.harvest2[j,11] + mean.harvest2[j,13] # eastern
  mean.harvest2[j,17] <- sum(mean.harvest2[j,15:16]) # total
  sd.harvest2[j,15] <- sqrt(sd.harvest2[j,1]^2 + sd.harvest2[j,8]^2 + sd.harvest2[j,9]^2 + sd.harvest2[j,10]^2 + 
                             sd.harvest2[j,12]^2 + sd.harvest2[j,14]^2)
  sd.harvest2[j,16] <- sqrt(sd.harvest2[j,2]^2 + sd.harvest2[j,3]^2 + sd.harvest2[j,4]^2 + sd.harvest2[j,5]^2 + 
                             sd.harvest2[j,6]^2 + sd.harvest2[j,7]^2 + sd.harvest2[j,11]^2 + sd.harvest2[j,13]^2)
  sd.harvest2[j,17] <- sqrt(sd.harvest2[j,15]^2 + sd.harvest2[j,16]^2)
}

colnames(mean.harvest2) <- c("Year",states,"Central","Eastern","Total")
colnames(sd.harvest2) <- c("Year",states,"Central","Eastern","Total")
mean.harvest2[1:n.yrs,1] <- seq(2013,1964,-1)
sd.harvest2[1:n.yrs,1] <- seq(2013,1964,-1)

write.csv(mean.harvest2,file="Saunders_harvest_estimates.csv")
write.csv(sd.harvest2,file="Saunders_harvest_estimates_SD.csv")