#####################################################
# Author: Johannes Bracher, johannes.bracher@uzh.ch

# This reproduces a shortened version of the analysis in Bracher and Held (2017):
# Periodically stationary multivariate autoregressive models.

# install.packages("surveillance")
# # installation of the hhh4addon package from github:
# library(devtools)
# install_github("jbracher/hhh4addon", build_vignettes = TRUE)

library(surveillance)
library(hhh4addon)
library(RColorBrewer)
library(forecast)

######################################################
# Getting and plotting the data:
data("noroBL")
noroBL <- noroBL[1:312, ] # restrict to subset used in article

par(mfrow = c(1, 2))
plot(noroBL@observed[, 1], type = "h", main = "(a) Bremen", las = 1,
     xlab = "time", ylab = "weekly no. of cases", axes = FALSE); box()
axis(1, at = 0:6*51 + 1, labels = 2011:2017)
axis(2, las = 1)

plot(noroBL@observed[, 2], type = "h", main = "(b) Lower Saxony", las = 1,
     xlab = "time", ylab = "weekly no. of cases", axes = FALSE); box()
axis(1, at = 0:6*51 + 1, labels = 2011:2017)
axis(2, las = 1)

######################################################
# analysis using geometric lags:

# defining control, christmas break and offsets:
offsets_ne <- matrix(c(67, 790), ncol = 2, nrow = nrow(noroBL@observed), byrow = TRUE)
x <- matrix(1*((1:nrow(noroBL@observed))%%52 %in% c(0, 1)), ncol = 2, nrow = nrow(noroBL@observed))
x_ne <- x; x_ne[,2] <- 0
control_lg <- list(end = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE), S = c(1, 1))),
                   ar = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE) + fe(x, unitSpecific = TRUE), S = c(1, 1))),
                   ne = list(f = addSeason2formula(~ 1, S = 1), offset = offsets_ne),
                   subset = 6:nrow(noroBL@observed), family = "NegBinM")
# fit model:
fit_lg <- profile_par_lag(noroBL, control_lg)

# get stationary moments:
stat_mom_lg <- hhh4addon:::stationary_moments(fit_lg, return_mu_decomposed = TRUE, return_Sigma = TRUE, n_seasons = 2)

# Plots of fit (simplified, the plots in the article are notproduced by the generic functions):
plot(fit_lg, unit = 1)
plot(fit_lg, unit = 2)

# Plot of stationary means and observations:
hhh4addon:::plotHHH4lag_stationary(fit_lg, unit = 1)
hhh4addon:::plotHHH4lag_stationary(fit_lg, unit = 2)

# fit simpler models:
# no cross-lags:
control_lg_ar <- list(end = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE), S = c(1, 1))),
                      ar = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE) + fe(x, unitSpecific = TRUE), S = c(1, 1))),
                      subset = 6:nrow(noroBL@observed), family = "NegBinM")
fit_lg_ar <- profile_par_lag(noroBL, control_lg_ar)
# no geometric lags:
control_l1 <- list(end = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE), S = c(1, 1))),
                   ar = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE) + fe(x, unitSpecific = TRUE), S = c(1, 1))),
                   ne = list(f = addSeason2formula(~ 1, S = 1), offset = offsets_ne),
                   subset = 6:nrow(noroBL@observed), family = "NegBinM")
fit_l1 <- hhh4(noroBL, control_l1)
# neither cross-lags nor geometric lags:
control_l1_ar <- list(end = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE), S = c(1, 1))),
                      ar = list(f = addSeason2formula(~ -1 + fe(1, unitSpecific = TRUE) + fe(x, unitSpecific = TRUE), S = c(1, 1))),
                      subset = 6:nrow(noroBL@observed), family = "NegBinM")
fit_l1_ar <- hhh4(noroBL, control_l1_ar)

# compare the four AIC values
AIC(fit_lg)
AIC(fit_lg_ar)
AIC(fit_l1)
AIC(fit_l1_ar)

