# 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

##' Analysis using Muddy Creek m0 model
##' Alternate TAS begins line 445

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

setwd(dirname(rstudioapi::getActiveDocumentContext()$path))

library(pomp2)
library(tidyverse)
library(gridExtra)

theme_set(theme_bw())
set.seed(594709947L)

load("./out/muddy_m0_swe_04Aug.RData")

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_blank(),
        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"))
}

m0 <- pomp(dat, 
           time = "dyear",
           t0 = 2004,
           rprocess = discrete_time(rproc, delta.t = 1), 
           rinit = initlz,
           dmeasure = dmeas,
           rmeasure = rmeas,
           covar = covariate_table(covar, times = "cyear"),
           statenames = c("S", "I", "R1", "R2", "truesp", "N"),
           paramnames = c("beta","mu", "sigma", "rho","gamma", "theta","iota",
                          "S_0","I_0","R1_0","R2_0"),
           partrans = parameter_trans(log = log_params, logit = logit_params))

#=========================
# look at global results
#=========================

summary(results_global$loglik) 

library(plyr)
all <- ldply(list(guess = guesses, result = subset(results_global, loglik > max(loglik) - 50)))

pairs(~loglik + beta + mu + sigma + rho + gamma + theta + iota,
      data = all, col = ifelse(all$.id == "guess", grey(0.5), "red"), pch = 16)

### to write:
png("./Figs/muddy_m0_pairs_guess_params.png", width = 2000, height = 2000, units = "px", res = 275)
pairs(~loglik + beta + mu + + sigma + rho + gamma + theta + iota,
      data = all, col = ifelse(all$.id == "guess", grey(0.5), "red"), pch = 16)
dev.off()

pairs(~loglik + S_0 + I_0 + R1_0 + R2_0,
      data = all, col = ifelse(all$.id == "guess", grey(0.5), "red"), pch = 16)

### to write:
png("./Figs/muddy_m0_pairs_guess_SIRR.png", width = 2000, height = 2000, units = "px", res = 275)
pairs(~loglik + S_0 + I_0 + R1_0 + R2_0, data = all,
      col = ifelse(all$.id == "guess", grey(0.5), "red"), pch = 16)
dev.off()

#=================================================================
# simulate again with the 'mle' from the global search 
#=================================================================

paramnames <- c("beta", "mu",  "sigma", "rho", "gamma", "theta", "iota",
                "S_0", "I_0", "R1_0", "R2_0")

results <- rbind(results_local, results_global)

(m0_mle <- results[which.max(results[,"loglik"]),][paramnames])

m0_mle_vals <- c(beta = m0_mle[,"beta"], 
                 mu = m0_mle[,"mu"],
                 sigma = m0_mle[,"sigma"], 
                 rho = m0_mle[,"rho"],
                 gamma = m0_mle[,"gamma"], 
                 theta = m0_mle[,"theta"],
                 iota = m0_mle[,"iota"],
                 S_0 = m0_mle[,"S_0"], 
                 I_0 = m0_mle[,"I_0"], 
                 R1_0 = m0_mle[,"R1_0"], 
                 R2_0 = m0_mle[,"R2_0"])

wils <- cbind(dat, tot = covar$tot[-1])
wils <- wils%>%
  select(-dpop)%>%
  mutate(lo = NA, hi = NA)

library(Hmisc)

for(i in 1:nrow(wils)){
  if(wils$pos[i] %in% NA) next
  tmp95 <- binconf(wils$pos[i], wils$tot[i], alpha = 0.1, method = "wilson")
  wils$lo[i] <- tmp95[2]
  wils$hi[i] <- tmp95[3]
}

#####################
### 90% quantiles ###
#####################

