final.size <- function(R0, contact.mat, eigenmax) {
  G <- contact.mat
  #G <- G/max(eigen(G)$values) # normalised to max eigenvalue
  G <- G/eigenmax # normalised to max eigenvalue of default matrix with relinf
  s0 <- rep(1, nrow(contact.mat))
  
  fnr <- function(zzs) s0*(1-exp(-R0*apply(G*zzs, MARGIN = 2, sum))) - zzs
  tG <- t(G)
  zs <- s0*(1-exp(-R0*apply(G*s0, MARGIN = 2, sum)))
  fz <- fnr(zs)
  
  ctr <- 1
  while(max(abs(fz)) > 1e-09 & ctr < 200) {
    tmp <- apply(G*zs, MARGIN = 2, sum)
    Jac <- R0*s0*exp(-R0*tmp) 
    Jac <- tG*Jac - diag(nrow(tG))
    
    invJ <- solve(Jac)
    zs <- zs - as.vector(invJ%*%fz)
    fz <- fnr(zs)
    ctr <- ctr + 1
  }
  
  return(attack = zs)
}

contactpattern2matrix <- function(contactpattern.data, var.name = "c_smt") {
  sex <- contactpattern.data %>% names %>% str_detect("gender") %>% any
  if(sex) {
    contactpattern.data <- contactpattern.data %>%   
      select(part_gender, cnt_gender, part_age, cnt_age, var.name) %>%   
      mutate(part = paste(part_gender, part_age),
             cnt = paste(cnt_gender, cnt_age)) %>%
      select(-part_age, - part_gender, -cnt_age, -cnt_gender) 
  } else {
    contactpattern.data <- contactpattern.data %>% 
      select(part_age, cnt_age, var.name) %>%
      mutate(part = part_age,
             cnt = cnt_age) %>%
      select(-part_age, -cnt_age) 
  }
  part_order <- contactpattern.data$part %>% unique
  contact.mat <- contactpattern.data %>% spread(key = part, value = var.name)
  rownames <- contact.mat$cnt
  contact.mat <- contact.mat %>% select(-cnt) %>% as.matrix 
  rownames(contact.mat) <- rownames
  contact.mat <- contact.mat[part_order, part_order]
  return(t(contact.mat))
}

