##############################################################################
# Calculate instantaneous and case reproduction number from epicurve
# - sample number of admissions from expectation
# - construct incidence of known symptom onsets on each day for each sample
# - backsample unknown symptom onsets from admission date
#   (taking into account whether infection occurred abroad or not)
# - calculate instantaneous and case reproduction number + confidence interval
##############################################################################


calculate_R <- function(epicurve, SO2adm, serial_interval = NULL, nsample = 1000) {

  if(length(serial_interval) == 0) {
    # probability density of serial interval discretized by 1 day (note: first day from 0 to 0.5)
    serial_interval <- diff(pgamma(q = c(0, seq(0.5, 10.5, 1)), scale = 1.0, shape = 4))
    serial_interval <- serial_interval/sum(serial_interval)
    names(serial_interval) <- 0:10
  }
  g <- serial_interval

  
  # sample non-observed admissions due to reporting delay from neg binomial
  admission_matrix <- matrix(0, nrow = nrow(epicurve), ncol = nsample)
  for(j in 1:nrow(epicurve)) {
    if(epicurve$admission_expected[j] != epicurve$admission[j]) {
      inctot <-  if(epicurve$admission[j] == 0) 1 else epicurve$admission[j]
      admission_matrix[j, ] <- rnbinom(n = nsample, size = inctot , prob = inctot/epicurve$admission_expected[j]) + inctot
    } else {
      admission_matrix[j, ] <- rep(epicurve$admission[j], nsample)
    }
  }
  
  # known SO make up fixed part of incidence matrix
  incidence_matrix <- matrix(rep(epicurve$incidence_knownSO, nsample), ncol = nsample)
  incidence_infectedabroad_matrix <- matrix(rep(epicurve$incidence_knownSO_infectedabroad, nsample), ncol = nsample)
  
  # sampling symtpom onsets
  for(j in 1:nrow(incidence_matrix)) {
    for(i in 1:nsample) {
      unknownSO <- admission_matrix[j,i] - epicurve$admission_knownSO[j]
      unknownSO_infectedabroad <- epicurve$admission_infectedabroad[j] - epicurve$admission_knownSO_infectedabroad[j]
      if(unknownSO != 0) {
        tmp <- table(j - sample(as.integer(names(SO2adm)), size = unknownSO - unknownSO_infectedabroad, prob = SO2adm, replace = TRUE))
        tmp <- tmp[names(tmp) > 0]
        incidence_matrix[as.integer(names(tmp)), i] <- incidence_matrix[as.integer(names(tmp)), i] + tmp
      }
      if(unknownSO_infectedabroad != 0) {
        tmp <- table(j - sample(as.integer(names(SO2adm)), size = unknownSO_infectedabroad, prob = SO2adm, replace = TRUE))
        tmp <- tmp[names(tmp) > 0]
        incidence_matrix[as.integer(names(tmp)), i] <- incidence_matrix[as.integer(names(tmp)), i] + tmp
        incidence_infectedabroad_matrix[as.integer(names(tmp)), i] <- incidence_infectedabroad_matrix[as.integer(names(tmp)), i] + tmp
      }
    }
  }
  
  # for most recent dates divide by sum(SO2adm2) to correct for symptom onset of admissions in future (implicitly assuming they remain constant)
  for(j in (nrow(incidence_matrix) - max(as.integer(names(SO2adm)))):nrow(incidence_matrix)) {
    SO2adm2 <- SO2adm[1:min(length(SO2adm), nrow(incidence_matrix)-j+1)]
    incidence_matrix[j,] <- incidence_matrix[j,]/sum(SO2adm2)
  }
  
  # construct instantaneous Rt from incidence
  Rt <- matrix(0, nrow = nrow(incidence_matrix), ncol = nsample)
  for(t in 1:nrow(incidence_matrix)) {
    g2 <- g[as.integer(names(g)) < t & as.integer(names(g)) >= t - nrow(incidence_matrix)]
    # those showing symptoms at t, but not infected abroad
    infected <- incidence_matrix[t, ] - incidence_infectedabroad_matrix[t, ]
    # are infected by (note: at day 0 (at t) possible infectors needs to be one less (as not to infect yourself))
    infectors <- incidence_matrix[t - as.integer(names(g2)), , drop = FALSE]- as.integer(names(g2) == 0)
    Rt[t,] <- as.numeric(infected/colSums(infectors*as.vector(g2)))/sum(g2) # note: dividing by sum(g2) only has effect when SI comprises negative values
  }
  Rt[Rt == Inf] <- NA # occurs when incidence != 0 but no incidence in previous days
  
  # construct case Ru from instantaneous Rt 
  Ru <- matrix(0, nrow = nrow(incidence_matrix), ncol = nsample)
  for(u in 1:nrow(Rt)) {
    g2 <- g[as.integer(names(g)) > - u & as.integer(names(g)) <= nrow(incidence_matrix) - u]
    Ru[u,] <- colSums(as.vector(g2)*Rt[u + as.integer(names(g2)), ,drop = FALSE])/sum(g2) # divide by sum(g2) to account for not observed Rt's in future
  }
  Ru[is.na(Ru)] <- 0 # only necessary when SI comprises strictly positive values
  
  # construct confidence bounds by sampling from poisson
  Rtupper <- rep(0, nrow(incidence_matrix))
  Rtlower <- rep(0, nrow(incidence_matrix))
  logRtvar <- rep(0, nrow(incidence_matrix))
  for(j in 1:nrow(incidence_matrix)) {
    # moet eigenlijk op gedeeld door (incidence_matrix[j,i] - incidence_infectedabroad_matrix[j,i]
    tmp <- sapply(1:nsample, function(i) if(incidence_matrix[j,i] != 0 & Rt[j,i] !=Inf & !is.na(Rt[j,i])) rpois(100, incidence_matrix[j,i]*Rt[j,i])/incidence_matrix[j,i])
    Rtlower[j] <- quantile(unlist(tmp), probs = 0.025)
    Rtupper[j] <- quantile(unlist(tmp), probs = 0.975)
    logRtvar[j] <- var(log(Rt[j, Rt[j,] > 0]))
  }

  # construct confidence bounds by sampling from poisson
  Ruupper <- rep(0, nrow(incidence_matrix))
  Rulower <- rep(0, nrow(incidence_matrix))
  logRuvar <- rep(0, nrow(incidence_matrix))
  for(j in 1:nrow(incidence_matrix)) {
    tmp <- sapply(1:nsample, function(i) if(incidence_matrix[j,i] != 0 & Ru[j,i] !=Inf) rpois(100, incidence_matrix[j,i]*Ru[j,i])/incidence_matrix[j,i])
    Rulower[j] <- quantile(unlist(tmp), probs = 0.025)
    Ruupper[j] <- quantile(unlist(tmp), probs = 0.975)
    logRuvar[j] <- var(log(Ru[j, Ru[j,] > 0]))
  }
  
  # last Ru's are 0, change to NA to prevent plotting them
  Ru[nrow(incidence_matrix),] <- NA
  Rulower[nrow(incidence_matrix)] <- NA
  Ruupper[nrow(incidence_matrix)] <- NA
  
  
  
  epicurve <- epicurve %>%
    mutate(incidence_mean = apply(incidence_matrix, MARGIN = 1, mean),
           incidence_median = apply(incidence_matrix, MARGIN = 1, median),
           incidence_lower = apply(incidence_matrix, MARGIN = 1, quantile, probs = 0.025),
           incidence_upper = apply(incidence_matrix, MARGIN = 1, quantile, probs = 0.975),
           instR = apply(Rt, MARGIN = 1, mean),
           instRlower = Rtlower,
           instRupper = Rtupper,
           caseR = apply(Ru, MARGIN = 1, mean),
           caseRlower = Rulower,
           caseRupper = Ruupper,
           instRlogvar = logRtvar,
           caseRlogvar = logRuvar
    )
  
  return(list(epicurve = epicurve, incidence = incidence_matrix))
}