library(dplyr)
library(cowplot)
library(deSolve)
set.seed(1234)
SIRsim <- function(beta,gamma){
mat <- matrix(0,12,12)
mat[sample(2:11,1),sample(2:11,1)] <- 1
out <- vector("list",100)
out[[1]] <- mat[2:11,2:11]
for(t in 2:100){
  for(i in 2:11){
    for(j in 2:11){
      #I ->R
      if(mat[i,j]==1){mat[i,j] <- rbinom(1,1,gamma)+1}
      #S -> I
      if(mat[i-1,j]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
      if(mat[i+1,j]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
      if(mat[i,j-1]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
      if(mat[i,j+1]==1 & mat[i,j]==0){mat[i,j] <- rbinom(1,1,beta)}
      }
  }
  out[[t]] <- mat[2:11,2:11]
}
SIR <-data.frame(time=1:100,
                 S =sapply(out,function(x)sum(x==0)),
                 I =sapply(out,function(x)sum(x==1)),
                 R =sapply(out,function(x)sum(x==2))) %>%
  tidyr::gather(state,populations,-time) %>%
  mutate(state=factor(state,levels=c("S","I","R")))
return(list(matlist=out,SIR=SIR))
}

res_sim <- SIRsim(0.1,0.01)

p2 <-ggplot(res_sim$SIR,aes(x=time,y=populations,colour=state))+
  geom_line(size=1)

ggsave("SIRsim.png",p2)
# dat <-res_sim$SIR %>%
#   tidyr::spread(state,populations)

cols = c("white","grey30","grey70")

#i <-which.max(dat$S==0)
library(animation)
saveGIF({
for(i in 1:100){
tmp <-res_sim$matlist[[i]] %>%
  as.data.frame() %>%
  data.table::setnames(LETTERS[1:10]) %>%
  mutate(col=factor(1:10)) %>%
  tidyr::gather(row,state,-col)

if(all(tmp$state!=2)){
  tmp_col=  cols[1:2]
}else if(all(tmp$state!=0)){
  tmp_col = cols[2:3]
  }else{
  tmp_col = cols[1:3]
}

p1<-ggplot(tmp,aes(x=factor(col),y=factor(row),fill=factor(state)))+
  geom_tile()+
  scale_fill_manual(values = tmp_col)+
  theme(legend.position = "none")+
  xlab("")+ylab("")
print(p1)
}},interval=0.3)