# save the serop quantiles
serop_1 <- m0 %>%
  simulate(params=m0_mle_vals,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,apparent))%>%
  mutate(data=.id=="data") %>%
  ddply(~dyear+data,plyr::summarize,
        p=c(0.05, 0.5, 0.95),q=quantile(apparent,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi")),
         data=mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')

disease_3 <- serop_1%>%
  ggplot()+
  geom_ribbon(aes(x=dyear,y=med,color=data,fill=data,ymin=lo,ymax=hi), alpha=0.2, show.legend = FALSE)+
  gg.gavin.theme()+
  labs(x = NULL, y = "Seroprevalence")+
  ylim(c(0,1))+
  ggtitle("Muddy Creek")+
  geom_errorbar(wils, mapping = aes(x = dyear, ymin = lo, ymax = hi, width = 0.2))+
  scale_x_discrete(limits = c(2008,2012, 2016))

# save the pop trend
popul_1 <- m0 %>%
  simulate(params=m0_mle_vals,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,dpop))%>%
  mutate(data=.id=="data") %>%
  ddply(~dyear+data,plyr::summarize,
        p=c(0.05, 0.5, 0.95),q=quantile(dpop,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi")),
         data=mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')%>%
  rbind(data.frame(dyear=2004,data="data",hi=325,lo=325,med=325))

pop_3 <- popul_1%>%
  ggplot()+
  geom_ribbon(aes(x=dyear,y=med,color=data,fill=data,ymin=lo,ymax=hi), alpha=0.2, show.legend = FALSE)+
  gg.gavin.theme()+
  labs(x = NULL, y = "Female count")+
  ggtitle("")+
  scale_x_discrete(limits = c(2008,2012,2016))

grid.arrange(disease_3, pop_3, ncol = 2)

#####################
### FOI/SIRR plot ###
#####################

nsim = 2e3

sims3 <- simulate(m0,
                  params = m0_mle_vals,
                  nsim = nsim,
                  format = "data.frame", include.data = TRUE)

# double check the ordering here
sims3 <- sims3[order(sims3$.id, sims3$dyear),]

foi_df <- sims3%>%
  filter(.id != "data")%>%
  mutate(beta_foi = m0_mle[,"beta"] * (I + m0_mle[,"iota"]) / N^m0_mle[,"theta"])%>%
  ddply(~dyear,plyr::summarize,
        p=c(0.05, 0.5, 0.95),qBeta=quantile(beta_foi,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qBeta')

width=1.5

muddy_foi <- ggplot(foi_df, aes(x = dyear))+
  geom_ribbon(aes(y = med, ymin = lo, ymax = hi), color = 'red', fill = "red", alpha = 0.2, show.legend = FALSE)+
  geom_line(aes(y=med),linetype="dashed",size=width,color='red')+
  gg.gavin.theme()+
  labs(x = NULL, y = expression(hat(lambda)))+
  ggtitle("Muddy Creek")+
  scale_y_continuous(breaks = c(0,0.1,0.2))
muddy_foi



testS <- m0 %>%
  simulate(params=m0_mle_vals,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, S))%>%
  subset(.id!="data")%>%
  ddply(~dyear,plyr::summarize,
        p=c(0.05, 0.5, 0.95),qS=quantile(S,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qS')

testI <- m0 %>%
  simulate(params=m0_mle_vals,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, I))%>%
  subset(.id!="data")%>%
  ddply(~dyear,plyr::summarize,
        p=c(0.05, 0.5, 0.95),qI=quantile(I,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qI')

testR1 <- m0 %>%
  simulate(params=m0_mle_vals,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,R1))%>%
  subset(.id!="data")%>%
  ddply(~dyear,plyr::summarize,
        p=c(0.05, 0.5, 0.95),qR1=quantile(R1,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR1')

testR2 <- m0 %>%
  simulate(params=m0_mle_vals,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, R2))%>%
  subset(.id!="data")%>%
  ddply(~dyear,plyr::summarize,
        p=c(0.05, 0.5, 0.95),qR2=quantile(R2,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=mapvalues(p,from=c(0.05, 0.5, 0.95),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR2')

# 16%?
sum(testI$med)/ (sum(testI$med) + sum(testR1$med))
sum(testI$lo)/ (sum(testI$lo) + sum(testR1$lo))
sum(testI$hi)/ (sum(testI$hi) + sum(testR1$hi))

muddy_sirr <- ggplot()+
  geom_ribbon(data = testS, aes(x=dyear,y=med,ymin=lo,ymax=hi), 
              color = "green", fill = "green", alpha=0.2, show.legend = FALSE)+
  geom_line(data = testS, aes(x=dyear,y=med), size=width, linetype="dashed", color = "green")+
  geom_ribbon(data = testI, aes(x=dyear,y=med,ymin=lo,ymax=hi), 
              color = "red", fill = "red", alpha=0.2, show.legend = FALSE)+
  geom_line(data = testI, aes(x=dyear,y=med), size=width, linetype="dashed", color = "red")+
  
  geom_ribbon(data = testR1, aes(x=dyear,y=med,ymin=lo,ymax=hi), 
              color = "orange", fill = "orange", alpha=0.2, show.legend = FALSE)+
  geom_line(data = testR1, aes(x=dyear,y=med), size=width, linetype="dashed", color = "orange")+
  
  geom_ribbon(data = testR2, aes(x=dyear,y=med,ymin=lo,ymax=hi), 
              color = "grey", fill = "grey", alpha=0.2, show.legend = FALSE)+
  geom_line(data = testR2, aes(x=dyear,y=med), size=width, linetype="dashed", color = "grey")+
  gg.gavin.theme()+
  theme(plot.margin=unit(c(1,1,0,1.2), unit="cm"))+
  labs(x=NULL, y="Individuals")+
  ggtitle("SIRR composition")
muddy_sirr

grid.arrange(muddy_foi, muddy_sirr, ncol = 2)

# AIC calc
results[which(results$loglik == max(results$loglik)),]
(ll <- max(results$loglik))
# loglik and se columns aren't params, hence: -2
k <- length(results_global) - 2

(AIC_muddy_m0 <- -2 * ll + 2*k) 



#############################################
### Test and slaughter regime comparisons ###
#############################################


# Estimate the number of infectious elk removed ####
rproc1 <- Csnippet("
                   double foi, muI, muR1;
                   double rate[9], trans[9];
                   
                   // rem I and R1
                   if(t == 2006) {
                   muI = 0;
                   muR1 = 0;
                   } else if (t == 2007) {
                   muI = 0;
                   muR1 = 0;
                   } else if (t == 2008) {
                   muI = 0;
                   muR1 = 0;
                   } else if (t == 2009) {
                   muI = 0;
                   muR1 = 0;
                   } else if (t == 2010) {
                   muI = 0;
                   muR1 = 0;
                   } else {
                   muI = 0;
                   muR1 = 0;
                   }; 
                   // expected force of infection
                   foi = beta * (I + iota) / pow(pop, theta);
                   
                   rate[0] = foi * dt;        // S to I/R1, force of infection
                   rate[1] = mu;              // S natural death
                   rate[2] = sigma;           // I to R1, rate of recovery
                   rate[3] = mu;              // I natural death
                   rate[4] = muI;             // I management removals      
                   rate[5] = gamma;           // R1 to R2, rate of detectable antibody loss
                   rate[6] = mu;              // R1 natural death
                   rate[7] = muR1;            // R1 management removals
                   rate[8] = mu;              // R2 natural death
                   
                   // transitions between classes
                   reulermultinom(2, S, &rate[0], dt, &trans[0]);
                   reulermultinom(3, I, &rate[2], dt, &trans[2]);
                   reulermultinom(3, R1, &rate[5], dt, &trans[5]);
                   reulermultinom(1, R2, &rate[8], dt, &trans[8]);
                   
                   S += births   - trans[0] - trans[1];
                   I += nearbyint(trans[0] * (1 - rho)) - trans[2] - trans[3] - trans[4];
                   R1 += nearbyint(trans[0] * rho) + trans[2] - trans[5] - trans[6] - trans[7];
                   R2 += trans[5] - trans[8]; 
                   N = S + I + R1 + R2;
                   truesp = (I + R1) / N;               // true seroprevalence
                   ")




m01 <- pomp(dat, 
            time = "dyear",
            t0 = 2004,
            rprocess = discrete_time(rproc1, delta.t = 1), 
            rinit = initlz,
            dmeasure = dmeas,
            rmeasure = rmeas,
            covar = covariate_table(covar, times = "cyear"),
            statenames = c("S", "I", "R1", "R2", "truesp", "N"),
            paramnames = c("beta","mu", "sigma", "rho","gamma","theta","iota", 
                           "S_0","I_0","R1_0","R2_0"),
            partrans = parameter_trans(log = log_params, logit = logit_params))

m0_mle_vals2 <- c(beta = m0_mle[,"beta"], 
                  mu = m0_mle[,"mu"],
                  sigma = m0_mle[,"sigma"], 
                  rho = m0_mle[,"rho"],
                  gamma = m0_mle[,"gamma"], 
                  theta = m0_mle[,"theta"],
                  iota = m0_mle[,"iota"],
                  S_0 = m0_mle[,"S_0"], 
                  I_0 = m0_mle[,"I_0"], 
                  R1_0 = m0_mle[,"R1_0"], 
                  R2_0 = m0_mle[,"R2_0"])

# save the disease trend
serop_2 <- m01 %>%
  simulate(params=m0_mle_vals2,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,apparent))%>%
  mutate(data=.id=="data") %>%
  plyr::ddply(~dyear+data,plyr::summarize,
              p=c(0.25, 0.5, 0.75),q=quantile(apparent,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi")),
         data=plyr::mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')

# save the pop trend
popul_2 <- m01 %>%
  simulate(params=m0_mle_vals2,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,dpop))%>%
  mutate(data=.id=="data") %>%
  plyr::ddply(~dyear+data,plyr::summarize,
              p=c(0.25, 0.5, 0.75),q=quantile(dpop,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi")),
         data=plyr::mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')

sims4 <- simulate(m01,
                  params = m0_mle_vals2,
                  nsim = nsim,
                  format = "data.frame", include.data = TRUE)

# check the ordering here
sims4 <- sims4[order(sims4$.id, sims4$dyear),]

foi_df2 <- sims4%>%
  filter(.id != "data")%>%
  mutate(beta_foi = m0_mle[,"beta"] * (I + m0_mle[,"iota"]) / N^m0_mle[,"theta"])%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qBeta=quantile(beta_foi,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qBeta')



# SIRR

testS2 <- m01 %>% 
  simulate(params=m0_mle_vals2,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, S))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qS=quantile(S,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qS')

testI2 <- m01 %>% 
  simulate(params=m0_mle_vals2,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, I))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qI=quantile(I,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qI')

testR12 <- m01 %>% 
  simulate(params=m0_mle_vals2,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,R1))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qR1=quantile(R1,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR1')

testR22 <- m01 %>% 
  simulate(params=m0_mle_vals2,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, R2))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qR2=quantile(R2,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR2')

# total:
sum(testI2$med - testI$med) 

# 2006 - 2010:
sum(testI2$med[2:6] - testI$med[2:6]) 

# 2011 - 2018:
sum(testI2$med[7:14] - testI$med[7:14]) 







##########################################
### 10% seronegative removal, S and R2 ###
##########################################

rproc2 <- Csnippet("
                   double foi, muNeg;
                   double rate[9], trans[9];
                   
                   // rem S and R2
                   if(t >= 2006 && t <= 2010) {
                   muNeg = rem;
                   } else {
                   muNeg = 0;
                   }; 
                   // expected force of infection
                   foi = beta * (I + iota) / pow(pop, theta);
                   
                   rate[0] = foi * dt;        // S to I/R1, force of infection
                   rate[1] = mu;              // S natural death
                   rate[2] = muNeg;             // S and R2 management removals      
                   rate[3] = sigma;           // I to R1, rate of recovery
                   rate[4] = mu;              // I natural death
                   rate[5] = gamma;           // R1 to R2, rate of detectable antibody loss
                   rate[6] = mu;              // R1 natural death
                   rate[7] = mu;              // R2 natural death
                   rate[8] = muNeg;
                   
                   
                   // transitions between classes
                   reulermultinom(3, S, &rate[0], dt, &trans[0]);
                   reulermultinom(2, I, &rate[3], dt, &trans[3]);
                   reulermultinom(2, R1, &rate[5], dt, &trans[5]);
                   reulermultinom(2, R2, &rate[7], dt, &trans[7]);
                   
                   S += births   - trans[0] - trans[1]  - trans[2];
                   I += nearbyint(trans[0] * (1 - rho)) - trans[3] - trans[4];
                   R1 += nearbyint(trans[0] * rho) + trans[3] - trans[5] - trans[6];
                   R2 += trans[5] - trans[7]- trans[8]; 
                   N = S + I + R1 + R2;
                   truesp = (I + R1) / N;               // true seroprevalence
                   ")

initlz2 <- Csnippet("
                    double m = pop / (S_0 + I_0 + R1_0 + R2_0);
                    S = nearbyint(m * S_0);
                    I = nearbyint(m * I_0);
                    R1 = nearbyint(m * R1_0);
                    R2 = nearbyint(m * R2_0);
                    truesp = (I + R1) / pop;
                    ")

dmeas2 <- Csnippet("if (ISNA(pos)) {
                   lik = (give_log) ? 0 : 1;
                   } else {
                   lik = dbinom(pos, tot, truesp, give_log) +
                   dnorm(dpop, N, 20, give_log); 
                   }")

rmeas2 <- Csnippet("dpop = rnorm(N, 20);
                   pos = rbinom(tot, truesp);
                   apparent = pos/tot;")

# transform parameters for use with unconstrained optimizer
log_params2 <- c("beta", "iota")                                              # params > 0
logit_params2 <- c("mu", "rem", "sigma", "rho", "gamma", "theta", 
                   "S_0", "I_0", "R1_0", "R2_0")                              # 1 > params > 0                                  

m02 <- pomp(dat, 
            time = "dyear",
            t0 = 2004,
            rprocess = discrete_time(rproc2, delta.t = 1), 
            rinit = initlz2,
            dmeasure = dmeas2,
            rmeasure = rmeas2,
            covar = covariate_table(covar, times = "cyear"),
            statenames = c("S", "I", "R1", "R2", "truesp", "N"),
            paramnames = c("beta","mu", "rem", "sigma", "rho","gamma","theta","iota", 
                           "S_0","I_0","R1_0","R2_0"),
            partrans = parameter_trans(log = log_params2, logit = logit_params2))


m0_mle_vals3 <- c(beta = m0_mle[,"beta"], 
                  mu = m0_mle[,"mu"],
                  rem = 0.10,
                  sigma = m0_mle[,"sigma"], 
                  rho = m0_mle[,"rho"],
                  gamma = m0_mle[,"gamma"], 
                  theta = m0_mle[,"theta"],
                  iota = m0_mle[,"iota"],
                  S_0 = m0_mle[,"S_0"], 
                  I_0 = m0_mle[,"I_0"], 
                  R1_0 = m0_mle[,"R1_0"], 
                  R2_0 = m0_mle[,"R2_0"])

# save disease trend:
serop_3 <- m02 %>%
  simulate(params=m0_mle_vals3,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,apparent))%>%
  mutate(data=.id=="data") %>%
  plyr::ddply(~dyear+data,plyr::summarize,
              p=c(0.25, 0.5, 0.75),q=quantile(apparent,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi")),
         data=plyr::mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')

# save the pop trend:
popul_3 <- m02 %>%
  simulate(params=m0_mle_vals3,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,dpop))%>%
  mutate(data=.id=="data") %>%
  plyr::ddply(~dyear+data,plyr::summarize,
              p=c(0.25, 0.5, 0.75),q=quantile(dpop,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi")),
         data=plyr::mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')


sims5 <- simulate(m02,
                  params = m0_mle_vals3,
                  nsim = nsim,
                  format = "data.frame", include.data = TRUE)

sims5 <- sims5[order(sims5$.id, sims5$dyear),]

foi_df3 <- sims5%>%
  filter(.id != "data")%>%
  mutate(beta_foi = m0_mle[,"beta"] * (I + m0_mle[,"iota"]) / N^m0_mle[,"theta"])%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qBeta=quantile(beta_foi,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qBeta')



# SIRR

testS3 <- m02 %>% 
  simulate(params=m0_mle_vals3,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, S))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qS=quantile(S,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qS')

testI3 <- m02 %>% 
  simulate(params=m0_mle_vals3,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, I))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qI=quantile(I,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qI')

testR13 <- m02 %>% 
  simulate(params=m0_mle_vals3,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,R1))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qR1=quantile(R1,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR1')

testR23 <- m02 %>% 
  simulate(params=m0_mle_vals3,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, R2))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qR2=quantile(R2,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR2')


# total:
sum(testI3$med - testI$med) 

# 2006 - 2010:
sum(testI3$med[2:6]) - sum(testI$med[2:6]) 

# next 8 years:
sum(testI3$med[7:14] - testI$med[7:14]) 


#################
### 7.5% cull ###
#################

rproc3 <- Csnippet("
                   double foi, muRem;
                   double rate[11], trans[11];
                   
                   // rem S and R2
                   if(t >= 2006 && t <= 2010) {
                   muRem = rem;
                   } else {
                   muRem = 0;
                   }; 
                   // expected force of infection
                   foi = beta * (I + iota) / pow(pop, theta);
                   
                   rate[0] = foi * dt;        // S to I/R1, force of infection
                   rate[1] = mu;              // S natural death
                   rate[2] = muRem;           // S management removals  
                   
                   rate[3] = sigma;           // I to R1, rate of recovery
                   rate[4] = mu;              // I natural death
                   rate[5] = muRem;           // I management removals   
                   
                   rate[6] = gamma;           // R1 to R2, rate of detectable antibody loss
                   rate[7] = mu;              // R1 natural death
                   rate[8] = muRem;           // R1 management removals   
                   
                   rate[9] = mu;              // R2 natural death
                   rate[10] = muRem;           // R2 management removals
                   
                   
                   // transitions between classes
                   reulermultinom(3, S, &rate[0], dt, &trans[0]);
                   reulermultinom(3, I, &rate[3], dt, &trans[3]);
                   reulermultinom(3, R1, &rate[6], dt, &trans[6]);
                   reulermultinom(2, R2, &rate[9], dt, &trans[9]);
                   
                   S += births - trans[0] - trans[1]  - trans[2];
                   I += nearbyint(trans[0] * (1 - rho)) - trans[3] - trans[4] - trans[5];
                   R1 += nearbyint(trans[0] * rho) + trans[3] - trans[6] - trans[7] - trans[8];
                   R2 += trans[6] - trans[9]- trans[10]; 
                   N = S + I + R1 + R2;
                   truesp = (I + R1) / N;               // true seroprevalence
                   ")

initlz3 <- Csnippet("
                    double m = pop / (S_0 + I_0 + R1_0 + R2_0);
                    S = nearbyint(m * S_0);
                    I = nearbyint(m * I_0);
                    R1 = nearbyint(m * R1_0);
                    R2 = nearbyint(m * R2_0);
                    truesp = (I + R1) / pop;
                    ")

dmeas3 <- Csnippet("if (ISNA(pos)) {
                   lik = (give_log) ? 0 : 1;
                   } else {
                   lik = dbinom(pos, tot, truesp, give_log) +
                   dnorm(dpop, N, 20, give_log); 
                   }")

rmeas3 <- Csnippet("dpop = rnorm(N, 20);
                   pos = rbinom(tot, truesp);
                   apparent = pos/tot;")

# transform parameters for use with unconstrained optimizer
log_params3 <- c("beta", "iota")                                              # params > 0
logit_params3 <- c("mu", "rem", "sigma", "rho", "gamma", "theta", 
                   "S_0", "I_0", "R1_0", "R2_0")                              # 1 > params > 0                                  

m03 <- pomp(dat, 
            time = "dyear",
            t0 = 2004,
            rprocess = discrete_time(rproc3, delta.t = 1), 
            rinit = initlz3,
            dmeasure = dmeas3,
            rmeasure = rmeas3,
            covar = covariate_table(covar, times = "cyear"),
            statenames = c("S", "I", "R1", "R2", "truesp", "N"),
            paramnames = c("beta","mu", "rem", "sigma", "rho","gamma","theta","iota", 
                           "S_0","I_0","R1_0","R2_0"),
            partrans = parameter_trans(log = log_params3, logit = logit_params3))


m0_mle_vals4 <- c(beta = m0_mle[,"beta"], 
                  mu = m0_mle[,"mu"],
                  rem = 0.075,
                  sigma = m0_mle[,"sigma"], 
                  rho = m0_mle[,"rho"],
                  gamma = m0_mle[,"gamma"], 
                  theta = m0_mle[,"theta"],
                  iota = m0_mle[,"iota"],
                  S_0 = m0_mle[,"S_0"], 
                  I_0 = m0_mle[,"I_0"], 
                  R1_0 = m0_mle[,"R1_0"], 
                  R2_0 = m0_mle[,"R2_0"])

# save disease trend:
serop_4 <- m03 %>%
  simulate(params=m0_mle_vals4,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,apparent))%>%
  mutate(data=.id=="data") %>%
  plyr::ddply(~dyear+data,plyr::summarize,
              p=c(0.25, 0.5, 0.75),q=quantile(apparent,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi")),
         data=plyr::mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')

# save pop trend:
popul_4 <- m03 %>%
  simulate(params=m0_mle_vals4,nsim=2e4,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,dpop))%>%
  mutate(data=.id=="data") %>%
  plyr::ddply(~dyear+data,plyr::summarize,
              p=c(0.25, 0.5, 0.75),q=quantile(dpop,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi")),
         data=plyr::mapvalues(data,from=c(TRUE,FALSE),to=c("data","simulation"))) %>%
  reshape2::dcast(dyear+data~p,value.var='q')


sims6 <- simulate(m03,
                  params = m0_mle_vals4,
                  nsim = nsim,
                  format = "data.frame", include.data = TRUE)

# check the ordering here
sims6 <- sims6[order(sims6$.id, sims6$dyear),]

foi_df4 <- sims6%>%
  filter(.id != "data")%>%
  mutate(beta_foi = m0_mle[,"beta"] * (I + m0_mle[,"iota"]) / N^m0_mle[,"theta"])%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qBeta=quantile(beta_foi,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qBeta')




# SIRR

testS4 <- m03 %>% 
  simulate(params=m0_mle_vals4,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, S))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qS=quantile(S,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qS')

testI4 <- m03 %>% 
  simulate(params=m0_mle_vals4,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, I))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qI=quantile(I,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qI')

testR14 <- m03 %>% 
  simulate(params=m0_mle_vals4,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id,R1))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qR1=quantile(R1,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR1')

testR24 <- m03 %>% 
  simulate(params=m0_mle_vals4,nsim=nsim,format="data.frame",include.data=TRUE)%>%
  subset(select=c(dyear,.id, R2))%>%
  subset(.id!="data")%>%
  plyr::ddply(~dyear,plyr::summarize,
              p=c(0.25, 0.5, 0.75),qR2=quantile(R2,prob=p,names=FALSE, na.rm=T)) %>%
  mutate(p=plyr::mapvalues(p,from=c(0.25, 0.5, 0.75),to=c("lo","med","hi"))) %>%
  reshape2::dcast(dyear~p,value.var='qR2')



# total:
sum(testI4$med - testI$med) 

# 2006 - 2010:
sum(testI4$med[2:6] - testI$med[2:6]) 

# 2006 - 2010:
sum(testI4$med[7:14] - testI$med[7:14]) 



tas_sims <- data.frame(year = testI$dyear,
                       n.inf.tas.m = testI$med,
                       n.inf.tas.l = testI$lo,
                       n.inf.tas.h = testI$hi,
                       n.inf.nada.m = testI2$med,
                       n.inf.nada.l = testI2$lo,
                       n.inf.nada.h = testI2$hi,
                       n.inf.neg.m = testI3$med,
                       n.inf.neg.l = testI3$lo,
                       n.inf.neg.h = testI3$hi,
                       n.inf.cul.m = testI4$med,
                       n.inf.cul.l = testI4$lo,
                       n.inf.cul.h = testI4$hi
)


width=1.5

p1 <- ggplot(data = tas_sims, aes(x = year))+
  geom_ribbon(aes(y=n.inf.tas.m, ymin=n.inf.tas.l,ymax=n.inf.tas.h), 
              color = "red", fill = "red", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=n.inf.tas.m), color = "red", size=width, linetype="dashed")+
  
  geom_ribbon(aes(y=n.inf.nada.m, ymin=n.inf.nada.l,ymax=n.inf.nada.h), 
              color = "grey", fill = "grey", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=n.inf.nada.m), color = "grey", size=width, linetype="dashed")+
  
  geom_ribbon(aes(y=n.inf.neg.m, ymin=n.inf.neg.l,ymax=n.inf.neg.h), 
              color = "blue", fill = "blue", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=n.inf.neg.m), color = "blue", size=width, linetype="dashed")+
  
  geom_ribbon(aes(y=n.inf.cul.m, ymin=n.inf.cul.l,ymax=n.inf.cul.h), 
              color = "orange", fill = "orange", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=n.inf.cul.m), color = "orange", size=width, linetype="dashed")+
  
  gg.gavin.theme()+
  labs(x = NULL, y = "Number of infectives")+
  ggtitle("")+
  geom_vline(xintercept= c(2006, 2010), linetype = "dashed")+
  theme(plot.margin = unit(c(0,0.5,0,0.5), "cm"))
p1


serop_1 <- serop_1%>%
  filter(data == "simulation")
serop_2 <- serop_2%>%
  filter(data == "simulation")
serop_3 <- serop_3%>%
  filter(data == "simulation")
serop_4 <- serop_4%>%
  filter(data == "simulation")

tas_serop <- data.frame(year = serop_1$dyear,
                        tas.m = serop_1$med,
                        tas.l = serop_1$lo,
                        tas.h = serop_1$hi,
                        nada.m = serop_2$med,
                        nada.l = serop_2$lo,
                        nada.h = serop_2$hi,
                        neg.m = serop_3$med,
                        neg.l = serop_3$lo,
                        neg.h = serop_3$hi,
                        cul.m = serop_4$med,
                        cul.l = serop_4$lo,
                        cul.h = serop_4$hi)

p2 <- ggplot(data = tas_serop, aes(x = year))+
  geom_ribbon(aes(y=tas.m, ymin=tas.l,ymax=tas.h), 
              color = "red", fill = "red", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=tas.m), color = "red", size = width, linetype="dashed")+
  
  geom_ribbon(aes(y=nada.m, ymin=nada.l,ymax=nada.h), 
              color = "grey", fill = "grey", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=nada.m), color = "grey", size = width, linetype="dashed")+
  
  geom_ribbon(aes(y=neg.m, ymin=neg.l,ymax=neg.h), 
              color = "blue", fill = "blue", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=neg.m), color = "blue", size = width, linetype="dashed")+
  
  geom_ribbon(aes(y=cul.m, ymin=cul.l,ymax=cul.h), 
              color = "orange", fill = "orange", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=cul.m), color = "orange", size = width, linetype="dashed")+
  
  gg.gavin.theme()+
  labs(x = NULL, y = "Seroprevalence")+
  ggtitle("")+
  geom_vline(xintercept= c(2006, 2010), linetype = "dashed")+
  theme(plot.margin = unit(c(0,0.5,0,0.5), "cm"))


###########
### FOI ###
###########

tas_foi <- data.frame(year = foi_df$dyear,
                      tas.m = foi_df$med,
                      tas.l = foi_df$lo,
                      tas.h = foi_df$hi,
                      nada.m = foi_df2$med,
                      nada.l = foi_df2$lo,
                      nada.h = foi_df2$hi,
                      neg.m = foi_df3$med,
                      neg.l = foi_df3$lo,
                      neg.h = foi_df3$hi,
                      cul.m = foi_df4$med,
                      cul.l = foi_df4$lo,
                      cul.h = foi_df4$hi)

p3 <- ggplot(data = tas_foi, aes(x = year))+
  geom_ribbon(aes(y=tas.m, ymin=tas.l,ymax=tas.h),
              color = "red", fill = "red", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=tas.m), size = width, color = "red", linetype = "dashed")+
  
  geom_ribbon(aes(y=nada.m, ymin=nada.l,ymax=nada.h),
              color = "grey", fill = "grey", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=nada.m), size = width, color = "grey", linetype = "dashed")+
  
  geom_ribbon(aes(y=neg.m, ymin=neg.l,ymax=neg.h),
              color = "blue", fill = "blue", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=neg.m), size = width, color = "blue", linetype = "dashed")+
  
  geom_ribbon(aes(y=cul.m, ymin=cul.l,ymax=cul.h),
              color = "orange", fill = "orange", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=cul.m), size = width, color = "orange", linetype = "dashed")+
  
  gg.gavin.theme()+
  labs(x = NULL, y = "FOI")+
  ggtitle("")+
  geom_vline(xintercept= c(2006, 2010), linetype = "dashed")+
  theme(plot.margin = unit(c(0,0.5,0,0.5), "cm"))
p3

popul_1 <- popul_1%>%
  filter(data == "simulation")
popul_2 <- popul_2%>%
  filter(data == "simulation")
popul_3 <- popul_3%>%
  filter(data == "simulation")
popul_4 <- popul_4%>%
  filter(data == "simulation")

tas_popul <- data.frame(year = popul_1$dyear,
                        tas.m = popul_1$med,
                        tas.l = popul_1$lo,
                        tas.h = popul_1$hi,
                        nada.m = popul_2$med,
                        nada.l = popul_2$lo,
                        nada.h = popul_2$hi,
                        neg.m = popul_3$med,
                        neg.l = popul_3$lo,
                        neg.h = popul_3$hi,
                        cul.m = popul_4$med,
                        cul.l = popul_4$lo,
                        cul.h = popul_4$hi)

p4 <- ggplot(data = tas_popul, aes(x = year))+
  geom_ribbon(aes(y=tas.m, ymin=tas.l,ymax=tas.h), 
              color = "red", fill = "red", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=tas.m), size = width, color = "red", linetype = "dashed")+
  
  geom_ribbon(aes(y=nada.m, ymin=nada.l,ymax=nada.h), 
              color = "grey", fill = "grey", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=nada.m), size = width, color = "grey", linetype = "dashed")+
  
  geom_ribbon(aes(y=neg.m, ymin=neg.l,ymax=neg.h), 
              color = "blue", fill = "blue", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=neg.m), size = width, color = "blue", linetype = "dashed")+
  
  geom_ribbon(aes(y=cul.m, ymin=cul.l,ymax=cul.h), 
              color = "orange", fill = "orange", alpha=0.2, show.legend = FALSE)+
  geom_line(aes(y=cul.m), size = width, color = "orange", linetype = "dashed")+
  
  gg.gavin.theme()+
  labs(x = NULL, y = "Female count")+
  ggtitle("")+
  geom_vline(xintercept= c(2006, 2010), linetype = "dashed")+
  theme(plot.margin = unit(c(0,0.5,0,0.5), "cm"))
p4



grid.arrange(p1, p2, p4, ncol = 3, widths=c(0.99,1.015,1.025))