# This software has been approved for release by the U.S. Geological Survey (USGS). 
# Although the software has been subjected to rigorous review, the USGS reserves 
# the right to update the software as needed pursuant to further analysis and 
# review. No warranty, expressed or implied, is made by the USGS or the U.S. 
# Government as to the functionality of the software and related material nor 
# shall the fact of release constitute any such warranty. Furthermore, the 
# software is released on condition that neither the USGS nor the U.S. Government 
# shall be held liable for any damages resulting from its authorized or 
# unauthorized use.
#
# Supplemental material for:
# Cotterill et al., 2020
# Parsing the effects of demography, climate, and management on recurrent brucellosis outbreaks in elk
# gavin.cotterill@aggiemail.usu.edu

rm(list = ls())
dev.off()

setwd(dirname(rstudioapi::getActiveDocumentContext()$path))
set.seed(219)

library(tidyverse)

#' In order to recreate the full figure, you have to run this script 9 times:
#' for model m0, m1, m2, sites 1-3.

# choose model: "m0" = internal, "m1" = combined, "m2" = external
mod <- "m0"

# choose site: 1, 2, or 3
site <- 1

greys <- read.csv("../data/covar.csv", header = T)%>%
  filter(Feedground == "Greys River")

gg.gavin.theme <- function(textsize = 40){
  theme(axis.line = element_line(colour = "black"),
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank(),
        panel.border = element_blank(),
        panel.background = element_blank(),
        legend.title = element_text(size = textsize*0.5, face = "bold"),
        legend.text = element_text(size = textsize*0.5),
        legend.key = element_rect(color = "transparent", fill = "white"),
        legend.key.height = unit(3, "line"),
        legend.key.width = unit(3, "line"),
        legend.key.size = unit(1, "cm"),
        plot.title = element_text(size = textsize*0.7),
        axis.title.x = element_text(size = textsize*0.7, face = "bold",
                                    margin = margin(t = 40)),
        axis.title.y = element_text(size = textsize*.7, face = "bold",
                                    margin = margin(r = 40)),
        axis.text.x = element_text(size = 30*0.8), 
        axis.text.y = element_text(size = 30*0.8), 
        plot.margin = unit(c(1, 1, 0, 1.2), "cm"))
}

## settings
n.years <- 25 * 4
swe <- sample(greys$MJ, 100, replace=T)


# transmission terms vary by model, all taken from Greys River MLEs
# m0 vals:
beta0 <- 0.02 # transmission beta
theta0 <- 0.18 # freq/dens
iota0 <- 0.11 # imported infections
sigma0 <- 0.95 # I to R1 recovery
rho0 <- 0.51 # S to R1 (never abort)
gamma0 <- 0.10 # R1 to R2 seroreversion

# m1 vals:
beta1 <- 0.23 # transmission beta
bp1 <- 0.04 # climate beta prime
theta1 <- 0.63 # freq/dens
iota1 <- 0.35 # imported infections
sigma1 <- 0.93 # I to R1 recovery
rho1 <- 0.49 # S to R1 (never abort)
gamma1 <- 0.12 # R1 to R2 seroreversion

# m2 vals:
bp2 <- 0.015 # climate beta prime
theta2 <-  0.29 # freq/dens
iota2 <- 0.41 # imported infections
sigma2 <- 0.33 # I to R1 recovery
rho2 <- 0.49 # S to R1 (never abort)
gamma2 <- 0.10 # R1 to R2 seroreversion


# constants:
n.0 <- 400 # starting pop size
# starting values
S.0 <- 0.75
I.0 <- 0.01
R1.0 <- 0.15
R2.0 <- 0.09

if(site == 1){
  mu <- 1/8 # death
}else if(site == 2){
  mu <- 1/9
}else if(site == 3){
  mu <- 1/10
}


# to go from starting proportions (they don't need to sum to one) to integers:
m <- n.0 /(S.0 + I.0 + R1.0 + R2.0)
S.0 <- round(m * S.0, 0)
I.0 <- round(m * I.0, 0)
R1.0 <- round(m * R1.0, 0)
R2.0 <- round(m* R2.0, 0)

# containers for output
S <- c(S.0, rep(0, n.years - 1))
I <- c(I.0, rep(0, n.years - 1))
R1 <- c(R1.0, rep(0, n.years - 1))
R2 <- c(R2.0, rep(0, n.years - 1))
N <- c(n.0, rep(0, n.years - 1))

births <- rep(0, n.years - 1)
foi <- rep(0, n.years - 1)
new.inf <- rep(0, n.years - 1)
trans.I <- rep(0, n.years - 1)
trans.R1 <- rep(0, n.years - 1)
recover <- rep(0, n.years - 1)
revert <- rep(0, n.years - 1)


