# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
## MODEL FITTING ===============================================================

# Description:
#     Fit interaction models with different combinations of covariates
#     (temperature anomaly, long lag SPI and short lag SPI) with disease cases
#     as the response variable and population as the offset.

# Paper:
#     Compound and cascading effects of climatic extremes on dengue outbreak
#     risk in the Caribbean: an impact-based modelling framework with long-lag
#     and short-lag interactions

# Script authors:
#     Chloe Fletcher        ORCID: 0000-0002-6705-7605
#     Prof. Rachel Lowe     ORCID: 0000-0003-3939-7343

# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''

## Source packages, functions and data -----------------------------------------

# load packages and functions
source("functions/00_packages_functions.R")

# read in harmonised data
data <- read.csv("data/barbados_data.csv")


## Baseline models -------------------------------------------------------------

# INTERCEPT MODEL
# - intercept

# run intercept only model
intercept <- "cases ~ 1"
int.mod <- runinlamod(formula(intercept), data)

# save intercept only model
# (un-comment for the first run, but want to retain same comparative model)
# saveRDS(int.mod, file = "outputs/int_mod.rds")

# read in intercept only model
int.mod <- readRDS("outputs/int_mod.rds")


# LAGGED DIR MODEL (BASE)
# - intercept
# - monthly cyclic random effect using random walk (RW2)
# - interannual random effects using random walk (RW2)
# - logarithm of dengue incidence rate plus 1 lagged by 4 months

re1 <- paste("f(month, model = 'rw2', cyclic = TRUE, constr = TRUE,",
             "scale.model = TRUE, hyper = precision.prior)")
re2 <- "f(year_index, model = 'rw2', hyper = precision.prior)"
fe <- "logdir.4"
baseformula <- paste(c("cases ~ 1", re1, re2, fe), collapse = " + ")

# run base model with negative binomial and poisson likelihoods
base.mod <- runinlamod(formula(baseformula), data)
base.mod.p <- runinlamod(formula(baseformula), data, family = "poisson")

# compare the waic across base models with different likelihoods
waic(base.mod) - waic(base.mod.p)

# save base 
# (un-comment for the first run, but want to retain same comparative model)
# saveRDS(base.mod, file = "outputs/base_mod.rds")

# read in base model
base.mod <- readRDS("outputs/base_mod.rds")


## Create list of variables ----------------------------------------------------

# create list of variables for analysis
tas_vars <- grep("tas.*\\.\\d+$", colnames(data), value=TRUE)
spi_vars <- grep("spi", colnames(data), value = TRUE)
spi1_vars_short <- grep("spi1\\.[1-3]$", colnames(data), value=TRUE)
spi3_vars_long <- grep("spi3.*\\.[4-6]$", colnames(data), value=TRUE)
spi3_vars_short <- grep("spi3.*\\.[1-3]$", colnames(data), value=TRUE)
spi6_vars_long <- grep("spi6.*\\.[4-6]$", colnames(data), value=TRUE)
spi6_vars_short <- grep("spi6.*\\.[1-3]$", colnames(data), value=TRUE)
spi12_vars_long <- grep("spi12.*\\.[4-6]$", colnames(data), value=TRUE)


## Univariable analysis --------------------------------------------------------

# create dataframe of univariable models
uni.mod <- data.frame(vars = c(tas_vars, spi_vars))

# create list with gof dataframe, fitted values, fixed effects, random effects
uni.mod.out <- create.mod.out(uni.mod)

# run models
uni.mod.out <- fit.models(uni.mod.out, data, base.mod, int.mod, 
                          fname="univariable_mods")


## Long SPI-3 short SPI-1 ------------------------------------------------------

