# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
## COMBINED CLIMATE CONTRIBUTION ===============================================

# Description:
#     Extract parameters from the final model posterior distribution, and then
#     calculate the combined contribution of the climatic variables on the 
#     dengue incidence rate (mean and 95% confidence interval) under two
#     temperature scenarios (cool and warm) and all SPI-6 combinations.

# 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
#     Dr Giovenale Moirano  ORCID: 0000-0001-8748-3321
#     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")

# read in mean variables for interpretation
mean_tmp <- readRDS("data/temperature_means.rds")


## Script inputs ---------------------------------------------------------------

# set up variables
vars <- c("tas.3mon.4", "spi6.5", "spi6.1")

# model formulation
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"
form <- paste(c("cases ~ 1", re1, re2, paste(vars, collapse=' * '), fe),
              collapse=' + ')

model <- runinlamod(formula(form), data, config=TRUE)


## Extract model parameters ----------------------------------------------------

# define number of samples to generate
s <- 1000

# extract posterior samples
sam <- inla.posterior.sample(s, model)
sam.extract <- inla.posterior.sample.eval(
  (function(...) {
    beta.1 <- get(vars[1])
    beta.2 <- get(vars[2])
    beta.3 <- get(vars[3])
    beta.4 <- get(fe)
    beta.12 <- get(paste(vars[1:2], collapse=':'))
    beta.13 <- get(paste(vars[c(1,3)], collapse=':'))
    beta.23 <- get(paste(vars[2:3], collapse=':'))
    beta.123 <- get(paste(vars, collapse=':'))
    delta <- get("month")
    gamma <- get("year_index")
    return(c(Intercept, beta.1, beta.2, beta.3, beta.4, beta.12, beta.13,
             beta.23, beta.123, delta, gamma, theta[1]))
  }), sam)

# assign to individual parameters
alpha <- sam.extract[1,]
beta.1 <- sam.extract[2,]
beta.2 <- sam.extract[3,]
beta.3 <- sam.extract[4,]
beta.4 <- sam.extract[5,]
beta.12 <- sam.extract[6,]
beta.13 <- sam.extract[7,]
beta.23 <- sam.extract[8,]
beta.123 <- sam.extract[9,]
delta <- sam.extract[10:21,]
gamma <- sam.extract[22:(21 + max(data$year_index)), ]
theta <- sam.extract[(22 + max(data$year_index)), ]


## Calculate forecast scenarios ------------------------------------------------

# set up climate variables
tmp <- c(quantile(data[, vars[1]], 0.1),
         median(data[, vars[1]]),  # low, med and high temp
         quantile(data[, vars[1]], 0.9))
spi <- seq(-2.5, 2.5, 0.25)  # spi values from -2.5 to +2.5

# create all long-short lag spi combinations
lsl_comb <- expand.grid(spi, spi)
names(lsl_comb) <- c("spi6.long", "spi6.short")

# calculate combined covariate effect on dengue risk in barbados
lsl_comb <- model.cce(tmp, 1, lsl_comb, new_col = "tmp_low")  # low temperature
lsl_comb <- model.cce(tmp, 2, lsl_comb, new_col = "tmp_med")  # med temperature
lsl_comb <- model.cce(tmp, 3, lsl_comb, new_col = "tmp_hgh")  # high temperature


## Visualise combined climate contribution -------------------------------------

# set up plot axes
y <- spi
x <- spi
z1 <- array(lsl_comb$tmp_low, c(length(x), length(y)))
z2 <- array(lsl_comb$tmp_med, c(length(x), length(y)))
z3 <- array(lsl_comb$tmp_hgh, c(length(x), length(y)))

# set up contour plot
pal <- rev(brewer.pal(11, "RdGy"))
levels <- pretty(c(z1, z2, z3),  20)
col1 <- colorRampPalette(pal[1:6])
col2 <- colorRampPalette(pal[6:11])
cols <- c(col1(sum(levels < 0)), col2(sum(levels > 0)))

# contour plot of lagged temperature and long-short lag spi interaction
# for cool 3-month mean temperature
g1 <- filled.contour(x, y, z1,
                     xlab = "SPI-6 long lag", ylab = "SPI-6 short lag",
                     main = paste0("Cool (3-Month Mean Temp = ",
                                   round(tmp[1] + mean_tmp["tas.3mon"], 1)," °C)"),
                     col = cols, levels = levels,
                     plot.axes = { axis(1)
                       axis(2)},
                     key.title = title(main = expression(log(RR[cc])), cex.main = 1))