##############################################################
# The remainder of this file produces some of the figures in the manuscript.
# These are not generated by plotting functions from the package
# but are specifically written for the manuscript
# The code is therefore somewhat more lengthy.


##############################################################
# Elaborate plots of unconditional properties: Fig. 3

# matrices containing realizations from one year per row:
matr_bremen <- matrix(c(noroBL@observed[, "Bremen"], rep(NA, 52 - nrow(noroBL@observed)%%52)), ncol = 52, byrow = TRUE)
matr_ls <- matrix(c(noroBL@observed[, "Lower.Saxony"], rep(NA, 52 - nrow(noroBL@observed)%%52)), ncol = 52, byrow = TRUE)

# get direct empirical estimates of means and sds:
emp_means_bremen <- colMeans(matr_bremen, na.rm = TRUE)
emp_means_ls <- colMeans(matr_ls, na.rm = TRUE)
emp_sds_bremen <- apply(matr_bremen, 2, sd, na.rm = TRUE)
emp_sds_ls <- apply(matr_ls, 2, sd, na.rm = TRUE)

###########################################
# Plot of stationary means (a and b):
# helper function:
ominus <- function(a, b, l){
  ifelse(a > b, a - b, l - ((b - a)%%l))
}

# get parameters:
param <- hhh4addon:::lambda_tilde_complex_neighbourhood(fit_lg, periodic = TRUE)
mu <- stat_mom_lg$mu_matrix
n_units <- ncol(mu)
length_of_period <- 52
max_lag <- length(fit_lg$distr_lag)
timepoints <- seq(from = 2011 + 6/52, by = 1/52, length.out = 308)

# compute the decomposition of the mean:
end <- stat_mom_lg$mu_decomposed[,, "endemic"]
ar_contributions <- ne_contributions <- array(dim = c(length_of_period, n_units, max_lag))
dimnames(ar_contributions) <- dimnames(ne_contributions) <- list(NULL, c("Bremen", "Lower Saxony"), paste0("lag", 1:max_lag))

for(t in 1:length_of_period){
  lambda_temp <- hhh4addon:::only_ar(param$lambda[,,t])
  phi_temp <- param$lambda[,,t] - lambda_temp
  mu_preceding <- mu[ominus(t, max_lag:1, length_of_period), ]
  for(lag in 1:max_lag){
    inds <- seq(to = n_units*(max_lag - lag + 1), length.out = n_units)
    lambda_this_lag <- lambda_temp; lambda_this_lag[,-inds] <- 0
    phi_this_lag <- phi_temp; phi_this_lag[,-inds] <- 0

    ar_contributions[t, , lag] <- lambda_this_lag%*%as.vector(t(mu_preceding))
    ne_contributions[t, , lag] <- phi_this_lag%*%as.vector(t(mu_preceding))
  }
}


# arrange all in one array:
end_ar_ne <- array(dim = c(dim(ar_contributions)[1:2], 2*dim(ar_contributions)[3] + 1))
end_ar_ne[,,1] <- end[1:52, ]
end_ar_ne[,,2:6] <- ar_contributions
end_ar_ne[,,7:11] <- ne_contributions

# apply cumulative sum:
# Bremen:
stacked_temp_B <- t(apply(end_ar_ne[,1,], 1, cumsum))
# Lower Saxony:
stacked_temp_L <- t(apply(end_ar_ne[,2,], 1, cumsum))

par(mfrow = c(3, 2), mar = c(4, 4.5, 3, 1))

inds_plot <- c(27:52, 1:26)
lty_shading <- "solid"
cols <- brewer.pal(6, "Paired")
labels_weeks <- c(27, 40, 1, 14, 26); at_weeks <- c(1, 14, 27, 40, 52)

plot(NULL, xlim = c(1, 52), ylim = c(0, 40), xlab = "calendar week t", ylab = expression(hat(mu[t])), las = 1, main = "(a) Stationary means, Bremen", axes = FALSE)
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()

