## Set Working Directory to file directory - RStudio approach
setwd(dirname(rstudioapi::getActiveDocumentContext()$path))

weighted <- 'True'
#conda_install(envname = "r-reticulate", packages="networkx")
#conda_install(envname = "r-reticulate", packages="matplotlib")
#conda_install(envname = "r-reticulate", packages="pydot")
#conda_install(envname = "r-reticulate", packages="pygraphviz")

seirds.f <- function(t, y, k) {
  S <- y[1]
  E <- y[2]
  I <- y[3]
  R <- y[4]
  D <- y[5]

  beta = k[1]
  delta = k[2]
  gamma_r = k[3]
  gamma_d = k[4]
  mu = k[5]
  epsilon = k[6]
  waning = k[7]

  # Saying infec rate of S in contact with E same as contact with I
  dS <- epsilon - (beta*E + beta*I)*S + waning*R - gamma_d*S
  dE <- (beta*E + beta*I)*S - (delta+gamma_d)*E
  dI <- delta*E - (1+gamma_d)*I
  dR <- (1-mu)*I - (waning+gamma_d)*R
  dD <- mu*I
  
  return(as.matrix(c(dS, dE, dI, dR, dD)))
}

library(reticulate)
source_python('prep_model.py')
model_data <- prep_seirds(weighted)

S <- unlist(model_data)[1]
E <- unlist(model_data)[2]
I <- unlist(model_data)[3]
R <- unlist(model_data)[4]
D <- unlist(model_data)[5]

beta <- unlist(model_data)[6]
delta <- unlist(model_data)[7]
gamma_r <- unlist(model_data)[8]
gamma_d <- unlist(model_data)[9]
mu <- unlist(model_data)[10]
epsilon <- unlist(model_data)[11]
omega <- unlist(model_data)[12]

# Obtained from prep_model.py
seirds.params <- c(beta,     # beta
                   delta,    # delta
                   gamma_r,  # gamma_r
                   gamma_d,  # gamma_d
                   mu,       # mu
                   epsilon,  # epsilon
                   omega)    # waning

S <- S/length(color_d)
E <- E/length(color_d)
I <- (I_R+I_D)/length(color_d)
R <- R/length(color_d)
D <- D/length(color_d)

tmin <- 0
tmax <- 5

library(pracma)
seirds.ode.sol <- ode45(f = function(t,y){seirds.f(t,y,k=seirds.params)},
                     y = c(S, E, I, R, D),
                     t0 = tmin, tfinal = tmax)

# plot
library(reshape2)
library(ggplot2)
plot.seirds <- function(sol, method){
  # - pred/prey columns melted to 1 column called "value"
  sol.melted <- melt(data.frame(time=sol$t,
                                S=sol$y[,1],
                                E=sol$y[,2],
                                I=sol$y[,3],
                                R=sol$y[,4],
                                D=sol$y[,5]),
                     id = "time") # melt based on time
  # New column created with variable names, called “variable”
  colnames(sol.melted)[2] <- "Group" # used in legend
  g <- ggplot(data = sol.melted,
              aes(x = time, y = value, color = Group))
  g <- g + geom_point(size=2)
  g <- g + xlab("time") + ylab("Population")
  g <- g + ggtitle(paste("SEIRDS Model Using", method))
  show(g)
}

plot.seirds(seirds.ode.sol, "ODE45")
ggsave("SEIRDS.pdf")
# Sanity check: Make sure sums to ~1.0
sum(tail(seirds.ode.sol$y,1))