library(mgcv)
library(dplyr)
library(marginaleffects)
library(gratia)
library(ggplot2); theme_set(theme_bw())

# Load some pre-prepared Portal rodent data and filter to only include the Desert
# Pocket Mouse
portal_ts <- read.csv('https://raw.githubusercontent.com/nicholasjclark/EFI_seminar/main/data/portal_data.csv', as.is = T) %>%
  dplyr::filter(species == 'PP')

dplyr::glimpse(portal_ts)

# Number of timepoints
max(portal_ts$time)

# Plot the capture time series, noting there is clear seasonality
# (though this can be hard to see due to the fact that some observations are NA)
ggplot(portal_ts, aes(x = time, y = captures)) +
  geom_point() + 
  geom_smooth(method = gam, formula = y ~ s(x, k = 60))

# Building on from the distributed lag example, we can set up a model that
# allows the effect of the predictor to change smoothly over moving averages
# of lagged temperature. First, a function to create the matrix of lagged moving 
# averages of different orders
mv_mat <- function(x, n_lag = 6) {
  n <- length(x)
  X <- matrix(NA, n, n_lag)
  for (i in 1:n_lag){
    # Aligning to the 'right' in zoo::rollmean() will ensure each 
    # value represents an average of the last i preceeding values
    X[1:n, i] <- zoo::rollmean(x, 
                               k = i, 
                               na.pad = TRUE, 
                               align = 'right')
  } 
  X
}

# Create a matrix of lagged minimum temperature moving average values, with lags up to 
# 9 lunar months in the past
mintemp_mat <- mv_mat(portal_ts$mintemp, 9)
dim(mintemp_mat)
head(mintemp_mat) # Note the NAs; these will be silently omitted when fitting a model

# View the increasingly smoothed predictors
matplot(mintemp_mat, type = 'l')

# Carrying on from the interactions script, we can use a tensor product of this lag
# matrix and another matrix that indexes by lag
lag_mat <- matrix(0:8, nrow(portal_ts), 9, byrow = TRUE)
dim(lag_mat)
head(lag_mat)

# Gather the data in list() format as before
lag_data <- list(
  captures = portal_ts$captures,
  mintemp_mat = mintemp_mat,
  lag_mat = lag_mat
)

# Fit with bam() so we can include autocorrelated residuals as before
mod <- bam(captures ~ 
             te(mintemp_mat, lag_mat),
           data = lag_data,
           family = nb(),
           rho = 0.8,
           discrete = TRUE)
summary(mod)
appraise(mod, method = 'simulate', type = 'deviance')
draw(mod)

# Predict the mean response values (expectations, which ignore the overdispersion
# parameter) and plot against observations
preds <- predict(mod, newdata = lag_data, se.fit = TRUE,
                 type = 'response')
portal_ts %>%
  bind_cols(data.frame(pred = preds$fit,
                       upper = preds$fit + 2 * preds$se.fit,
                       lower = preds$fit - 2 * preds$se.fit)) %>%
  ggplot(., aes(x = time, y = captures)) +
  geom_line(aes(y = pred), linewidth = 1, alpha = 0.4) +
  geom_ribbon(aes(ymax = upper,
                  ymin = lower),
              alpha = 0.4) +
  geom_point()

# Is this model a better fit compared to the distributed lag model?
lagard <- function(x, n_lag = 6) {
  n <- length(x)
  X <- matrix(NA, n, n_lag)
  for (i in 1:n_lag){
    X[i:n, i] <- x[i:n - i + 1]
  } 
  X
}
mintemp_mat <- lagard(portal_ts$mintemp, 9)
lag_data <- list(
  captures = portal_ts$captures,
  mintemp_mat = mintemp_mat,
  lag_mat = lag_mat
)
mod2 <- bam(captures ~ 
              te(mintemp_mat, lag_mat),
            data = lag_data,
            family = nb(),
            rho = 0.8,
            discrete = TRUE)
AIC(mod); AIC(mod2)
anova(mod, mod2, test = 'Chisq')

#### Bonus tasks ####
# 1. Compare this model to the distributed lag example using AIC() and 
#    anova.gam(...). Why might these models give similar results?

# 2. If you have mvgam installed, refit mod but use a poisson() observation
#    family instead (simultaneously estimating AR1 and overdispersion parameters is 
#    very challenging!). Be sure to exclude the rows of data in which the lag matrices
#    are NAs first
library(mvgam)
mintemp_mat <- mv_mat(portal_ts$mintemp, 9)
lag_mat <- matrix(0:8, nrow(portal_ts), 9, byrow = TRUE)

# Gather the data in list() format as before
lag_data <- list(
  captures = portal_ts$captures,
  mintemp_mat = mintemp_mat,
  lag_mat = lag_mat
)
idx <- which(complete.cases(mintemp_mat))
lag_data_mvgam <- list(
  captures = portal_ts$captures[idx],
  mintemp_mat = mintemp_mat[idx, ],
  lag_mat = lag_mat[idx, ],
  time = portal_ts$time[idx],
  series = factor(rep('series1', length(idx)))
)
mod2 <- mvgam(captures ~ 1,
              trend_formula = ~ te(mintemp_mat, lag_mat),
              data = lag_data_mvgam,
              trend_model = AR(),
              family = nb())
summary(mod2, include_betas = FALSE)
plot(mod2, type = 'residuals')

#    Plot the estimated distributed lag smooth
plot(mod2, type = 'smooths', trend_effects = TRUE)
draw(mod2$trend_mgcv_model)

#    Plot the hindcast distribution
plot(hindcast(mod2))

#    Use pp_check() to see how the model would predict if we 
#    assumed the dynamic AR process was stationary
pp_check(mod2, type = 'ribbon')

# Some additional plots / investigations
pp_check(mod2, type = 'dens_overlay', ndraws = 20)
pp_check(mod2, type = 'ecdf_overlay', ndraws = 20)
mcmc_plot(mod2, variable = 'ar', regex = TRUE, type = 'hist')
mcmc_plot(mod2, variable = 'phi', regex = TRUE)