for(t in 2:n.years){

  births[t] <- (S[t-1] + I[t-1] + R1[t-1] + R2[t-1]) * runif(1,0.1,0.4) / 2 # halved for female-only

  if(mod == "m0"){
    iota <- rbinom(n = 1, size = 1, prob = iota0)
    foi[t-1] <- beta0 * (I[t-1] + iota)/(N[t-1]^theta0)
    new.inf[t] <- pomp2::reulermultinom(2, S[t-1], foi[t-1], 1)[1]
    trans.I[t] <- as.integer(round(new.inf[t] * (1 - rho0) * (1 - mu), 0))
    trans.R1[t] <- as.integer(round(new.inf[t] * rho0 * (1 - mu), 0))
    recover[t] <- as.integer(round(I[t-1] * sigma0 * (1 - mu), 0))
    revert[t] <- as.integer(round(R1[t-1] * gamma0 * (1 - mu), 0))
    S[t] <- round((S[t-1] * (1 - mu)) - new.inf[t] + births[t], 0)
    I[t] <- round((I[t-1] * (1 - mu) * (1 - sigma0) + trans.I[t]), 0)
    R1[t] <- round(R1[t-1] * (1 - mu) * (1 - gamma0) + trans.R1[t] + recover[t], 0)
    R2[t] <- round(R2[t-1] * (1 - mu) + revert[t], 0)
    N[t] <- S[t] + I[t] + R1[t] + R2[t]
  }else if(mod == "m1"){
    iota <- rbinom(n = 1, size = 1, prob = iota1)
    foi[t-1] <- (beta1 + (bp1 * swe[t-1])) * (I[t-1] + iota)/(N[t-1]^theta1)
    new.inf[t] <- pomp2::reulermultinom(2, S[t-1], foi[t-1], 1)[1]
    trans.I[t] <- as.integer(round(new.inf[t] * (1 - rho1) * (1 - mu), 0))
    trans.R1[t] <- as.integer(round(new.inf[t] * rho1 * (1 - mu), 0))
    recover[t] <- as.integer(round(I[t-1] * sigma1 * (1 - mu), 0))
    revert[t] <- as.integer(round(R1[t-1] * gamma1 * (1 - mu), 0))
    S[t] <- round((S[t-1] * (1 - mu)) - new.inf[t] + births[t], 0)
    I[t] <- round((I[t-1] * (1 - mu) * (1 - sigma1) + trans.I[t]), 0)
    R1[t] <- round(R1[t-1] * (1 - mu) * (1 - gamma1) + trans.R1[t] + recover[t], 0)
    R2[t] <- round(R2[t-1] * (1 - mu) + revert[t], 0)
    N[t] <- S[t] + I[t] + R1[t] + R2[t]
  }else if(mod == "m2"){    
    iota <- rbinom(n = 1, size = 1, prob = iota2)
    foi[t-1] <- (bp2 * swe[t-1]) * (I[t-1] + iota)/(N[t-1]^theta2)
    new.inf[t] <- pomp2::reulermultinom(2, S[t-1], foi[t-1], 1)[1]
    trans.I[t] <- as.integer(round(new.inf[t] * (1 - rho2) * (1 - mu), 0))
    trans.R1[t] <- as.integer(round(new.inf[t] * rho2 * (1 - mu), 0))
    recover[t] <- as.integer(round(I[t-1] * sigma2 * (1 - mu), 0))
    revert[t] <- as.integer(round(R1[t-1] * gamma2 * (1 - mu), 0))
    S[t] <- round((S[t-1] * (1 - mu)) - new.inf[t] + births[t], 0)
    I[t] <- round((I[t-1] * (1 - mu) * (1 - sigma2) + trans.I[t]), 0)
    R1[t] <- round(R1[t-1] * (1 - mu) * (1 - gamma2) + trans.R1[t] + recover[t], 0)
    R2[t] <- round(R2[t-1] * (1 - mu) + revert[t], 0)
    N[t] <- S[t] + I[t] + R1[t] + R2[t]
  }
  
}

out <- data.frame(site = site, mod = mod, time = c(1:n.years), S = S, I = I, R1 = R1, R2 = R2, N = N,
                  births = births)

# seroprevalence and population trends
ggplot(data = out, aes(x = time, y = (I + R1)/N))+
  geom_line(color = "blue", size = 1) +
  ylim(0, 1)+
  gg.gavin.theme()

ggplot(data = out, aes(x = time, y = N))+
  geom_line(color = "black", size = 1)+
  gg.gavin.theme()

# SIRR composition
ggplot()+
  geom_line(data = out, aes(x = time, y = S), color = "green", size = 1)+
  geom_line(data = out, aes(x = time, y = I), color = "red", size = 1)+
  geom_line(data = out, aes(x = time, y = R1), color = "orange", size = 1)+
  geom_line(data = out, aes(x = time, y = R2), color = "grey", size = 1)+
  ylab("")+
  gg.gavin.theme()

# # to save:
# write.csv(out, paste0("./site",site,"mod",mod,".csv"), row.names = F)