countContactsLoc <- function(contact.data, 
                          pop.data,
                          loc = NULL,
                          age.breaks = c(seq(from = 0, to = 99, by = 1), Inf),
                          sex = TRUE) {
  
  # check if cnt_age_cor available
  if(!("cnt_age_cor" %in% names(contact.data))) stop("Variable ", "\"", "cnt_age_cor", "\"", " not found: run ", "\"", "correctContactAgeDigitPreference", "\""," first")
  
  # Count participants, contacts and population ----
  
  # Participants
  part_count.data <- contact.data %>%
    filter(!is.na(cnt_location)) %>% # ignore 108 contacts without location
    # Filter out duplicated participants (i.e. the contacts), because we are counting participants here
    filter(!duplicated(part_id)) %>% 
    # Cut part_age in provided age categories, open on the right [left, right)
    mutate(part_age = part_age %>% cut(breaks = age.breaks,right = FALSE, include.lowest = TRUE)) %>% 
    # Group by age categories, drop none if they appear to be empty
    group_by(part_age, .drop = FALSE) %>% 
    # Group by part_gender if sex = TRUE
    {if (sex) group_by(., part_gender, add = TRUE, .drop = FALSE) else .} %>% 
    # Count number of participants
    count(name = "n_part")

  # Contacts
  cont_count.data <- contact.data %>%
    filter(!is.na(cnt_location)) %>% # ignore 108 contacts without location
    {if(length(loc) != 0) filter(., cnt_location == loc) else .} %>%
    # Filter out missing cnt_age_cor or cnt_gender. These contacts are useless
    # Then drop empty level "Missing" from cnt_gender
    filter(!is.na(cnt_age_cor) & cnt_gender != "(Missing)") %>% 
    droplevels %>% 
    # Cut part_age and cont_age_cor in provided age categories, open on the right [left, right)
    mutate(
      part_age = part_age    %>% cut(breaks = age.breaks, right = FALSE, include.lowest = TRUE),
      cnt_age  = cnt_age_cor %>% cut(breaks = age.breaks, right = FALSE, include.lowest = TRUE)) %>% 
    # Group by age categories, drop none if they appear to be empty
    group_by(part_age, cnt_age, .drop = FALSE) %>% 
    # Group by part_gender and cnt_gender if sex = TRUE
    {if (sex) group_by(., part_gender, cnt_gender, add = TRUE, .drop = FALSE) else .} %>% 
    # Count number of contacts
    count(name = "n_cnt")
  
  # Population
  pop_count.data <- pop.data %>% 
    # Cut age in provided age categories, open on the right [left, right)
    mutate(age = age %>% cut(breaks = age.breaks, right = FALSE, include.lowest = TRUE)) %>% 
    # Group by age categories, drop none if they appear to be empty
    group_by(age, .drop = FALSE) %>% 
    # Group by gender if sex = TRUE, and make sure gender factor levels from contact.data and pop.data are aligned
    {if (sex) group_by(., gender = gender %>% fct_relevel(levels(part_count.data$part_gender)), add = TRUE, .drop = FALSE)
      else .} %>% 
    # Count the population number
    summarize(n_pop = sum(population)) %>%
    ungroup %>%
    # Calculate population fraction
    mutate(f_pop = n_pop/sum(n_pop))
  
  # Join the tibbles above into one tibble ----
  # Calculate observed contact intensities and rates
  
  # Create contactpattern.data by joining part_cont.data and cont_count.data
  contactpattern.data <- 
    full_join(part_count.data, cont_count.data) %>%
    # The next three steps depend on sex. Therefore we use an if statemant
    # 1. Join with pop_count.data
    #    Note: because the number of contacts is proportional with population,
    #          explicitly join age and gender in pop_count.data with that in cont_count.data
    # 2. For modelling the contact rates, be sure the records are arranged as given
    # 3. Reorder columns as given
    {if (sex) {
      full_join(., pop_count.data, by = c("cnt_age" = "age", "cnt_gender" = "gender")) %>% 
        arrange(part_gender, cnt_gender, part_age, cnt_age) %>% 
        select(part_gender, cnt_gender, part_age, cnt_age, n_part, n_pop, f_pop, n_cnt)
    } else {
      full_join(., pop_count.data, by = c("cnt_age" = "age")) %>% 
        arrange(part_age, cnt_age) %>% 
        select(part_age, cnt_age, n_part, n_pop, f_pop, n_cnt)
    }} %>%
    mutate(
      # Calculate observed contact intensties (m_obs) and rates (c_obs)
      # Note: based on population fractions and not population numbers!
      m_obs = n_cnt/n_part,
      c_obs = m_obs/f_pop) %>% 
    # Ungroup data
    ungroup
  
  # for age categories of 1 year, use numbers for participant and contact age
  if(all(diff(age.breaks) == Inf | diff(age.breaks) == 1)) {
    levels(contactpattern.data$part_age) <- age.breaks[-length(age.breaks)]
    levels(contactpattern.data$cnt_age) <- age.breaks[-length(age.breaks)]
  }
  
  return(contactpattern.data)  
}

