library(tidyverse)

options(scipen = 999)

# http://epirecip.es/epicookbook/chapters/sir/r_desolve

# Load deSolve library
library(deSolve)

################################################################################

url <- "https://api.covid19data.dk/ssi_hospitalized"

hosp_raw <- jsonlite::fromJSON(url)

hosp <- hosp_raw %>% 
  mutate(timestamp = timestamp %>% lubridate::ymd_hm(),
         date = timestamp %>% as.Date())

summed_data <- hosp %>%
  group_by(date) %>%
  summarise_at(c("hospitalized", "critical", "respirator"), sum, na.rm  = TRUE)


## Time-Dependent R_0 ##########################################################

# Define the model function
sir_ode <- function(times, init, parms){
  
  with(as.list(c(parms, init)), {
    
    # ODEs
    dSdt = -beta(times) * S * I / N
    dIdt = beta(times) * S * I / N - gamma * I
    dRdt = gamma * I
    
    list(c(dSdt, dIdt, dRdt))
  })
}

L = 31
N = 5.8 * 10 ^ 6
D = 7 # infections lasts x days
gamma = 1.0 / D

R_0 <- function(t){ifelse(t < L, .8, .9)}
beta <- function(t){R_0(t) * gamma}

# initial conditions: one exposed
S0 = N - 535; I0 = 535 ; R0 = 0

# Setup the model inputs
parms <- c(N = N, beta = beta, gamma = gamma)

init <- c(S = S0, I = I0, R = R0) 

times <- 1:100

# Run the model
sir_out <- ode(init, times, sir_ode, parms)

# Plot model
sir_out_df <- sir_out %>% as.data.frame() %>% as_tibble()

sir_out_df$time <- (sir_out_df$time - 1) + as.Date("2020-04-01")

pd <- sir_out_df %>% filter(time > "2020-04-15", time < "2020-05-15")

hosp_pd <- summed_data %>% filter(date >= "2020-04-01", date <= "2020-04-15")
hosp_pd_pred <- summed_data %>% filter(date > "2020-04-15")

pp <- ggplot() +
  
  geom_line(data = pd, aes(x = time, y = I), linetype = "dashed", color = "blue") +
  geom_line(data = pd, aes(x = time, y = I * .25), linetype = "dashed", color = "red") +
  
  geom_point(data = hosp_pd, aes(date, hospitalized), color = "blue") +
  geom_point(data = hosp_pd, aes(date, critical), color = "red") + 
  
  geom_point(data = hosp_pd_pred, aes(date, hospitalized), color = "blue", shape = 1) +
  geom_point(data = hosp_pd_pred, aes(date, critical), color = "red", shape = 1) + 
  
  theme_minimal() +
  labs(y = "", x = "") + 
  scale_x_date(date_breaks = "1 week", date_labels = "%d/%m") +
  
  annotate(geom = "text", x = as.Date("2020-04-15") + .5, y = 362, label = "362 indlagte i alt", color = "blue", hjust = 0) + 
  annotate(geom = "text", x = as.Date("2020-04-15"), y = 89, label = "89 indlagte på intensiv", color = "red", vjust = -1, hjust = .05) +
  
  annotate(geom = "text", x = as.Date("2020-05-01	"), y = 227, label = "227", color = "blue", vjust = 2) + 
  annotate(geom = "text", x = as.Date("2020-05-01	"), y = 57 , label = "57", color = "red", vjust = -1) +
  
  annotate(geom = "text", x = as.Date("2020-05-14"), y = 185, label = "185", color = "blue", vjust = 2) + 
  annotate(geom = "text", x = as.Date("2020-05-14"), y = 46, label = "46", color = "red", vjust = -1) + 
  
  theme(panel.grid = element_blank())
  
pp