# ''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''''
## FORECAST MODEL ==============================================================

# Description:
#     Produce dengue risk forecast 3 months ahead in Barbados using final model.
#     The year index is calculated such that the forecast issue month is month 9 
#     and the target is month 12. This forecast was issued as of March 2024 for 
#     dengue outbreak risk in June 2024 for the ICC Men's T20 Cricket World Cup.

# 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/prediction/barbados_fc_2024-03.csv")
data$dir <- round(data$cases / data$population * 10**5, 3)


## Set model parameters and update dataset -------------------------------------

# forecast values
vars <- c("tas.3mon.4", "spi6.5", "spi6.1")
fc_issue_year <- 2024  # calendar year forecast issued
fc_issue_month <- 3  # calendar month forecast issued
fc_issue_month_txt <- ifelse(fc_issue_month < 10, paste0("0", fc_issue_month),
                             as.character(fc_issue_month))  # add leading zero

fc_horizon <- 3
month <- ifelse(fc_horizon + fc_issue_month > 12,  # actual target month value
                (fc_horizon + fc_issue_month) %% 12,
                fc_horizon + fc_issue_month)

# update dataframe
start_id <- match(month + 1, data$cal_month)
data <- data[start_id:nrow(data), ]

nyear <- length(unique(data$cal_year)) - 1
data$month <- rep(1:12, nyear, length.out = nrow(data))
data$year_index <- rep(1:nyear, each=12, length.out = nrow(data))

# extract last year index and population
current_year <- max(data$year_index)
pop <- data[nrow(data), "population"]


## Extract parameters of updated model -----------------------------------------

# 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 samples from posterior distribution
s <- 1000
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 negative binomial distribution means ------------------------------

# create empty array for distribution means
fit <- array(NA,
             dim = c(s, 1),
             dimnames = list(paste0("s.", 1:s), paste0("t", 1)))

# calculate distribution means
fit[1:s, 1] <-
  exp(log(pop/100000) +
        alpha +
        beta.1 * data[nrow(data), vars[1]] +
        beta.2 * data[nrow(data), vars[2]] +
        beta.3 * data[nrow(data), vars[3]] +
        beta.4 * data[nrow(data), fe] +
        beta.12 * data[nrow(data), vars[1]] * data[nrow(data), vars[2]] +
        beta.13 * data[nrow(data), vars[1]] * data[nrow(data), vars[3]] +
        beta.23 * data[nrow(data), vars[2]] * data[nrow(data), vars[3]] +
        beta.123 * data[nrow(data), vars[1]] * data[nrow(data), vars[2]] * data[nrow(data), vars[3]] +
        delta[12, ] +
        gamma[(current_year), ])


## Generate posterior predictive distribution ----------------------------------

# create empty array for generated data points
pred <- array(NA,
              dim = c(s, 1),
              dimnames = list(paste0("s.", 1:s), paste0("t", 1)))

# generate data points from negative binomial distribution
for (n in 1:s){
  pred[n,1] <- rnbinom(1, mu = fit[n,1], size = theta[n])
}

# save pred output
fname <- glue("outputs/prediction/pred_t20_{fc_issue_year}{fc_issue_month_txt}.rds")
#saveRDS(pred, fname)


## Produce forecast visualisations ---------------------------------------------

# read in pred
pred <- readRDS(fname)
max_year_index <- max(data$year_index, na.rm = TRUE)

# calculate moving monthly 75th percentile threshold
pct_75 <- data %>%
  filter(year_index < max_year_index) %>%  # excl. prediction year
  group_by(month) %>%
  summarise(q75 = round(quantile(dir, 0.75) * pop / 10**5, 2))

m <- 12
oprob <- round(sum(pred[,1] >= pct_75[m, "q75"][[1]]) * 100 / s, 3)

# plot gauge
q75 <- pct_75[12, 2]
dens <- density(pred, from = min(pred)) %$%
  data.frame(x = x, y = y) %>%
  mutate(area = x >= q75[[1]])

g1 <- gg.gauge(oprob, breaks = c(0, 29, 50, 70, 100))

# plot probability density function 
g2 <- ggplot(dens, aes(x = x, ymin = 0, ymax = y, fill = area)) + 
  geom_ribbon() + 
  scale_fill_manual(values = c("grey90", "#5BC0EB")) +
  geom_line(aes(y = y)) +
  geom_vline(xintercept = q75[[1]], color = "#AA1155") +
  labs(x="Predicted Dengue Cases", y="Density") +
  theme_bw(base_size = 14) +
  theme(panel.border = element_blank(), plot.title = element_text(hjust=0.5),
        panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        legend.position = "none")

# combined risk visualisations
ggdraw() +
  draw_plot(g1, x = 0, y = 0.5, width = 1, height = 0.5) +
  draw_plot(g2, x = 0, y = 0, width = 1, height = 0.5) +
  draw_plot_label(label = c("A", "B"), size = 15,
                  x = c(0, 0), y = c(1, .5)) +
  theme_bw(base_size = 16) +
  theme(panel.border = element_blank(), plot.title=element_text(hjust=0.5),
        panel.grid.major = element_blank(), panel.grid.minor = element_blank(),
        legend.position="none", axis.text=element_blank(), axis.ticks=element_blank())
#ggsave("outputs/figs/F_05.png", width=10, height=10, dpi=300)


## END