countExposureContactsLoc <- function(contact.data, 
                                  pop.data,
                                  loc = NULL,
                                  age.breaks = c(seq(from = 0, to = 80, by = 10), Inf),
                                  sex = FALSE) {
  
  # check if cnt_age_cor available
  if(!("cnt_age_cor" %in% names(contact.data))) stop("Variable ", "\"", "cnt_age_cor", "\"", " not found: run ", "\"", "correctContactAgeDigitPreference", "\""," first")
  
  # Count participants, contacts and population ----
  
  # Participants
  part_count.data <- contact.data %>%
    filter(!is.na(cnt_location)) %>% # ignore 108 contacts without location
    # Filter out duplicated participants (i.e. the contacts), because we are counting participants here
    filter(!duplicated(part_id)) %>% 
    # Cut part_age in provided age categories, open on the right [left, right)
    mutate(part_age = part_age %>% cut(breaks = age.breaks,right = FALSE, include.lowest = TRUE)) %>% 
    # Group by age categories, drop none if they appear to be empty
    group_by(part_age, .drop = FALSE) %>% 
    # Group by part_gender if sex = TRUE
    {if (sex) group_by(., part_gender, add = TRUE, .drop = FALSE) else .} %>% 
    # Count number of participants
    count(name = "n_part")
  
  # Contacts
  cont_count.data <- contact.data %>%
    filter(!is.na(cnt_location)) %>% # ignore 108 contacts without location
    {if(length(loc) != 0) filter(., cnt_location == loc) else .} %>%
    # Filter out missing cnt_age_cor or cnt_gender. These contacts are useless
    # Then drop empty level "Missing" from cnt_gender
    filter(!is.na(cnt_age_cor) & cnt_gender != "(Missing)") %>% 
    droplevels %>% 
    # Cut part_age and cont_age_cor in provided age categories, open on the right [left, right)
    mutate(
      duration = duration_multi %>% fct_recode("5" = "less than 5 minutes", 
                                               "10" = "5 - 15 mins",
                                               "30" = "15 mins - 1 hour",
                                               "120" =  "1 hour - 4 hours",
                                               "240" = "4 hours or more",
                                               "60" = "(Missing)") %>% as.character %>% as.numeric,
      part_age = part_age    %>% cut(breaks = age.breaks, right = FALSE, include.lowest = TRUE),
      cnt_age  = cnt_age_cor %>% cut(breaks = age.breaks, right = FALSE, include.lowest = TRUE)) %>% 
    # Group by age categories, drop none if they appear to be empty
    group_by(part_age, cnt_age, .drop = FALSE) %>% 
    # Group by part_gender and cnt_gender if sex = TRUE
    {if (sex) group_by(., part_gender, cnt_gender, add = TRUE, .drop = FALSE) else .} %>% 
    # Count number of contacts
    # count(name = "n_cnt")
    summarise(n_cnt = sum(duration)/60)
  
  # Population
  pop_count.data <- pop.data %>% 
    # Cut age in provided age categories, open on the right [left, right)
    mutate(age = age %>% cut(breaks = age.breaks, right = FALSE, include.lowest = TRUE)) %>% 
    # Group by age categories, drop none if they appear to be empty
    group_by(age, .drop = FALSE) %>% 
    # Group by gender if sex = TRUE, and make sure gender factor levels from contact.data and pop.data are aligned
    {if (sex) group_by(., gender = gender %>% fct_relevel(levels(part_count.data$part_gender)), add = TRUE, .drop = FALSE)
      else .} %>% 
    # Count the population number
    summarize(n_pop = sum(population)) %>%
    ungroup %>%
    # Calculate population fraction
    mutate(f_pop = n_pop/sum(n_pop))
  
  # Join the tibbles above into one tibble ----
  # Calculate observed contact intensities and rates
  
  # Create contactpattern.data by joining part_cont.data and cont_count.data
  contactpattern.data <- 
    full_join(part_count.data, cont_count.data) %>%
    # The next three steps depend on sex. Therefore we use an if statemant
    # 1. Join with pop_count.data
    #    Note: because the number of contacts is proportional with population,
    #          explicitly join age and gender in pop_count.data with that in cont_count.data
    # 2. For modelling the contact rates, be sure the records are arranged as given
    # 3. Reorder columns as given
    {if (sex) {
      full_join(., pop_count.data, by = c("cnt_age" = "age", "cnt_gender" = "gender")) %>% 
        arrange(part_gender, cnt_gender, part_age, cnt_age) %>% 
        select(part_gender, cnt_gender, part_age, cnt_age, n_part, n_pop, f_pop, n_cnt)
    } else {
      full_join(., pop_count.data, by = c("cnt_age" = "age")) %>% 
        arrange(part_age, cnt_age) %>% 
        select(part_age, cnt_age, n_part, n_pop, f_pop, n_cnt)
    }} %>%
    mutate(
      # Calculate observed contact intensties (m_obs) and rates (c_obs)
      # Note: based on population fractions and not population numbers!
      m_obs = n_cnt/n_part,
      c_obs = m_obs/f_pop) %>% 
    # Ungroup data
    ungroup
  
  # for age categories of 1 year, use numbers for participant and contact age
  if(all(diff(age.breaks) == Inf | diff(age.breaks) == 1)) {
    levels(contactpattern.data$part_age) <- age.breaks[-length(age.breaks)]
    levels(contactpattern.data$cnt_age) <- age.breaks[-length(age.breaks)]
  }
  
  return(contactpattern.data)  
}

