# JAGS model for nowcasting
#
# Author: ___UITZONDERINGSGROND_2___ - RIVM
#
# This model is a slightly modified version of the Höhle and an der Heiden (2012) model
# The original model can be found on https://github.com/cran/surveillance/blob/master/inst/jags/bhpm.bugs
#
# Remark about indexing:
# Since d = 0 ... D and JAGS cannot handle [0] as an index, we use d + 1 in the indices
# The value of d still is 0 ... D

jags_nowcast <- function() {
  
  #
  # Likelihood ----
  #
  
  # Loop over all time points in the moving window
  # t is defined here the time point in the moving window, i.e. t = 1:T_window
  for (t in 1:T_window) {
    
    #
    # Reporting trapezoid and corresponding epicuve ----
    #
    
    # The number of cases in the reporting trapezoid y_rep[t, d] has a Negative Binomial distribution
    # with parameters mu_rep[t, d] (= mean) and theta_rep (= overdispersion parameter)
    # mu_rep[t, d] itself is the product of the fraction reported p_rep[t, d] and the expected epicurve mu_epi[t]
    for (d in 0:D) {
      y_rep[t, d + 1] ~ dnegbin(prob_rep[t, d + 1], theta_rep)
      mu_rep[t, d + 1] <- p_rep[t, d + 1]*mu_epi[t]
      
      # Note that JAGS uses the parametrisation of the Neg. Bin. distribution in terms of a
      # probability and overdispersion parameter
      prob_rep[t, d + 1] <- theta_rep/(theta_rep + mu_rep[t, d + 1])
    }
    
    # Epicurve model. This is a smooth function of time
    log(mu_epi[t]) <- beta_epi[1] +
      inprod(X_epi[t, ], beta_epi[2:n_beta_epi]) + inprod(Z_epi[t, ], b_epi[])
    
    # Nowcast. This is the sum of y_rep[t, d] over all delays d = 0 ... D at time t
    y_epi[t] <- sum(y_rep[t, ])
    
    #
    # Time-dependent delay distribution ----
    #
    
    # The delay distribution (i.e. the fraction reported) p_rep[t, d] is function of:
    # - Delay (obviously)
    # - Time (optionally)
    # - Covariates, such as day of the week
    #
    # The delay distribution is modelled by a discrete time-to-event model
    # The underlying hazard (in terms of a log-odds) is smooth function of delay and time (in terms of log-odds ratios)
    
    # For d = 0
    p_rep[t, 1] <- h[t, 1] # By definition
    
    # For d = 1 ... D - 1
    for (d in 1:(D - 1)) {
      p_rep[t, d + 1] <- (1 - sum(p_rep[t, 1:d]))*h[t, d + 1]
    }
    
    # For d = D
    h[t, D + 1] <- 1 # By definition. Everything has been reported at d = D
    p_rep[t, D + 1] <- (1 - sum(p_rep[t, 1:D]))*h[t, D + 1]
  }
  
  # The log(OR) function for delay d = 0:(D - 1)
  for (d in 0:(D - 1)) {
    logit_h_delay_d[d + 1] <- inprod(X_delay_d[d + 1, ], beta_delay_d[]) + inprod(Z_delay_d[d + 1, ], b_delay_d[])
  }
  for (t in 1:T_window) {
    # The log(OR) function for time t = 1:T_window
    logit_h_delay_t[t] <- inprod(X_delay_t[t, ], beta_delay_t[]) + inprod(Z_delay_t[t, ], b_delay_t[])
    for (d in 0:(D - 1)) {
      # The log(OR) for the covariates for time t = 1:T_window and delay d = 0:(D - 1)
      logit_h_delay_c[t, d + 1] <- inprod(X_delay_c[t, d + 1, ], beta_delay_c[])
      # The total log-odds is the sum of an intercept and the log(OR)'s
      logit(h[t, d + 1]) <- beta_delay +
        #logit_h_delay_t[t] +
        logit_h_delay_d[d + 1] +
        logit_h_delay_c[t, d + 1]
    }
  }
  
  #
  # Priors ----
  #
  
  # Prior overdispersion parameter
  theta_rep <- exp(log.theta_rep)
  log.theta_rep ~ dnorm(0, 0.01)
  
  # Priors epicuve P-spline in time
  for (i in 1:n_beta_epi) {
    beta_epi[i] ~ dnorm(0, 0.01)
  }
  for (i in 1:n_b_epi) {
    b_epi[i] ~ dnorm(0, tau_b_epi)
  }
  tau_b_epi ~ dscaled.gamma(5, 1)
  sigma_b_epi <- 1/sqrt(tau_b_epi)
  
  # Prior delay distribution intercept
  beta_delay ~ dnorm(0, 0.01)
  
  # Priors delay distribution P-spine in delay
  for (i in 1:n_beta_delay_d) {
    beta_delay_d[i] ~ dnorm(0, 0.01)
  }
  for (i in 1:n_b_delay_d) {
    b_delay_d[i] ~ dnorm(0, tau_b_delay_d)
  }
  tau_b_delay_d ~ dscaled.gamma(5, 1)
  sigma_b_delay_d <- 1/sqrt(tau_b_delay_d)
  
  # Priors delay distribution P-spine in time
  for (i in 1:n_beta_delay_t) {
    beta_delay_t[i] ~ dnorm(0, 0.01)
  }
  for (i in 1:n_b_delay_t) {
    b_delay_t[i] ~ dnorm(0, tau_b_delay_t)
  }
  tau_b_delay_t ~ dscaled.gamma(5, 1)
  sigma_b_delay_t <- 1/sqrt(tau_b_delay_t)
  
  # Priors delay distribution covariates
  for (i in 1:n_beta_delay_c) {
    beta_delay_c[i] ~ dnorm(0, 0.01)
  }
  
  # End model
}