polygon(c(1, 1:52, 52),
        c(0, stat_mom_lg$mu_matrix[inds_plot, 1], 0),
        col = cols[2])
polygon(c(1, 1:52, 52),
        c(0, stacked_temp_B[inds_plot, 7], 0),
        col = cols[1], border = cols[1])
polygon(c(1, 1:52, 52),
        c(0, stat_mom_lg$mu_decomposed[inds_plot, 1, 1] + stat_mom_lg$mu_decomposed[inds_plot, 1, 2], 0),
        col = cols[6], border = cols[6])
polygon(c(1, 1:52, 52),
        c(0, stacked_temp_B[inds_plot, 2], 0),
        col = cols[5], border = cols[5])


polygon(c(1, 1:52, 52), c(0, stat_mom_lg$mu_decomposed[inds_plot, 1, 1], 0), col = "grey", border = "grey")
lines(stat_mom_lg$mu_matrix[inds_plot, 1])
points(emp_means_bremen[inds_plot], pch = 23, cex = 0.5, bg = "white")

legend("topleft", col = c(NA, cols[5], cols[1], "grey", cols[6], cols[2], "black"),
       legend = c("Decomposition:", "same reg., lag 1", "other reg., lag 1", "endemic",
                  "same reg., lag > 1", "other reg., lag > 1"), pch = 15, ncol = 2,
       bty = "n")

plot(NULL, xlim = c(1, 52), ylim = c(0, 400), xlab = "calendar week t", ylab = expression(hat(mu[t])), las = 1, main = "(b) Stationary means, Lower Saxony", axes = FALSE)
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()

polygon(c(1, 1:52, 52),
        c(0, stat_mom_lg$mu_decomposed[inds_plot, 2, 1] + stat_mom_lg$mu_decomposed[inds_plot, 2, 2], 0),
        col = cols[6], border = cols[6])
polygon(c(1, 1:52, 52),
        c(0, stacked_temp_L[inds_plot, 2], 0),
        col = cols[5], border = cols[5])

polygon(c(1, 1:52, 52), c(0, stat_mom_lg$mu_decomposed[inds_plot, 2, 1], 0), col = "grey", border = "grey")
lines(stat_mom_lg$mu_matrix[inds_plot, 2])
points(emp_means_ls[inds_plot], pch = 23, cex = 0.5, bg = "white")

legend("topleft", legend = c("model-based", "empirical"),
       lty = c(1, NA), pch = c(NA, 23), bty = "n", pt.cex = 0.5)

###########################################
# plot of standard deviations (c and d)
plot(NULL, xlim = c(1, 52), ylim = c(0, 20), xlab = "calendar week t", ylab = expression(hat(sigma[t])), las = 1, axes = FALSE,
     main = "(c) Stationary standard deviations, Bremen")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
lines(sqrt(stat_mom_lg$var_matrix[inds_plot, 1]), lty = "solid")
points(emp_sds_bremen[inds_plot], pch = 23, cex = 0.5, bg = "white")

plot(NULL, xlim = c(1, 52), ylim = c(0, 200), xlab = "calendar week t", ylab = expression(hat(sigma[t])), las = 1, axes = FALSE,
     main = "(d) Stationary standard deviations, Lower Saxony")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
lines(sqrt(stat_mom_lg$var_matrix[inds_plot, 2]))
points(emp_sds_ls[inds_plot], pch = 23, cex = 0.5, bg = "white")
legend("topleft", legend = c("model-based", "empirical"),
       lty = c(1, NA), pch = c(NA, 23), bty = "n", pt.cex = 0.5)

###########################################
# plot of approximated distributions (e and f)
plot(NULL, xlim = c(1, 52), ylim = c(0, 80), xlab = "calendar week t", ylab = expression(Y[t]), las = 1, axes = FALSE,
     main = "(e) Approximated stationary distributions, Bremen")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
stat_mom_lg_oneS <- stationary_moments(fit_lg, return_mu_decomposed = TRUE, return_Sigma = TRUE,
                                                   start = 27)