ContactPatternLocation <- function(contact.data, pop.data, age.breaks = c(seq(from = 0, to = 80, by = 10), Inf), sex = FALSE, countWhat = countContactsLoc) {
  contactpattern <- countWhat(contact.data = contact.data,
                              pop.data = pop.data,
                              age.breaks = age.breaks,
                              sex = sex)
  contactpattern <- modelContacts(contactpattern.data = contactpattern, posterior.sample = 0, verbose = FALSE)
  for(loc in levels(contact.data$cnt_location)) {
    # tmp <- contact.data %>% filter(cnt_location == loc)
    tmp <- countWhat(contact.data = contact.data,
                     pop.data = pop.data,
                     loc = loc,
                     age.breaks = age.breaks,
                     sex = sex)
    # tmp <- modelContacts(contactpattern.data = tmp, posterior.sample = 0, verbose = TRUE)
    contactpattern[[paste0("n_cnt_",loc)]] <- tmp$n_cnt
  }
  # contactpattern[["m_smt_tot"]] <- rowSums(contactpattern %>% select(paste0("m_smt_", levels(contact3$cnt_location))))
  return(contactpattern)
}


EffectOfMeasureOnContacts <- function(contactpattern, par = tibble(H_1=1, H_2=1, H_3=1, H_4=1, H_5=1, H_6=1, H_7=1, H_8=1, H_9=1,
                                                                   W_1=1, W_2=1, W_3=1, W_4=1, W_5=1, W_6=1, W_7=1, W_8=1, W_9=1,
                                                                   S_1=1, S_2=1, S_3=1, S_4=1, S_5=1, S_6=1, S_7=1, S_8=1, S_9=1,
                                                                   T_1=1, T_2=1, T_3=1, T_4=1, T_5=1, T_6=1, T_7=1, T_8=1, T_9=1,
                                                                   L_1=1, L_2=1, L_3=1, L_4=1, L_5=1, L_6=1, L_7=1, L_8=1, L_9=1,
                                                                   O_1=1, O_2=1, O_3=1, O_4=1, O_5=1, O_6=1, O_7=1, O_8=1, O_9=1,
                                                                   I_1=1, I_2=1, I_3=1, I_4=1, I_5=1, I_6=1, I_7=1, I_8=1, I_9=1)) {
  
  tmp <- contactpattern %>% 
    # effect on contact rate
    mutate(n_cnt_Home = ifelse(as.integer(part_age) == 1, par$H_1*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 2, par$H_2*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 3, par$H_3*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 4, par$H_4*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 5, par$H_5*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 6, par$H_6*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 7, par$H_7*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 8, par$H_8*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(part_age) == 9, par$H_9*n_cnt_Home, n_cnt_Home),
           n_cnt_Work = ifelse(as.integer(part_age) == 1, par$W_1*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 2, par$W_2*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 3, par$W_3*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 4, par$W_4*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 5, par$W_5*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 6, par$W_6*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 7, par$W_7*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 8, par$W_8*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(part_age) == 9, par$W_9*n_cnt_Work, n_cnt_Work),
           n_cnt_School = ifelse(as.integer(part_age) == 1, par$S_1*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 2, par$S_2*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 3, par$S_3*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 4, par$S_4*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 5, par$S_5*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 6, par$S_6*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 7, par$S_7*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 8, par$S_8*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(part_age) == 9, par$S_9*n_cnt_School, n_cnt_School),
           n_cnt_Transport = ifelse(as.integer(part_age) == 1, par$T_1*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 2, par$T_2*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 3, par$T_3*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 4, par$T_4*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 5, par$T_5*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 6, par$T_6*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 7, par$T_7*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 8, par$T_8*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(part_age) == 9, par$T_9*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 1, par$L_1*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 2, par$L_2*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 3, par$L_3*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 4, par$L_4*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 5, par$L_5*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 6, par$L_6*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 7, par$L_7*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 8, par$L_8*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(part_age) == 9, par$L_9*n_cnt_Leisure, n_cnt_Leisure),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 1, par$O_1*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 2, par$O_2*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 3, par$O_3*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 4, par$O_4*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 5, par$O_5*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 6, par$O_6*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 7, par$O_7*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 8, par$O_8*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(part_age) == 9, par$O_9*`n_cnt_Other place`, `n_cnt_Other place`),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 1, par$H_1*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 2, par$H_2*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 3, par$H_3*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 4, par$H_4*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 5, par$H_5*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 6, par$H_6*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 7, par$H_7*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 8, par$H_8*n_cnt_Home, n_cnt_Home),
           n_cnt_Home = ifelse(as.integer(cnt_age) == 9, par$H_9*n_cnt_Home, n_cnt_Home),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 1, par$W_1*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 2, par$W_2*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 3, par$W_3*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 4, par$W_4*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 5, par$W_5*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 6, par$W_6*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 7, par$W_7*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 8, par$W_8*n_cnt_Work, n_cnt_Work),
           n_cnt_Work = ifelse(as.integer(cnt_age) == 9, par$W_9*n_cnt_Work, n_cnt_Work),
           n_cnt_School = ifelse(as.integer(cnt_age) == 1, par$S_1*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 2, par$S_2*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 3, par$S_3*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 4, par$S_4*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 5, par$S_5*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 6, par$S_6*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 7, par$S_7*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 8, par$S_8*n_cnt_School, n_cnt_School),
           n_cnt_School = ifelse(as.integer(cnt_age) == 9, par$S_9*n_cnt_School, n_cnt_School),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 1, par$T_1*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 2, par$T_2*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 3, par$T_3*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 4, par$T_4*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 5, par$T_5*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 6, par$T_6*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 7, par$T_7*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 8, par$T_8*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Transport = ifelse(as.integer(cnt_age) == 9, par$T_9*n_cnt_Transport, n_cnt_Transport),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 1, par$L_1*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 2, par$L_2*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 3, par$L_3*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 4, par$L_4*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 5, par$L_5*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 6, par$L_6*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 7, par$L_7*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 8, par$L_8*n_cnt_Leisure, n_cnt_Leisure),
           n_cnt_Leisure = ifelse(as.integer(cnt_age) == 9, par$L_9*n_cnt_Leisure, n_cnt_Leisure),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 1, par$O_1*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 2, par$O_2*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 3, par$O_3*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 4, par$O_4*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 5, par$O_5*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 6, par$O_6*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 7, par$O_7*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 8, par$O_8*`n_cnt_Other place`, `n_cnt_Other place`),
           `n_cnt_Other place` = ifelse(as.integer(cnt_age) == 9, par$O_9*`n_cnt_Other place`, `n_cnt_Other place`))

  return(tmp)
}



