library(igraph)
library(grid)
library(gridBase)
library(gridExtra)

############
###### Data Prep Diagram
# create bipartite graph
edges = c(1,6,
          1,7,
          2,8,
          3,7,
          3,8,
          3,10,
          4,7,
          4,8,
          5,6,
          5,9)
g = graph.bipartite(c(rep(T,5),rep(F,5)),edges)

par(mfrow=c(1,3))
# (A) data table
data = data.frame(get.edgelist(g))
data$X2 = toupper(letters[data$X2-5])
colnames(data) = c('Event Code (EC)','Identity Code (IC)')
plot.new()
vps <- baseViewports()
pushViewport(vps$figure)
vp1 <-plotViewport()
grid.table(data)

# (B) bipartite person-event graph
layout = matrix(c(rep(seq(4,0,-1),2),rep(0,5),rep(1,5)),ncol=2)
labels = c( 1:5, toupper(letters[1:5]) )
cols = c( rep('#1a9850',5), rep('#1f78b4',5) )
plot(g, layout=layout[,2:1], vertex.color=cols,edge.color='black',
     vertex.frame.color=NA,vertex.size=30,vertex.label=labels,
     vertex.label.color='white',vertex.label.family='sans')

# (C) unipartite person-person graph
g2 = bipartite.projection(g)$proj1
layout.g2 = layout.auto(g2)
plot(g2, vertex.color='#1f78b4', edge.color='black',layout=layout.g2,
     vertex.frame.color=NA,vertex.size=30,vertex.label=toupper(letters[1:5]),
     vertex.label.color='white',vertex.label.family='sans')

#############
#### Hawkes Process Diagram
par(mfrow=c(1,3))

vics = data.frame(IC = toupper(letters[1:5]), 
                  Victim = c(T,F,F,F,T),
                  Day = c(3,NA,NA,NA,5))
colnames(vics) = c('Identity Code (IC)','Victim','Infection Date')

plot.new()
vps <- baseViewports()
pushViewport(vps$figure)
vp1 <-plotViewport()
grid.table(vics)

cols = rep('#1f78b4',vcount(g2))
cols[vics$Victim==T] = '#e41a1c'

plot(g2, vertex.color=cols, edge.color='black',layout=layout.g2,
     vertex.frame.color=NA,vertex.size=30,vertex.label=toupper(letters[1:5]),
     vertex.label.color='white',vertex.label.family='sans')


rate = function(x){
  print(x)
  r = rep(.3,length(x))
  times = c(1,4)
  for (time in times){
    r[x>=time] = r[x>=time] + 1.5*exp((time-x[x>=time])/4)
  }
  return(r)
}
plot(rate,from=0,to=10,ylim=c(0,3),lwd=0,bty='l',
     yaxt='n',xlab='Date',ylab='Infection Rate')
mtext("Date", line = 3, side = 1)
n = 1000
xs = c(0,seq(0,10,length.out=n),10)
ys = rate(xs)
ys[c(1,n+2)] = 0
polygon(xs,ys,col='#1f78b4',border=NA)