#####################
### ode modelcode ###
#####################

# required: many functions and parameters, created in masterfile

# compartments: S, I, and cumulative incidence
# I has letter A, to allow more I-compartments later, each with own letter
ODEcompartments <- function() return(c("S", "IA", "cumI"))

### rate of new infections into different age classes
covidlambda <- function(y, contacts, susceptibility, infectiousness, parms) {
  res <- rep(0, 9)
  for(i in ODEcompartments()) {
    if(substr(i, 1, 1) == "I") {
      res <- res + parms[[paste0("beta", substr(i, 2, 2))]] * 
        susceptibility *
        colSums(y[paste0(i, 1:9)] * infectiousness * contacts) * y[paste0("S", 1:9)]
    }
  }
  names(res) <- paste0("la", 1:9)
  return(res)
}

### model equations of SIR model: S, I, and cumulative incidence
covidODEmodel <- function(t, y, parms) {
  la <- covidlambda(y, parms$contacts, parms$susceptibility, parms$infectiousness, parms)
  
  res <- rep(0, length(y))
  names(res) <- names(y)
  
  # dS/dt
  for(i in 1:9) {
    res[paste0("S", i)] <- -la[paste0("la", i)]
  }
  
  # dIA/dt
  for(i in 1:9) {
    res[paste0("IA", i)] <- la[paste0("la", i)] - parms[["gammaA"]] * y[paste0("IA", i)]
  }
  
  # dcumI/dt
  for(i in 1:9) {
    res[paste0("cumI", i)] <- la[paste0("la", i)]
  }
  
  # finished
  return(list(res))
}

### function for a single simulation, returning time series of all compartments
covidsim <- function(contactcontrol = c("99", "99"),
                     infectivitycontrol = c("99", "99"),
                     endtimes = c(30, 250),
                     seeding = 10,
                     normalisedby = "Symp",
                     fittedrelinfs = 1,  
                     fittedrelSIs = 1, ...) {
  ### relinfs, relSIs
  if(length(fittedrelinfs) == 1) {
    fittedrelinfs <- rep(fittedrelinfs, length(endtimes))
  }
  if(length(fittedrelSIs) == 1) {
    fittedrelSIs <- rep(fittedrelSIs, length(endtimes))
  }
  
  ### scale transmission rates
  susceptibilities <- relsusceptibility(normalisedby, ...)
  iniinfectivities <- relinfectivity(normalisedby, ...)
  relbeta <- 1/eigen(t(contactmatrix(contactcontrol[1], ...) * 
                         agedist(...) * susceptibilities) * 
                       iniinfectivities, only.values = T)[[1]][1]
  nathistparms <- parmsNatHist(relinfstart = infectivitycontrol[1], ...)
  for(i in names(nathistparms)) {
    if(substr(i, 1, 4) == "beta") nathistparms[[i]] <- nathistparms[[i]] * relbeta
  }
  
  ### define all compartments
  compartments <- t(outer(ODEcompartments(), 1:9, paste0))
  
  ### initialize compartments
  initialstate <- rep(0, length(compartments))
  names(initialstate) <- compartments
  imports <- rep(seeding/9, 9) / popsize(...)
  initialstate[paste0("S", 1:9)] <- agedist(...) - imports
  initialstate[paste0("IA", 1:9)] <- imports
  
  ### simulate no control
  nathistparms2use <- nathistparms
  nathistparms2use$betaA <- nathistparms2use$betaA * fittedrelinfs[1]
  nathistparms2use$gammaA <- nathistparms2use$gammaA / fittedrelSIs[1]
  ressimul <- ode(
    func = covidODEmodel,
    y = initialstate,
    times = seq(0, endtimes[1], 1),
    parms = c(list(contacts = contactmatrix(contactcontrol[1], ...),
                   susceptibility = susceptibilities,
                   infectiousness = relinfectivity(normalisedby = normalisedby,
                                                   scenario = infectivitycontrol[1], ...)),
              nathistparms2use)
  )
  
  ### simulate controls
  for(i in 2:length(contactcontrol)) {
    nathistparms2use <- nathistparms
    nathistparms2use$betaA <- nathistparms2use$betaA * fittedrelinfs[i]
    nathistparms2use$gammaA <- nathistparms2use$gammaA / fittedrelSIs[i]
    ressimul <- rbind(
      head(ressimul, -1),
      ode(
        func = covidODEmodel,
        y = tail(ressimul, 1)[, compartments],
        times = seq(endtimes[i - 1], endtimes[i], 1),
        parms = c(list(contacts = contactmatrix(contactcontrol[i], ...),
                       susceptibility = susceptibilities,
                       infectiousness = relinfectivity(normalisedby = normalisedby,
                                                       scenario = infectivitycontrol[i], ...)),
                  nathistparms2use)
      )
    )
  }
  
  return(ressimul)
}  


