require(data.table)
require(xlsx)
require(readxl)
require(rjags)
require(snow)

# Importations of COVID-19 into African countries and risk of onward spread

# Correspondence to:
# Haoyang Sun (haoyang.sun@nus.edu.sg)
# Hannah Clapham (hannah.clapham@nus.edu.sg)

##### PART A: Estimating total number of imported cases #####

# Loaded OAG flight data "flights_17"
# Set Working Directory

deplist = c("Italy", "United Kingdom", "United States", "France", "Spain", "Germany", 
            "Turkey", "Switzerland", "Belgium", "Netherlands", "Austria", "Portugal")

africa = as.data.table(read_xlsx("COVID-19\\For Plot\\Africa-Asia.xlsx"))[
  Continent == "Africa", .(ARRCOUNTRY = Country)][order(ARRCOUNTRY)]
africa = CJ(ARRCOUNTRY = africa$ARRCOUNTRY, DEPCOUNTRY = deplist)

africa = flights_17[Month == 3, .(Pax = sum(Pax)), by = .(DEPCOUNTRY, ARRCOUNTRY)][
  africa, on = c(DEPCOUNTRY = "DEPCOUNTRY", ARRCOUNTRY = "ARRCOUNTRY")]
africa[is.na(Pax), Pax := 0]

trvban = as.data.table(read.xlsx("COVID-19\\Line list\\Africa Data - Updated on 09072020.xlsx",
                                 sheetName = "Travel restrictions", stringsAsFactors = F, encoding = "UTF-8"))[
                                   ARRCOUNTRY %in% africa$ARRCOUNTRY, .(DEPCOUNTRY, ARRCOUNTRY, trvbant = Start_Date)]
africa = trvban[africa, on = c(DEPCOUNTRY = "DEPCOUNTRY", ARRCOUNTRY = "ARRCOUNTRY")]
africa = trvban[DEPCOUNTRY == "All selected DEPCOUNTRY", .(ARRCOUNTRY, trvbant)][africa, on = c(ARRCOUNTRY = "ARRCOUNTRY")]
africa[is.na(trvbant), trvbant := i.trvbant]
africa[is.na(trvbant), trvbant := as.Date("2020-05-01")]
africa = africa[, .(ARRCOUNTRY, DEPCOUNTRY, Pax, trvbant, Date = trvbant + 9, sgARRCN = -1)]

# Loaded reported number & timing of imported cases in SG "importedtoSG"

for(i in 1:nrow(africa))
  africa$sgARRCN[i] = round(sum(importedtoSG[Country == africa$DEPCOUNTRY[i] & Date <= africa$Date[i]]$cases))

mhelp = unique(africa[, .(DEPCOUNTRY, sgARRCN)])[order(DEPCOUNTRY, sgARRCN)][, I := .I]
mhelp = flights_17[ARRCOUNTRY == "Singapore" & Month == 3, .(Pax = sum(Pax)), by = DEPCOUNTRY][
  mhelp, on = c(DEPCOUNTRY = "DEPCOUNTRY")]

myjagscript = 
"
model
{
  for(i in 1:N)
  {
  sgARRCN[i] ~ dpois(beta[i] * Pax[i])
  
  beta[i] ~ dunif(0, 1)
  }
}
"

mypoismcmc = function()
{
  mydatajags = list(sgARRCN = mhelp$sgARRCN, Pax = mhelp$Pax, N = nrow(mhelp))
  initjags = list(beta = rbeta(nrow(mhelp), 1, 10))   
  
  jagmod = jags.model(textConnection(myjagscript), inits = initjags, data = mydatajags, n.chains = 10, quiet = T)
  update(jagmod, n.iter = 2000, progress.bar = "none")
  mypost = coda.samples(jagmod, c("beta"), n.iter = 15000, thin = 30)
  
  return(mypost)
}

'
mypost = mypoismcmc()
summary(geweke.diag(mypost)[[1]]$z)
summary(gelman.diag(mypost)$psrf[, c("Point est.", "Upper C.I.")])
'

myoutmcmc = as.list(as.data.table(as.matrix(mypost)))

africa = mhelp[, .(DEPCOUNTRY, sgARRCN, I)][africa, on = c(DEPCOUNTRY = "DEPCOUNTRY", sgARRCN = "sgARRCN")]
africa = africa[, .(ARRCOUNTRY, DEPCOUNTRY, Pax, trvbant, Date, sgARRCN, I, I_africa = .I)]