fanplot_stationary(stat_mom_lg_oneS, unit = 1, timepoints = 1:52, add = TRUE, mean_lty = "solid", mean_col = NA)
matr_bremen2 <- matrix(c(rep(NA, 26), noroBL@observed[, "Bremen"], rep(NA, 26)), byrow = TRUE, ncol = 52)
for(i in 1:7){
  lines(1:52, matr_bremen2[i, ], pch = 23, cex = 0.5, bg = "white", type = "b")
}

plot(NULL, xlim = c(1, 52), ylim = c(0, 800), xlab = "calendar week t", ylab = expression(Y[t]), las = 1, axes = FALSE,
     main = "(e) Approximated stationary distributions, Lower Saxony")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
fanplot_stationary(stat_mom_lg_oneS, unit = 2, timepoints = 1:52, add = TRUE, mean_lty = "solid", mean_col = NA)
matr_ls2 <- matrix(c(rep(NA, 26), noroBL@observed[, "Lower.Saxony"], rep(NA, 26)), byrow = TRUE, ncol = 52)
for(i in 1:7){
  lines(1:52, matr_ls2[i, ], pch = 23, cex = 0.5, bg = "white", type = "b")
}
y_legend <- matrix(seq(from = 250,
                       to = 750, length.out = 99),
                   ncol = 2, nrow = 99)
library(fanplot)
fan(y_legend, start = 50,  fan.col = colorRampPalette(c("darkgreen",
                                                        "gray90")), rlab = c(1, 50, 99)/100, ln = 0.5)
text(44, 730, "stat. quantiles:")
legend("topleft", legend = "observations", pch = 23, bty = "n", pt.cex = 0.5)



##############################################################33
# Plot of auto- and cross-correlations (Fig  4)

# define colours and line types:
col_b <- "black" # "darkgreen"
col_l <- "black" # "darkred"

col_lag0 <- "black" # "darkblue"
col_lag1 <- "black" # "cornflowerblue"
col_lag2 <- "black" # "cyan"

lty_lag0 <- "dotted"
lty_lag1 <- "solid"
lty_lag2 <- "dashed"

# move from covariances to correlations:
cor_matr <- cov2cor(stat_mom_lg$Sigma)

# indices corresponding to observations from Bremen and Lower Saxony, respectively:
inds_b <- 2*1:104 - 1
inds_l <- 2*1:104
inds_plot <- 27:(26 + 52)

par(mfrow = c(2, 2), mar = c(4, 4.5, 3, 1))

# get the auto- and cross-correlation matrices for
cor_matr_bb <- cor_matr[inds_b, inds_b] # Bremen vs lagged Bremen
cor_matr_bl <- cor_matr[inds_b, inds_l] # Bremen vs lagged Lower Saxony
cor_matr_lb <- cor_matr[inds_l, inds_b] # Lower Saxony vs lagged Bremen
cor_matr_ll <- cor_matr[inds_l, inds_l] # Lower Saxony vs lagged Lower Saxony

# helper function to get the backwards auto-correlations:
extract_backwards_corr <- function(cor_matr, inds, lag){
  inds_rows <- inds - lag
  inds_cols <- inds
  ret <- diag(cor_matr[inds_rows, inds_cols])
  names(ret) <- paste(rownames(cor_matr)[inds_rows], colnames(cor_matr)[inds_cols], sep = "--")
  ret
}

bb0 <- extract_backwards_corr(cor_matr = cor_matr_bb, inds = inds_plot, lag = 0)
bb1 <- extract_backwards_corr(cor_matr = cor_matr_bb, inds = inds_plot, lag = 1)
bb2 <- extract_backwards_corr(cor_matr = cor_matr_bb, inds = inds_plot, lag = 2)
plot(NULL, xlim = c(1, 52), ylim = 0:1, axes = FALSE, xlab = "calendar week t", ylab = expression(Cor(Y[Bt], Y[B~","~t-d])), main = "(a) Autocorrelation Bremen")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
lines(bb0, col = col_lag0, lwd = 2, lty = lty_lag0)
lines(bb1, col = col_lag1, lwd = 2, lty = lty_lag1)
lines(bb2, col = col_lag2, lwd = 2, lty = lty_lag2)

