##############################################################################
#
# Description
# Estimates smooth and symmetric contact rates

# Arguments
# data = contact pattern data as generated by countContacts
# verbose = whether output is verbose
# posterior.sample = number of posterior samples to take after analysis
# 
# Details
# The function will automatically detect if data is sex-specific
#
# Value
# A tibble with contact pattern data:
# Two variables are added: m_smt and c_smt, smoothed versions of m_obs and c_obs
# n posterior samples are added as c_smt.1 ... c_smt.n
#
# Authors
# ___UITZONDERINGSGROND_2___ & ___UITZONDERINGSGROND_2___ - RIVM
#
##############################################################################

modelContacts <- function(contactpattern.data, verbose = FALSE, posterior.sample = 0) {

  # Initialize INLA ----

  # Check for INLA package
  # If not available, download and install it
  if (!"INLA" %in% .packages(all = TRUE)) {
    install.packages("INLA",
                     repos = "https://inla.r-inla-download.org/R/stable",
                     dependencies = TRUE)
  }
  
  library(INLA)
  
  # this uses inla.static to circumvent error message about missing GLIBC version
  INLA:::inla.dynload.workaround()
  
  # Define helper functions (within function, not pretty but they're not needed outside) ----
  
  # Function to construct node IDs in matrix form
  construct.nodeID <- function(n, sex = TRUE, vector = TRUE) {
    if (sex) {
      # Sex specific
      nn1 <- n*(n + 1)/2 # Number of nodes in MM and FF
      nn2 <- n^2         # Number of nodes in FM and MF
      node.MM <- node.FF <- matrix(0, nrow = n, ncol = n) # Empty matrix
      LT <- lower.tri(node.MM, diag = TRUE) # Lower tri of MM and FF with diag
      node.MM[ LT] <- 1:nn1           # Fill lower tri with node IDs
      node.MM[!LT] <- t(node.MM)[!LT] # Fill upper tri with transposed
      node.FM      <- matrix(nn1 + 1:nn2, nrow = n, ncol = n) # Fill entire matrix
      node.MF      <- t(node.FM)                              # Transposed
      node.FF[ LT] <- nn1 + nn2 + 1:nn1 # As node.MM
      node.FF[!LT] <- t(node.FF)[!LT]
      nodeID <- list(node.MM = node.MM, node.FM = node.FM, node.MF = node.MF, node.FF = node.FF)
    } else {
      # Not sex specific
      nn <- n*(n + 1)/2
      node <- matrix(0, nrow = n, ncol = n)
      LT <- lower.tri(node, diag = TRUE)
      node[ LT] <- 1:nn
      node[!LT] <- t(node)[!LT]
      nodeID <- list(node = node)
    }
    
    # Return as vector or as list of matrices?
    if (vector) {
      return(nodeID %>% unlist(use.names = FALSE))
    } else {
      return(nodeID)
    }
  }
  # Function to construct R-matrix
  construct.Rmat <- function(n, order = 2, sex = TRUE) {
    if (sex) {
      # Sex specific
      D.MM <- D.FF <- construct.Dmat(n, order, tri = TRUE )$D # MM and FF
      D.FM         <- construct.Dmat(n, order, tri = FALSE)$D # FM
      D <- bdiag(D.MM, D.FM, D.FF) # Block diagonal difference matrix D
      R <- crossprod(D)            # Structure matrix R
      R <- as(R, "dsCMatrix")      # R is symmetric sparse (column compressed) matrix
      return(list(D = D, R = R))
    } else {
      # Not sex specific
      D <- construct.Dmat(n, order, tri = TRUE)$D
      R <- crossprod(D)
      R <- as(R, "dsCMatrix")
      return(list(D = D, R = R))
    }
  }
  # Function to construct D-matrix
  construct.Dmat <- function(n, order, tri = FALSE) {
    require(Matrix)
    In <- Diagonal(n)            # n x n matrix
    D0 <- diff(In, diff = order) # Differences for single vector
    D1 <- kronecker(In, D0)      # Differences in horizontal direction
    D2 <- kronecker(D0, In)      # Differences in   vertical direction
    if (tri) {
      # For lower triangular matrix (MM and FF)
      ri1.mat <- matrix(1:((n - order)*n), nrow = n - order, ncol = n)
      ri2.mat <- matrix(1:(n*(n - order)), nrow = n, ncol = n - order)
      ci.mat  <- matrix(1:n^2, nrow = n, ncol = n)
      ri1 <- ri1.mat[row(ri1.mat) >=  col(ri1.mat)         ] # Rows to keep in D1
      ri2 <- ri2.mat[row(ri2.mat) >= (col(ri2.mat) + order)] # Rows to keep in D2
      ci  <-  ci.mat[row(ci.mat)  >=  col(ci.mat)          ] # Columns to keep in D1 & D2
      D <- rbind(D1[ri1, ci], D2[ri2, ci])
      return(list(D0 = D0, D1 = D1, D2 = D2, ri1 = ri1, ri2 = ri2, ci = ci, D = D))
    } else {
      # For full matrix (FM)
      D <- rbind(D1, D2)
      return(list(D0 = D0, D1 = D1, D2 = D2, D = D))
    }
  }
  
  # Settings and prepare data ----

  # Sex specific data or not: is "gender" somewhere in the column names?
  sex <- contactpattern.data %>% names %>% str_detect("gender") %>% any
  
  # Number of age categories
  n_age <- contactpattern.data %>% pull(part_age) %>% nlevels
  
  # Construct structure matrix R
  # Put penalty on second order differences (default)
  # Rank deficiency is equal to the number of zero eigenvalues of R
  ord <- 2
  rankdef <- ifelse(sex, yes = 3, no = 1)*ord^2
  R <- construct.Rmat(n = n_age, order = ord, sex = sex)$R
  
  # Some data preparation
  contactpattern.data <- contactpattern.data %>% 
    mutate(
      # Calculate denominator U = n_part x f_pop
      # If U is zero (because n_part = 0), then set U = 1 and n_cnt = NA,
      # so this record does not contribute to the likelihood 
      U = n_part*f_pop,
      U = ifelse(n_part == 0, yes = 1, no = U), 
      n_cnt = ifelse(n_part == 0, yes = NA, no = n_cnt),
      # Add node IDs, needed to force symmetric contact rates
      node.id  =  construct.nodeID(  n = n_age, sex = sex, vector = TRUE))
  
  # Model ----

  # Run model
  contact.inla <- inla(
    formula = n_cnt ~ 1 + f(node.id, model = "generic0", Cmatrix = R,
                            # log-Gamma(1, 0.1) prior on log-precision
                            hyper = list(prec = list(prior = "loggamma", param = c(1, 0.1))),
                            # Rank deficiency, sum-to-zero contstraint, and add small value to diagonal of precision
                            rankdef = rankdef, constr = TRUE, diagonal = 0.001),
    # E is the known component in the mean for the Negative binomial response likelihood
    E = U,
    # Negative binomial likelihood
    family = "nbinomial",
    # Data
    data = contactpattern.data,
    # Normal(0, 0.001) prior on intercept
    control.fixed = list(mean.intercept = 0, prec.intercept = 0.001),
    # Normal(0, 0.001) prior on log-dispersion parameter
    control.family = list(
      hyper = list(theta = list(prior = "gaussian", param = c(0, 0.001)))),
    # Integration strategy "eb" for faster computation
    control.inla = list(int.strategy = "eb"),
    # Compute linear predictor
    control.predictor = list(compute = TRUE, link = 1),
    # Other control settings
    control.compute = list(
      waic = TRUE,    # Compute WAIC
      cpo = TRUE,     # Compute cross-validated predictive measures
      config = TRUE), # Enable posterior sampling
    # Verbose output?
    verbose = verbose) 
  
  # Show summary
  print(summary(contact.inla))
  
  # Post processing ----

  contactpattern.data <- contactpattern.data %>%
    # Add smoothed c and m to data
    mutate(
      # Contact rate c = exp(linear predictor)
      c_smt = exp(contact.inla$summary.linear.predictor[, "mean"]),
      # Contact intensity m (note: based on population fraction f_pop)
      m_smt = f_pop*c_smt) %>% 
    # Remove unnecessary variables
    select(-U, -node.id)
  
  # add posterior samples
  if(posterior.sample > 0) {
    sim <- inla.posterior.sample(posterior.sample, contact.inla)
    for(i in 1:posterior.sample) {
      s <- sim[[i]]$latent[,1]
      contactpattern.data <- contactpattern.data %>% mutate(!!paste0("c_smt.",i) := exp(s[grepl("Predictor", names(s))]))
    }
  }
  
  # Return output
  return(contactpattern.data)

}