africa_import = rbindlist(lapply(africa$I_africa, function(x)
{
  with(africa, 
       data.table(ARRCOUNTRY = ARRCOUNTRY[x], DEPCOUNTRY = DEPCOUNTRY[x],
                  ARRCN = rpois(length(myoutmcmc[[1]]), Pax[x] * myoutmcmc[[I[x]]]),
                  I_mcmc = seq(myoutmcmc[[1]])))
}))


##### PART B: Simulating onward spread #####

myserial = list(meanX = 4.7, sdX = 2.9)

myserial$meanlog = with(myserial, log((meanX^2) / (sqrt(sdX^2 + meanX^2))))
myserial$sdlog = with(myserial, sqrt(log(1 + (sdX / meanX)^2)))

internalban = list()
internalban[[1]] = as.data.table(read.xlsx(
  "COVID-19\\Line list\\Africa Data - Updated on 09072020.xlsx",
  sheetName = "Local restrictions", stringsAsFactors = F, encoding = "UTF-8"))[
    Type %in% 2:3, .(ARRCOUNTRY, Type, Start_Date = as.Date(Start_Date), End_Date = End_Date_formatted)]

internalban[[2]] = internalban[[1]][Type == 2, .(ARRCOUNTRY, Start_Date)][
  data.table(ARRCOUNTRY = sort(unique(africa$ARRCOUNTRY))), on = c(ARRCOUNTRY = "ARRCOUNTRY")]
internalban[[2]][is.na(Start_Date), Start_Date := as.Date("2020-12-31")]

internalban[[3]] = rbindlist(lapply(sort(unique(africa$ARRCOUNTRY)), function(x)
{
  tmp = internalban[[1]][Type == 3 & ARRCOUNTRY == x]
  
  if(nrow(tmp) == 0) return(data.table(ARRCOUNTRY = x, internalbant = as.Date("2020-12-31")))
  
  tmp[End_Date == "NA", End_Date := "30/06/2020"]
  tmp[, End_Date := as.Date(End_Date, format = "%d/%m/%Y")]
  
  if(nrow(tmp) == 1) return(data.table(ARRCOUNTRY = x, internalbant = seq(tmp$Start_Date[1], tmp$End_Date[1], by = "days")))
  
  return(data.table(ARRCOUNTRY = x, internalbant = 
                      do.call("c", sapply(1:nrow(tmp), function(y) seq(tmp$Start_Date[y], tmp$End_Date[y], by = "days")))))
}))



w = rep(-1, 14)

for(i in seq(w))
  w[i] = mean(sapply(seq(i-1, i, length = 1000), function(x)
    plnorm(x + 1, meanlog = myserial$meanlog, sdlog = myserial$sdlog) - 
      plnorm(x, meanlog = myserial$meanlog, sdlog = myserial$sdlog)))