lb0 <- extract_backwards_corr(cor_matr = cor_matr_lb, inds = inds_plot, lag = 0)
lb1 <- extract_backwards_corr(cor_matr = cor_matr_lb, inds = inds_plot, lag = 1)
lb2 <- extract_backwards_corr(cor_matr = cor_matr_lb, inds = inds_plot, lag = 2)

# draw the plots:
plot(NULL, xlim = c(1, 52), ylim = 0:1, type = "l", axes = FALSE, xlab = "calendar week t", ylab = expression(Cor(Y[Bt], Y[L~","~t-d])),
     main = "(b) Cross-correlations")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
lines(lb0, col = col_lag0, lwd = 2, lty = lty_lag0)
lines(lb1, col = col_lag1, lwd = 2, lty = lty_lag1)
lines(lb2, col = col_lag2, lwd = 2, lty = lty_lag2)

legend("topleft", legend = c("d = 0", "d = 1", "d = 2"), lty = c("dotted", "solid", "dashed"), bty = "n")

bl0 <- extract_backwards_corr(cor_matr = cor_matr_bl, inds = inds_plot, lag = 0)
bl1 <- extract_backwards_corr(cor_matr = cor_matr_bl, inds = inds_plot, lag = 1)
bl2 <- extract_backwards_corr(cor_matr = cor_matr_bl, inds = inds_plot, lag = 2)
plot(NULL, xlim = c(1, 52), ylim = 0:1, type = "l", axes = FALSE, xlab = "calendar week t", ylab = expression(Cor(Y[Lt], Y[B~","~t-d])),
     col = col_lag0, main = "(c) Cross-correlations")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
lines(bl0, col = col_lag0, lwd = 2, lty = lty_lag0)
lines(bl1, col = col_lag1, lwd = 2, lty = lty_lag1)
lines(bl2, col = col_lag2, lwd = 2, lty = lty_lag2)

ll0 <- extract_backwards_corr(cor_matr = cor_matr_ll, inds = inds_plot, lag = 0)
ll1 <- extract_backwards_corr(cor_matr = cor_matr_ll, inds = inds_plot, lag = 1)
ll2 <- extract_backwards_corr(cor_matr = cor_matr_ll, inds = inds_plot, lag = 2)
plot(NULL, xlim = c(1, 52), ylim = 0:1, type = "l", axes = FALSE, xlab = "calendar week t", ylab = expression(Cor(Y[Lt], Y[L~","~t-d])),
     main = "(d) Autocorrelation Lower Saxony")
abline(v = 27, col = "grey")
axis(1, labels = labels_weeks, at = at_weeks); axis(2, las = 1); box()
lines(ll0, col = col_lag0, lwd = 2, lty = lty_lag0)
lines(ll1, col = col_lag1, lwd = 2, lty = lty_lag1)
lines(ll2, col = col_lag2, lwd = 2, lty = lty_lag2)


##############################################################
# Plot of residual autocorrelations (Fig 5)
par(mfrow = c(2, 2), mar = c(4, 4.5, 3, 1))

# compute the conditional variances (given past):
vars_fitted <- exp(-fit_lg$coefficients[c("-log(overdisp.Bremen)",
                                                   "-log(overdisp.Lower.Saxony)")])*fit_lg$fitted.values^2 +
  fit_lg$fitted.values
# compute Pearson residuals:
res_pearson_fitted <- residuals(fit_lg, type = "response")/sqrt(vars_fitted)

# compute Pearson residuals relative to stationary means and variances
stat_mom_lg2 <- stationary_moments(fit_lg, n_seasons = 7)
res_pearson_stationary <- ((noroBL@observed - stat_mom_lg2$mu_matrix)/sqrt(stat_mom_lg2$var_matrix))[-(1:5), ]