# contour plot of lagged temperature and long-short lag spi interaction
# for median 3-month mean temperature
g2 <- filled.contour(x, y, z2,
                     xlab = "SPI-6 long lag", ylab = "SPI-6 short lag",
                     main = paste0("Median (3-Month Mean Temp = ",
                                   round(tmp[2] + mean_tmp["tas.3mon"], 1)," °C)"),
                     col = cols, levels = levels,
                     plot.axes = { axis(1)
                       axis(2)},
                     key.title = title(main = expression(log(RR[cc])), cex.main = 1))

# contour plot of lagged temperature and long-short lag spi interaction
# for warm 3-month mean temperature
g3 <- filled.contour(x, y, z3,
                     xlab = "SPI-6 long lag", ylab = "SPI-6 short lag",
                     main = paste0("Warm (3-Month Mean Temp = ",
                                   round(tmp[3] + mean_tmp["tas.3mon"], 1)," °C)"),
                     col = cols, levels = levels,
                     plot.axes = { axis(1)
                       axis(2)},
                     key.title = title(main = expression(log(RR[cc])), cex.main = 1))


## Sensitivity analysis --------------------------------------------------------

# calculate combined covariate effect on dengue risk in barbados
lsl_comb <- ci.cce(tmp, 1, 1000, lsl_comb, new_col="tmp_low")  # low temperature
lsl_comb <- ci.cce(tmp, 2, 1000, lsl_comb, new_col="tmp_med")  # med temperature
lsl_comb <- ci.cce(tmp, 3, 1000, lsl_comb, new_col="tmp_hgh")  # high temperature

# add column about whether each combination non-zero credible interval
lsl_comb <- ss.cce(lsl_comb, col="tmp_low")
lsl_comb <- ss.cce(lsl_comb, col="tmp_med")
lsl_comb <- ss.cce(lsl_comb, col="tmp_hgh")

# plot non-zero credible interval as X on mean plots
A_low <- lsl_comb[lsl_comb$tmp_low.ss == TRUE, "spi6.long"]
B_low <- lsl_comb[lsl_comb$tmp_low.ss == TRUE, "spi6.short"]

A_med <- lsl_comb[lsl_comb$tmp_med.ss == TRUE, "spi6.long"]
B_med <- lsl_comb[lsl_comb$tmp_med.ss == TRUE, "spi6.short"]

A_hgh <- lsl_comb[lsl_comb$tmp_hgh.ss == TRUE, "spi6.long"]
B_hgh <- lsl_comb[lsl_comb$tmp_hgh.ss == TRUE, "spi6.short"]

levels <- pretty(c(z1, z2, z3),  20)
col1 <- colorRampPalette(pal[1:6])
col2 <- colorRampPalette(pal[6:11])
cols <- c(col1(sum(levels < 0)), col2(sum(levels > 0)))

#png('outputs/figs/F_04B.png', width=600, height=480)
filled.contour(x, y, z1,
               xlab = "SPI-6 long lag", ylab = "SPI-6 short lag",
               main = paste0("Cool (3-Month Mean Temp = ",
                             round(tmp[1] + mean_tmp["tas.3mon"], 1)," °C)"),
               col = cols, levels = levels,
               plot.axes = { axis(1)
                 axis(2)
                 points(A_low, B_low, pch = 4, 
                        col=rgb(red = 0, green = 0, blue = 0.3, alpha = 0.8))},
               key.title = title(main = expression(log(RR[cc])), cex.main = 1))
#dev.off()

filled.contour(x, y, z2,
               xlab = "SPI-6 long lag", ylab = "SPI-6 short lag",
               main = paste0("Median (3-Month Mean Temp = ",
                             round(tmp[2] + mean_tmp["tas.3mon"], 1)," °C)"),
               col = cols, levels = levels,
               plot.axes = { axis(1)
                 axis(2)
                 points(A_med, B_med, pch = 4,
                        col=rgb(red = 0, green = 0, blue = 0.3, alpha = 0.8))},
               key.title = title(main = expression(log(RR[cc])), cex.main = 1))

#png('outputs/figs/F_04C.png', width=600, height=480)
filled.contour(x, y, z3,
               xlab = "SPI-6 long lag", ylab = "SPI-6 short lag",
               main = paste0("Warm (3-Month Mean Temp = ",
                             round(tmp[3] + mean_tmp["tas.3mon"], 1)," °C)"),
               col = cols, levels = levels,
               plot.axes = { axis(1)
                 axis(2)
                 points(A_hgh, B_hgh, pch = 4,
                        col=rgb(red = 0, green = 0, blue = 0.3, alpha = 0.8))},
               key.title = title(main = expression(log(RR[cc])), cex.main = 1))
#dev.off()


## END