### function for a single simulation, starting 12 Feb (day 30 = 13 March, start of control), 
# returning observed numbers of cases in age classes:
#   * symptomatic incidence (could be reported)
#   * hospitalisation incidence
#   * ICU incidence
#   * immunity prevalence
#   * hospitalisation prevalence
#   * ICU prevalence
covidcontrolsim <- function(contactcontrol = c("99", "99"),
                            infectivitycontrol = c("99", "99"),
                            endtimes = c(30, 250),
                            seeding = 10, ...) {
  result <- covidsim(contactcontrol = contactcontrol,
                     infectivitycontrol = infectivitycontrol,
                     endtimes = endtimes,
                     seeding = seeding, ...)
  
  # get daily incidence (new infections)
  cumincidencecurves <- result[, paste0("cumI", 1:9)] * popsize(...) + seeding/9
  incidencecurves <- rbind(rep(seeding/9/popsize(...), 9) ,
                           tail(cumincidencecurves, -1) - head(cumincidencecurves, -1))
  
  # symptomatic
  tobesymptomatic <- t(t(incidencecurves) * probI2S(...))
  incsymptomatic <- sapply(1:(1 + tail(endtimes, 1)), 
                                 function(x) colSums(tobesymptomatic[max(1, x - 100):x, , drop = FALSE] * 
                                                       rev(delayI2S()[1:min(101, x)])
                                                     )) %>% t()
  
  # hospitalised
  tobehospitalised <- t(t(incsymptomatic) * probS2H(...))
  inchospitalised <- sapply(1:(1 + tail(endtimes, 1)), 
                                  function(x) colSums(tobehospitalised[max(1, x - 100):x, , drop = FALSE] * 
                                                        delayS2H()[min(101, x):1, , drop = FALSE]
                                                      )) %>% t()
  
  # dead without hospital
  tobedeadwithouthospital <- t(t(incsymptomatic) * probS2D(...))
  incdeadwithouthospital <- sapply(1:(1 + tail(endtimes, 1)), 
                                   function(x) colSums(tobedeadwithouthospital[max(1, x - 100):x, , drop = FALSE] * 
                                                         rev(delayS2D()[1:min(101, x)]) 
                                   )) %>% t()
  

  # hospital to ICU
  tobeICU <- t(t(inchospitalised) * probH2IC(...))
  incICU <- sapply(1:(1 + tail(endtimes, 1)), 
                   function(x) colSums(tobeICU[max(1, x - 100):x, , drop = FALSE] * 
                                         rev(delayH2IC()[1:min(101, x)])
                   )) %>% t()
  
  # hospital to dead
  tobedeadfromhospital <- t(t(inchospitalised) * probH2D(...))
  incdeadfromhospital <- sapply(1:(1 + tail(endtimes, 1)), 
                                function(x) colSums(tobedeadfromhospital[max(1, x - 100):x, , drop = FALSE] * 
                                                      rev(delayH2D()[1:min(101, x)])
                                )) %>% t()
  
  # hospital to recovery
  tobecuredfromhospital <- inchospitalised - tobeICU - tobedeadfromhospital
  inccuredfromhospital <- sapply(1:(1 + tail(endtimes, 1)), 
                                 function(x) colSums(tobecuredfromhospital[max(1, x - 100):x, , drop = FALSE] * 
                                                       rev(delayH2C()[1:min(101, x)])
                                 )) %>% t()
  
  # ICU to dead
  tobedeadfromICU <- t(t(incICU) * probIC2D())
  incdeadfromICU <- sapply(1:(1 + tail(endtimes, 1)), 
                           function(x) colSums(tobedeadfromICU[max(1, x - 100):x, , drop = FALSE] * 
                                                 rev(delayIC2D()[1:min(101, x)])
                           )) %>% t()
  
  # ICU to rehospitalised
  toberehospitalised <- t(t(incICU) * probIC2H())
  increhospitalised <- sapply(1:(1 + tail(endtimes, 1)), 
                              function(x) colSums(toberehospitalised[max(1, x - 100):x, , drop = FALSE] * 
                                                    rev(delayIC2H()[1:min(101, x)])
                              )) %>% t()
  
  # ICU to recovery
  tobecuredfromICU <- incICU - tobedeadfromICU - toberehospitalised
  inccuredfromICU <- sapply(1:(1 + tail(endtimes, 1)), 
                            function(x) colSums(tobecuredfromICU[max(1, x - 100):x, , drop = FALSE] * 
                                                  rev(delayIC2C()[1:min(101, x)])
                            )) %>% t()
  
  # rehospitalised to dead
  tobedeadfromrehospitalised <- t(t(increhospitalised) * probH22D())
  incdeadfromrehospitalised <- sapply(1:(1 + tail(endtimes, 1)), 
                                      function(x) colSums(tobedeadfromrehospitalised[max(1, x - 100):x, , drop = FALSE] * 
                                                            rev(delayH22D()[1:min(101, x)])
                                      )) %>% t()
  
  # rehospitalised to recovery
  tobecuredfromrehospitalised <- increhospitalised - tobedeadfromrehospitalised
  inccuredfromrehospitalised <- sapply(1:(1 + tail(endtimes, 1)), 
                                      function(x) colSums(tobecuredfromrehospitalised[max(1, x - 100):x, , drop = FALSE] * 
                                                            rev(delayH22C()[1:min(101, x)])
                                      )) %>% t()
  
  # total mortality
  incmortality <- incdeadwithouthospital + incdeadfromhospital + incdeadfromICU + incdeadfromrehospitalised
  
  # hospital prevalence
  prevhospitalised <- apply(inchospitalised, 2, cumsum) +
    apply(increhospitalised, 2, cumsum) -
    apply(incICU, 2, cumsum) -
    apply(incdeadfromhospital, 2, cumsum) - 
    apply(inccuredfromhospital, 2, cumsum) -
    apply(incdeadfromrehospitalised, 2, cumsum) -
    apply(inccuredfromrehospitalised, 2, cumsum)

  # ICU prevalence
  prevICU <- apply(incICU, 2, cumsum) -
    apply(incdeadfromICU, 2, cumsum) -
    apply(inccuredfromICU, 2, cumsum) -
    apply(increhospitalised, 2, cumsum) 
  
  # seroprevalence
  previmmune <- t(t(
    rbind(matrix(0, nrow = delayI2R(), ncol = 9),
          head(cumincidencecurves, -delayI2R()))) / (agedist(...) * popsize(...)))
  

  # combine and add names
  toreturn <- as_tibble(incsymptomatic) %>%
    rename_all(~rownames(contactmatrix())) %>%
    mutate(time = 0:tail(endtimes, 1)) %>%
    pivot_longer(1:9, names_to = "ageclass", values_to = "incsymp") %>%
    full_join(
      as_tibble(inchospitalised) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "inchosp"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(incICU) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "incICU"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(incmortality) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "incmort"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(incdeadfromhospital) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "incmorthosp"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(incdeadfromICU + incdeadfromrehospitalised) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "incmortIC"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(previmmune) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "immune"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(prevhospitalised) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "prevhosp"),
      by = c("time", "ageclass")
    ) %>%
    full_join(
      as_tibble(prevICU) %>%
        rename_all(~rownames(contactmatrix())) %>%
        mutate(time = 0:tail(endtimes, 1)) %>%
        pivot_longer(1:9, names_to = "ageclass", values_to = "prevICU"),
      by = c("time", "ageclass")
    ) 
  
  return(toreturn)
}