EffectOfMeasureOnInf <- function(mat, par = tibble(H_1=1, H_2=1, H_3=1, H_4=1, H_5=1, H_6=1, H_7=1, H_8=1, H_9=1,
                                                   W_1=1, W_2=1, W_3=1, W_4=1, W_5=1, W_6=1, W_7=1, W_8=1, W_9=1,
                                                   S_1=1, S_2=1, S_3=1, S_4=1, S_5=1, S_6=1, S_7=1, S_8=1, S_9=1,
                                                   T_1=1, T_2=1, T_3=1, T_4=1, T_5=1, T_6=1, T_7=1, T_8=1, T_9=1,
                                                   L_1=1, L_2=1, L_3=1, L_4=1, L_5=1, L_6=1, L_7=1, L_8=1, L_9=1,
                                                   O_1=1, O_2=1, O_3=1, O_4=1, O_5=1, O_6=1, O_7=1, O_8=1, O_9=1,
                                                   I_1=1, I_2=1, I_3=1, I_4=1, I_5=1, I_6=1, I_7=1, I_8=1, I_9=1)) {

  Ieff <- c(par$I_1, par$I_2, par$I_3, par$I_4, par$I_5, par$I_6, par$I_7, par$I_8, par$I_9)

  return(Ieff*mat)
}



