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

############
###### Data Prep Diagram
# create bipartite graph
edges = c(1,6,
          1,7,
          1,10,
          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), mar=c(1,1,1,1))
# (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
par(mar=rep(1,4))
g2 = bipartite.projection(g)$proj1
layout.g2 = matrix(c(0,.33,-.75,-.33,0,-1,0,1,.75,-.33),ncol=2,byrow=T)
plot(g2, vertex.color='#1f78b4', edge.color='black',layout=layout.g2,rescale=F,
     vertex.frame.color=NA,vertex.size=30,vertex.label=toupper(letters[1:5]),
     vertex.label.color='white',vertex.label.family='sans')

##############################################################################
##############################################################################
#### Hawkes Process Diagram
layout(matrix(c(1,2,3,3,3,3), ncol=2, byrow = TRUE))

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

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

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

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


# C 
par(mar=c(5.1, 4.1, 4.1, 2.1))
rate = function(x){
  print(x)
  r = rep(base+0.3,length(x))
  for (time in times){
    r[x>=time] = r[x>=time] + 1.5*exp((time-x[x>=time])/2)
  }
  return(r)
}

tmin = 1
tmax = 6
n = 1000
plot(rate,from=tmin,to=tmax,ylim=c(0,15),lwd=0,bty='l',col='white',
     yaxt='n',xlab='Date',ylab='Infection Rate')
for(i in 1:vcount(g2)){
  base <<- 3*(5-i)
  times <<- vics[neighbors(g2,i),3]
  times <<- times[!is.na(times)]
  xs = c(tmin,seq(tmin,tmax,length.out=n),tmax)
  ys = rate(xs)
  ys[c(1,n+2)] = base
  polygon(xs,ys,col='#1f78b4',border=NA)
}
axis(2, at=seq(0.3,12.3,3), lab=toupper(letters[seq(5,1,-1)]), tick=F, las=T,lwd=0)