#!/usr/bin/env Rscript

# Leila Gúl
# Dénes Türei (turei.denes@gmail.com)
# 2017

require(igraph)
require(qgraph)

# params
N = 100
simlength <- 60
p.t <- 0.3
plot.spread <- TRUE
make.movie <- TRUE
immune.t <- 6
infected.t <- 3
healthy.col <- 'aquamarine'
infected.col <- 'coral'
immune.col <- 'bisque'

# methods
generate.network.B <- function(N,links.per.step){
  L <- matrix(nrow=0,ncol=2)
  deg <- integer(N) 
  for (i in 2:N) { 
    n.new <- min(links.per.step,i-1) 
    n.new
    linkto <- sample(i-1,n.new,prob=deg[1:(i-1)]+1) 
    linkto
    newlinks <- cbind(rep(i,n.new),linkto)
    newlinks
    L <- rbind(L,newlinks)
    L
    deg[i] = deg[i] + n.new
    deg[linkto] = deg[linkto]+1
    deg
  }
  colnames(L) <- NULL
  L
}

plot.time <- function(g, lo, i, name = 'simulation-'){
    png(filename = paste(name, sprintf('%04d', i), '.png', sep = ''),
        pointsize = 12,
        width = 900, height = 600, type = 'cairo-png')
    plot(g, layout = lo, main = sprintf('Time = %04d', i),
         vertex.frame.color = NA, vertex.size = 11)
    dev.off()
}

# initialization
links <- generate.network.B(N,2)
g <- graph.edgelist(links,directed=FALSE)
V(g)$infected <- rep(0, N)
patientzero <- sample(N,1)
V(g)$infected[patientzero] <- infected.t
V(g)$immune   <- rep(0, N)
V(g)$immune[patientzero] <- immune.t
V(g)$color <- healthy.col
V(g)$color[patientzero] <- infected.col

el <- get.edgelist(g, names = FALSE)
lo <- qgraph.layout.fruchtermanreingold(el, vcount = vcount(g),
                                        area = vcount(g)^2.3,
                                        repulse.rad = vcount(g)^2.1,
                                        niter = 3000)

# plotting time #0
if(plot.spread){
    plot.time(g, lo, 0)
}

# run simulation
for(i in 1:simlength){
    V(g)$infected <- sapply(V(g)$infected - 1, function(x) max(x, 0))
    V(g)$immune   <- sapply(V(g)$immune   - 1, function(x) max(x, 0))
    discordant.links <- which(xor(V(g)$infected[links[,1]] & !V(g)$immune[links[,2]],
                                  V(g)$infected[links[,2]] & !V(g)$immune[links[,1]]))
    transmit <- rbinom(length(discordant.links), 1, p.t)
    transmitter.links <- discordant.links[transmit==1]
    nodes.of.transmitter.links <- setdiff(unique(as.vector(links[transmitter.links, 1:2])),
                                          which(V(g)$infected > 0))
    V(g)$infected[nodes.of.transmitter.links] <- infected.t
    V(g)$immune[nodes.of.transmitter.links]   <- immune.t
    if (plot.spread) {
        V(g)$color <- healthy.col
        V(g)$color[V(g)$infected > 0] <- infected.col
        V(g)$color[V(g)$infected == 0 & V(g)$immune > 0] <- immune.col
        plot.time(g, lo, i)
    }
}

# convert to movie
# you need mplayer (mencoder) for this
if(make.movie){
    system('mencoder mf://simulation-*.png\\
        -mf w=900:h=600:fps=1:type=png\\
        -ovc lavc -lavcopts vcodec=mpeg4:mbd=2:trell\\
        -oac copy -o simulation-1.avi')
}