EffectOnTransmission <- function(contactpattern, measures, relinf, locations) {
  
  eigenmax <- max(eigen(relinf*contactpattern2matrix(contactpattern.data = contactpattern, var.name = "m_smt"))$value)
  
  tmp <- contactpattern
  for(i in 1:nrow(measures)) {
    input <- measures[i,] #%>% select(-measure, -measureID)
    tmp <- EffectOfMeasureOnContacts(tmp, par = input)
  }
  tmp$n_cnt <- rowSums(tmp %>% select(paste0("n_cnt_", locations)))
  tmp <- modelContacts(contactpattern.data = tmp, posterior.sample = 0, verbose = FALSE)
  tmp <- relinf*contactpattern2matrix(tmp, var.name = "m_smt")
  
  for(i in 1:nrow(measures)) {
    input <- measures[i,] #%>% select(-measure, -measureID)
    tmp <- EffectOfMeasureOnInf(tmp, par = input)
  }
  
  neweigenmax <- max(eigen(tmp)$value)
  
  return(100*(neweigenmax - eigenmax)/eigenmax)
}

InfectiousContactPattern <- function(contactpattern, locations, relinf) {
  for(loc in locations) {
    contactpattern[[paste0("m_smt_inf_", loc)]] <- as.vector(t(relinf*contactpattern2matrix(contactpattern, var.name = paste0("m_smt_", loc))))
  }
  contactpattern[["m_smt_inf"]] <- as.vector(t(relinf*contactpattern2matrix(contactpattern, var.name = "m_smt")))
  contactpattern[["m_smt_inf_tot"]] <- as.vector(t(relinf*contactpattern2matrix(contactpattern, var.name = "m_smt_tot")))
  return(contactpattern)
}

plotContactMatrix <- function(contactpattern.data, var.name = "c_smt", n.ticks = 20, plot.contour = FALSE, min = NA, max = NA) {
  
  sex <- contactpattern.data %>% names %>% str_detect("gender") %>% any
  
  n.levels <- length(levels(contactpattern.data$part_age))
  which2plot <- rep(FALSE, n.levels)
  which2plot[seq(1, n.levels, by = ceiling(n.levels/n.ticks))] <- TRUE
  
  rate <- str_detect(var.name, pattern = "c_")
  smoothed <- str_detect(var.name, pattern = "smt")
  
  data = contactpattern.data %>% 
    mutate(variable = contactpattern.data[[var.name]])
  
  log10max <- if(is.na(max)) ceiling(log10(max(data$variable, na.rm = TRUE))) else log10(max)
  log10min <- if(is.na(min)) floor(log10(min(data$variable[data$variable > 0], na.rm = TRUE))) else log10(min)
  
  p <- ggplot(
    data = data,
    mapping = aes(x = part_age, y = cnt_age, fill = variable, z = variable)) +
    geom_tile() +
    # scale_fill_distiller(
    #   breaks = c(0, outer(c(1, 2, 5), 10^(log10min:log10max))),
    #   palette = "YlOrRd",
    #   na.value = NA,
    #   trans = if(smoothed) "log" else "sqrt",
    #   direction = 1) +
    scale_fill_viridis_c(
      limits = 10^c(log10min, log10max),
      breaks = c(0, outer(c(1, 2, 5), 10^(floor(log10min):ceiling(log10max)))),
      na.value = "#450256",
      trans = if(smoothed) "log" else "sqrt",
      direction = 1) +
    coord_equal() +
    labs(
      x = "Participant age",
      y = "Contact age",
      #fill = if(rate) expression("Contact\nrate") else expression("Contact\nintensity")) +
      fill = if(rate) expression("Contact\nrate") else expression("Number of\ncontacts per\nparticipant")) +
    scale_x_discrete(
      breaks = levels(contactpattern.data$part_age)[which2plot]) +
    scale_y_discrete(
      breaks = levels(contactpattern.data$cnt_age)[which2plot]) +
      {if(sex) facet_grid(
        rows = vars(cnt_gender),
        cols = vars(part_gender %>% fct_rev))} +
    theme_bw() +
    theme(
      legend.position = "none",
      panel.grid = element_blank(),
      legend.key.height = unit(0.09, "npc"),
      axis.text.x = element_text(
        angle = 90,
        hjust = 1,
        vjust = 0.5))
  
  if(plot.contour) {
    p <- p + 
      geom_contour(
        aes(x = as.integer(part_age), 
            y = as.integer(cnt_age)),
        size = 0.1,
        colour = "black",
        breaks = c(outer(c(1, 2, 5), 10^(log10min:log10max))))
  }
  return(p)
}