# create dataframe of spi31 models
spi31.mod <- expand.grid(tas_vars, spi3_vars_long, spi1_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi31.mod.out <- create.mod.out(spi31.mod)

# run models and save outputs
spi31.mod.out <- fit.models(spi31.mod.out, data, base.mod, int.mod,
                            fname="outputs/spi31_mods")

# save spi31 models
saveRDS(spi31.mod.out, file = "outputs/spi31_mods.rds")


## Long SPI-3 short SPI-3 ------------------------------------------------------

# create dataframe of spi3 models
spi3.mod <- expand.grid(tas_vars, spi3_vars_long, spi3_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi3.mod.out <- create.mod.out(spi3.mod)

# run models and save outputs
spi3.mod.out <- fit.models(spi3.mod.out, data, base.mod, int.mod,
                           fname="outputs/spi3_mods")

# save spi3 models
saveRDS(spi3.mod.out, file = "outputs/spi3_mods.rds")


## Long SPI-6 short SPI-1 ------------------------------------------------------

# create dataframe of spi61 models
spi61.mod <- expand.grid(tas_vars, spi6_vars_long, spi1_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi61.mod.out <- create.mod.out(spi61.mod)

# run models and save outputs
spi61.mod.out <- fit.models(spi61.mod.out, data, base.mod, int.mod,
                            fname="outputs/spi61_mods")

# save spi61 models
saveRDS(spi61.mod.out, file = "outputs/spi61_mods.rds")


## Long SPI-6 short SPI-3 ------------------------------------------------------

# create dataframe of spi63 models
spi63.mod <- expand.grid(tas_vars, spi6_vars_long, spi3_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi63.mod.out <- create.mod.out(spi63.mod)

# run models and save outputs
spi63.mod.out <- fit.models(spi63.mod.out, data, base.mod, int.mod,
                            fname="outputs/spi63_mods")

# save spi63 models
saveRDS(spi63.mod.out, file = "outputs/spi63_mods.rds")


## Long SPI-6 short SPI-6 ------------------------------------------------------

# create dataframe of spi6 models
spi6.mod <- expand.grid(tas_vars, spi6_vars_long, spi6_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi6.mod.out <- create.mod.out(spi6.mod)

# run models and save outputs
spi6.mod.out <- fit.models(spi6.mod.out, data, base.mod, int.mod,
                           fname="outputs/spi6_mods")

# save spi6 models
saveRDS(spi6.mod.out, file = "outputs/spi6_mods.rds")


## Long SPI-12 short SPI-1 ------------------------------------------------------

# create dataframe of spi121 models
spi121.mod <- expand.grid(tas_vars, spi12_vars_long, spi1_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi121.mod.out <- create.mod.out(spi121.mod)

# run models and save outputs
spi121.mod.out <- fit.models(spi121.mod.out, data, base.mod, int.mod,
                             fname="outputs/spi121_mods")

# save spi121 models
saveRDS(spi121.mod.out, file = "outputs/spi121_mods.rds")


## Long SPI-12 short SPI-3 ------------------------------------------------------

# create dataframe of spi123 models
spi123.mod <- expand.grid(tas_vars, spi12_vars_long, spi3_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi123.mod.out <- create.mod.out(spi123.mod)

# run models and save outputs
spi123.mod.out <- fit.models(spi123.mod.out, data, base.mod, int.mod,
                             fname="outputs/spi123_mods")

# save spi123 models
saveRDS(spi123.mod.out, file = "outputs/spi123_mods.rds")


## Long SPI-12 short SPI-6 ------------------------------------------------------

# create dataframe of spi126 models
spi126.mod <- expand.grid(tas_vars, spi12_vars_long, spi6_vars_short) %>% 
  mutate(vars=paste0(Var1, " * ", Var2, " * ", Var3),
         vars=str_replace(vars, "\\+$|^\\+", ""))

# create list with gof dataframe, fitted values, fixed effects, random effects
spi126.mod.out <- create.mod.out(spi126.mod)

# run models and save outputs
spi126.mod.out <- fit.models(spi126.mod.out, data, base.mod, int.mod,
                             fname="outputs/spi126_mods")

# save spi126 models
saveRDS(spi126.mod.out, file = "outputs/spi126_mods.rds")


## END