# compute auto-correlations in the two types of residuals:
acf_stationary <- Acf(res_pearson_stationary, lag.max = 15, xaxp = c(1, 21, 5), plot = FALSE)
acf_fitted <- Acf(res_pearson_fitted, xaxp = c(1, 21, 15), lag.max = 15, plot = FALSE)

yaxp <- c(-0.2, 0.8, 5)
m_lag <- 7

# draw the plots
plot(1:m_lag, acf_stationary$acf[2:(m_lag + 1),1,1],type = "h", main = "(a) Bremen", xlab = "d",
     ylab = "resid. autocorrelation", ylim = c(-0.2, 0.8),
     col = "grey", yaxp = yaxp, las = 1, xlim = c(0, m_lag))
thresh <- 1.96/sqrt(acf_fitted$n.used)
points(1:m_lag, acf_stationary$acf[2:(m_lag + 1), 1, 1], pch = 22, bg = "darkgrey", cex = 0.7)
points(1:m_lag, acf_fitted$acf[2:(m_lag + 1), 1, 1], type = "h")
points(1:m_lag, acf_fitted$acf[2:(m_lag + 1), 1, 1], pch = 15, cex = 0.7)
abline(h = 0)
abline(h = c(-thresh, thresh), lty = "dotted")
legend("topright", legend = c(expression(paste("unconditional Pearson resid. ", r[gt]^paste("u"))),
                              expression(paste("conditional Pearson resid. ", r[gt]))),
       col = c("black", "black"), pt.bg = c("darkgrey", NA),
       pch = c(22, 15), pt.cex = 0.7, bty  = "n")

plot(0:m_lag, acf_stationary$acf[1:(m_lag + 1),1,2],type = "h", main = "(b) Bremen & Lower Saxony", xlab = "d", ylab = "residual cross-correlation", ylim = c(-0.2, 0.8), col = "darkgrey", yaxp = yaxp, las = 1)
points(0:m_lag, acf_stationary$acf[1:(m_lag + 1), 1, 2], pch = 22, bg = "darkgrey", cex = 0.7)
points(0:m_lag, acf_fitted$acf[1:(m_lag + 1), 1, 2], type = "h")
points(0:m_lag, acf_fitted$acf[1:(m_lag + 1), 1, 2], pch = 15, cex = 0.7)
abline(h = 0)
abline(h = c(-thresh, thresh), lty = "dotted")

plot(0:m_lag, acf_stationary$acf[1:(m_lag + 1),2,1],type = "h", main = "(c) Lower Saxony & Bremen", xlab = "d",
     ylab = "resid. cross-correlation", ylim = c(-0.2, 0.8), col = "darkgrey", yaxp = yaxp, las = 1)
points(0:m_lag, acf_stationary$acf[1:(m_lag + 1), 2, 1], pch = 22, bg = "darkgrey", cex = 0.7)
points(0:m_lag, acf_fitted$acf[1:(m_lag + 1), 2, 1], type = "h")
points(0:m_lag, acf_fitted$acf[1:(m_lag + 1), 2, 1], pch = 15, cex = 0.7)
abline(h = 0)
abline(h = c(-thresh, thresh), lty = "dotted")

plot(1:m_lag, acf_stationary$acf[2:(m_lag + 1),2,2],type = "h", main = "(d) Lower Saxony", xlab = "d", ylab = "residual autocorrelation", ylim = c(-0.2, 0.8),
     col = "grey", yaxp = yaxp, las = 1, xlim = c(0, m_lag))
points(1:m_lag, acf_stationary$acf[2:(m_lag + 1), 2, 2], pch = 22, bg = "darkgrey", cex = 0.7)
points(1:m_lag, acf_fitted$acf[2:(m_lag + 1), 2, 2], type = "h")
points(1:m_lag, acf_fitted$acf[2:(m_lag + 1), 2, 2], pch = 15, cex = 0.7)
abline(h = 0)
abline(h = c(-thresh, thresh), lty = "dotted")