my_onward_spread = function(myarr, i, exclude, rho2, rho3, dispersion) 
{
  t0s = rbindlist(lapply(deplist[!deplist %in% ifelse(rep(exclude, 2), c("United Kingdom", "United States"), "None")], 
                         function(mydep)
                         {
                           tmp = importedtoSG[Country == mydep & Date <= africa[ARRCOUNTRY == myarr & DEPCOUNTRY == mydep]$Date]
                           
                           data.table(ARRCOUNTRY = myarr,
                                      DEPCOUNTRY = mydep,
                                      Date = tmp$Date,
                                      ARRCN = c(rmultinom(1, africa_import[
                                        DEPCOUNTRY == mydep & ARRCOUNTRY == myarr & I_mcmc == i]$ARRCN, tmp$cases)))
                         }))
  
  
  
  t0s = internalban[[2]][t0s, on = c(ARRCOUNTRY = "ARRCOUNTRY")]
  
  if(myarr != "Congo DRC") t0s[Date - 9 >= Start_Date, ARRCN := as.integer(round(ARRCN * rho2))] else
    t0s[Date - 9 >= Start_Date & DEPCOUNTRY %in% c("Italy", "France", "Germany"), ARRCN := as.integer(round(ARRCN * rho2))]
  
  if(sum(t0s$ARRCN) == 0) 
    return(data.table(Date_inf = as.Date("2020-01-01"), ARRCN = 0, LCN = 0, TCN = 0, R0 = 2.0, I_mcmc = i, I = 1))
  
  
  t0s = t0s[, .(ARRCN = sum(ARRCN)), keyby = .(Date_inf = Date - 9)]
  
  onwarddt = data.table(Date_inf = min(t0s[ARRCN != 0]$Date_inf) + (-14:(as.Date("2020-06-30") - min(t0s[ARRCN != 0]$Date_inf))))
  onwarddt = t0s[onwarddt, on = c(Date_inf = "Date_inf")]
  onwarddt[is.na(ARRCN), ARRCN := 0]
  
  setkey(onwarddt, Date_inf)
  
  onwarddt[, ":="(LCN = 0, TCN = 0, R0 = 2.0, I_mcmc = i, I = .I)]
  onwarddt[Date_inf %in% internalban[[3]][ARRCOUNTRY == myarr]$internalbant, R0 := R0 * rho3]
  
  
  for(j in 15:max(onwarddt$I))
  {
    tmp = sum(sapply(1:14, function(x) (onwarddt$ARRCN[j - x] + onwarddt$LCN[j - x]) * w[x]))
    
    if(tmp != 0) onwarddt$LCN[j] = rnbinom(1, size = tmp * dispersion, prob = 1 / (1 + onwarddt$R0[j] / dispersion))
    
    onwarddt$TCN[j] = sum(onwarddt$LCN[1:j]) + sum(onwarddt$ARRCN[1:j])
    
    if(onwarddt$TCN[j] >= 10000) return(onwarddt[I <= j])
  }
  
  return(onwarddt[I <= j])
}


arrlist = sort(unique(africa$ARRCOUNTRY))

bp_in = CJ(rho3 = c(0.75, 0.5), dispersion = c(0.58, 0.1), myarr = arrlist, i = seq(5, 5000, by = 5))
bp_in[, I_bp_in := .I]

'
mycl = makeCluster(5, type = "SOCK")
clusterEvalQ(mycl, {library(data.table)}) 
clusterExport(mycl, c("deplist", "importedtoSG", "africa_import", "africa",
"internalban", "w", "my_onward_spread", "bp_in"))

Sys.time()
bpout = clusterApply(mycl, 1:nrow(bp_in), function(x) 
my_onward_spread(bp_in$myarr[x], bp_in$i[x], 1, 0.00, bp_in$rho3[x], bp_in$dispersion[x]))
Sys.time()

stopCluster(mycl)
rm(mycl)
'

bpout_guide = unique(bp_in[, .(rho3, dispersion, myarr)])
bpout_guide[, I := .I]


bpout_process = function(ii, threshold_tcn = 1e4)
{
  tmp = bp_in[rho3 == bpout_guide$rho3[ii] & dispersion == bpout_guide$dispersion[ii] & myarr == bpout_guide$myarr[ii]]
  
  tmp = rbindlist(lapply(tmp$I_bp_in, function(x) bpout[[x]]))
  tmp = tmp[, .(rho3 = bpout_guide$rho3[ii], dispersion = bpout_guide$dispersion[ii], ARRCOUNTRY = bpout_guide$myarr[ii],
                Date_inf, TCN, I_mcmc, I)]
  
  rbindlist(lapply(3:6, function(month_cutoff)
  {
    tmp2 = tmp[month(Date_inf) <= month_cutoff]
    tmp2[, isselected := (I == max(I)), by = I_mcmc]
    tmp2 = tmp2[isselected == T]
    
    if(nrow(tmp2) != 1000) print(paste0("Error! ii = ", ii))
    
    tmp2[TCN < threshold_tcn, Date_inf := as.Date("2020-12-31")] # To facilitate sorting only (Actual value is ">06-30")
    tmp2 = tmp2[order(Date_inf)]
    
    tmp2 = tmp2[, .(month_cutoff = month_cutoff,
                    risk = mean(TCN >= threshold_tcn),
                    Date_inf025 = Date_inf[round(1000 * 0.025)],
                    Date_inf250 = Date_inf[round(1000 * 0.250)],
                    Date_inf500 = Date_inf[round(1000 * 0.500)],
                    Date_inf750 = Date_inf[round(1000 * 0.750)],
                    Date_inf975 = Date_inf[round(1000 * 0.975)]), 
                by = .(rho3, dispersion, ARRCOUNTRY)]
    return(tmp2)
  }))
}

bpout2 = rbindlist(lapply(1:nrow(bpout_guide), bpout_process))