### fimd9.R
### R code from Van Buuren, S. (2012).
###		Flexible Imputation of Missing Data.
###		CRC/Chapman & Hall, Boca Raton, FL.
### (c) 2012 Stef van Buuren, www.multiple-imputation.com
### Version 1, 21mar2012
### Version 2, 4nov2015 tested with mice 2.23
### Tested with Mac OS X 10.7.3, R2.14-2, mice 2.12

if (packageVersion("mice")<'2.12') stop("This code requires mice 2.12.")

### Section 9.2 SE Fireworks Disaster Study

### Section 9.2.1 Intention to treat

library("mice")
library("lattice")

### Table 9.1

yvars <- c("yc1","yc2","yc3", "yp1","yp2","yp3")
x <- cbind(fdd[,c("id","trt","pp",yvars)])
x

md.pattern(fdd[fdd$pp=="Y",yvars])
md.pattern(fdd[fdd$pp=="N",yvars])

### Section 9.2.2 Imputation model

pred <- fdd.pred
vars <-c("ypa1", "ypb1", "ypc1",
         "ypa2", "ypb2", "ypc2",
         "ypa3", "ypb3", "ypc3")
pred[vars[1:3],vars]

dry <- mice(fdd, maxit=0)
method <- dry$method
method["yc1"] <- "~I(yca1 + ycb1 + ycc1)"
method["yc2"] <- "~I(yca2 + ycb2 + ycc2)"
method["yc3"] <- "~I(yca3 + ycb3 + ycc3)"
method["yp1"] <- "~I(ypa1 + ypb1 + ypc1)"
method["yp2"] <- "~I(ypa2 + ypb2 + ypc2)"
method["yp3"] <- "~I(ypa3 + ypb3 + ypc3)"
imp <- mice(fdd,pred=pred,meth=method,maxit=20,seed=54434)

### Section 9.2.3 Inspecting imputations

lowi <- complete(imp, "long", inc=TRUE)
lowi <- data.frame(lowi,cbcl2=NA, cbin2=NA,cbex2=NA)
lolo <- reshape(lowi, idvar = 'id',
                varying = 11:ncol(lowi),
                direction = "long",
                new.row.names = 1:(nrow(lowi)*3),
                sep="")
lolo <- lolo[order(lolo$.imp, lolo$id, lolo$time),]
row.names(lolo) <- 1:nrow(lolo)

### Figure 9.1

iv <- is.na(lolo[lolo$.imp==0,]$yp)
ivn <- ifelse(iv,1,0)
col12  <- c("grey80","grey80",
            mdc(2),mdc(1),
            mdc(2),"transparent",
            mdc(2),"transparent",
            mdc(2),"transparent",
            mdc(2),"transparent")
ic <- unique(lolo$id[iv])
ss <- lolo$id %in% ic
grp <- 2*as.integer(lolo$.imp) - ivn
loloss <- data.frame(lolo, grp=grp)
trellis.par.set(strip.background=list(col="grey95"))
tp1 <- xyplot(yp~time|factor(id), data=loloss, type="l",
      groups=factor(.imp), col="grey80", subset=ss,
       ylab="UCLA-RI Parent Score", pch=19, cex=1,
       xlab="Time", xlim=c("T1","T2","T3"),
             as.table=TRUE)
tp2 <- xyplot(yp~time|factor(id), data=loloss, type="p",
       groups=grp, col=col12, subset=ss,
       ylab="UCLA-RI Parent Score",
       pch=19,
       cex=0.8,
       xlab="Time", xlim=c("T1","T2","T3"),
       as.table=TRUE)
print(tp1)
print(tp2, newpage=FALSE)

### Figure 9.2

means <- aggregate(lolo$yp, list(lolo$.imp!=0,lolo$trt,lolo$time), mean, na.rm=TRUE)
names(means) <- c(".imp","trt","time","yp")
levels(means$trt) <- c("EMDR","CBT")
tp <- xyplot(yp~time|trt, data=means, type="o", pch=19,
      groups=factor(.imp), col=c(mdc(4), mdc(6)),
       ylab="UCLA-RI Parent Score", lwd=2,
       xlab="Time", xlim=c("T1","T2","T3"))
print(tp)

### Figure 9.3

means <- aggregate(lolo$yc, list(lolo$.imp!=0,lolo$trt,lolo$time), mean, na.rm=TRUE)
names(means) <- c(".imp","trt","time","yc")
levels(means$trt) <- c("EMDR","CBT")
tp <- xyplot(yc~time|trt, data=means, type="o", pch=19,
      groups=factor(.imp), col=c(mdc(4),mdc(6)),
       ylab="UCLA-RI Child Score", lwd=2,
       xlab="Time", xlim=c("T1","T2","T3"))
print(tp)

### Section 9.3 Time raster imputation

library("mice")
library("splines")
library("lme4")

data <- tbc
md.pattern(data)

### remove those with nocc 1 or 2 or 3
data <- data[data$nocc>=3,]

### Section 9.3.3 Broken stick model

### specify break ages
brk <- c(0, 8/365, 1/3, 1, 2, 6, 10, 18, 29)
k <- length(brk)

### calculate B-spline
X <- bs(data$age,
        knots = brk,
        B = c(brk[1],brk[k]+0.0001),
        degree = 1)
X <- X[,-(k+1)]
dimnames(X)[[2]] <- paste("x",1:ncol(X),sep="")
data <- cbind(data,X)
round(head(X,3),2)

### fit broken stick model
fit <- lmer(wgt.z~0+x1+x2+x3+x4+x5+x6+x7+x8+x9+
                 (0+x1+x2+x3+x4+x5+x6+x7+x8+x9|id),
                 data=data)

### calculate size and increment per person
tsiz <- t(ranef(fit)$id) + fixef(fit)
tinc <- diff(tsiz)

### inspect size and increment estimates
round(head(t(tsiz)),2)
round(head(t(tinc)),2)

### broken stick for hgt and BMI
fit.wgt <- fit
fit.hgt <- lmer(hgt.z~0+x1+x2+x3+x4+x5+x6+x7+x8+x9+
                 (0+x1+x2+x3+x4+x5+x6+x7+x8+x9|id),
                 data=data)
fit.bmi <- lmer(bmi.z~0+x1+x2+x3+x4+x5+x6+x7+x8+x9+
                 (0+x1+x2+x3+x4+x5+x6+x7+x8+x9|id),
                 data=data)

### define time warp

warp.setup <- data.frame(age = brk, age2 = seq(0, 29, length.out=k))
warp.model <- lm(age2 ~ bs(age, knots=brk[c(-1,-k)], degree=1) - 1,
                 data=warp.setup, x=T, y=T)
warped.knots <- warp.model$y
maxage <- max(warped.knots)
age2  <- predict(warp.model,newdata=data)
data  <- cbind(data, age2=age2)
rm(age2)

### Figure: time warping function (not in book)

par(fin=c(3,3),omi=c(0,0,0,0),mai=c(1,1,1,1),mfrow=c(1,1),lwd=0.5)
eqscplot(x=brk, y=warp.model$y,
         ylab="Warped age",
         xlab="Age (years)",
         axes=F,type="o",lwd=1,pch=16)
axis(1, at=seq(0,32,4), cex=0.5, pos=c(-0.5,0))
axis(2, at=warp.setup$age2,
     labels=c("0y","8d","4m","1y","2y","6y","10y","18y","29y"),
     line=0.5,cex=0.5)
lines(x=c(0,29), y=c(0,29), lwd=1)

# Prepare for figure 9.5 (weight, height and BMI profiles for six cases)

wsiz.bs <- extractBS(fit.wgt)
hsiz.bs <- extractBS(fit.hgt)
bsiz.bs <- extractBS(fit.bmi)
id <- unique(fit.wgt@flist$id)

data2 <- appendbreak(data, brk, id=id, warp.model=warp.model, typ="pred")
data2[data2$typ=="pred",c("wgt.z","hgt.z","bmi.z")] <- round(cbind(wsiz.bs, hsiz.bs, bsiz.bs),3)

six <- c(1259,7019,2447,7460,8046,7646)
set.seed(23221)
sample <- unique(c(six,sample(unique(tbc$id),300)))

idx <- data2$id %in% six
pd  <- data2[idx,]
pd$id <- as.factor(pd$id)
pd <- cbind(pd,rowname=row.names(pd))
pdr <- reshape(pd,
               varying = c("wgt.z","hgt.z","bmi.z"),
               idvar = "rowname",
               direction="long",
               v.names="y")
pdr$grp <- pdr$time
pdr$grp[pdr$typ=="pred"] <- pdr$grp[pdr$typ=="pred"] + 3

### Figure 9.5 Only weight

tbcstick <- xyplot(y~age2|id, data=pdr, xlim=c(-1,maxage+1), ylim=c(-4,4),
       as.table = TRUE,
       scales = list(x=list(at=warped.knots,
              labels=c("0","8d","4m","1y","2y","6y","10y","18y","29y"))),
                   xlab = "Age", ylab = "Weight SDS",
                   layout = c(2,3),
                   subset = (grp==1 | grp==4),
                   groups = grp,
                   pch=19,
                   type=c(rep("p",3),rep("l",3)),
                   distribute.type = TRUE,
                   lwd=2,
                   col=c(rep(mdc(1),3),rep(mdc(5),3)),
                   panel = function(...) {
         panel.abline(v=warped.knots, lty=2, col="grey80")
         panel.abline(h=c(-2,0,2), lty=1, col="grey80")
         panel.xyplot(...)
       },
       strip = strip.custom(bg="grey95"))
print(tbcstick)

### Same plot, but then with height and BMI added (not in book)

tbcstick <- xyplot(y~age2|id, data=pdr, xlim=c(-1,maxage+1), ylim=c(-4,4),
       as.table = TRUE,
       scales = list(x=list(at=warped.knots,
              labels=c("0","8d","4m","1y","2y","6y","10y","18y","29y"))),
       xlab = "Age", ylab = "Standard Deviation Score (SDS)",
       layout = c(2,3),
       groups = grp,
       pch=20,
       type=c(rep("p",3),rep("l",3)), distribute.type = TRUE,
       lwd=1.5,
       col=c(mdc(1), hcl(80,100,40,0.7),hcl(160,100,40,0.7)),
       panel = function(...) {
         panel.abline(v=warped.knots, lty=2, col="grey80")
         panel.abline(h=c(-2,0,2), lty=1, col="grey80")
         panel.xyplot(...)
       },
       strip = strip.custom(bg="grey95"),
       key = list(
         text   = list(c("Weight","Height","BMI"), cex=0.8),
         lines  = list(lty=1, lwd=2, type="l"),
         col  = c(mdc(1), hcl(80,100,40,0.7),hcl(160,100,40,0.7)),
         columns = 3,
         space  = "bottom",
         border = FALSE)
       )
print(tbcstick)

### Section 9.3.6 Imputation

### append break ages before imputation
id <- unique(data$id)
data2 <- appendbreak(data, brk, id=id, warp.model=warp.model, typ="sup")
table(data2$typ)

options(digits=2)
head(data2)

### specify imputation methods
Y <- c("hgt.z","wgt.z","bmi.z")
imp <- mice(data2, maxit=0)
meth <- imp$method
meth[1:length(meth)] <- ""
meth[Y] <- "2l.norm"

## Note: The following statements are outdated
##       and replaced by the intercept=FALSE option
##       in the call to mice()
## mice.impute.2l.norm.noint <- mice.impute.2l.norm
## meth[Y] <- "2l.norm.noint"

### specify predictor matrix
pred <- imp$pred
pred[1:nrow(pred),1:ncol(pred)] <- 0
pred[Y,"id"] <- (-2) # class variable
pred[Y,"sex"] <- 1 # fixed effect
pred[Y, paste("x",1:9,sep="")] <- 2
pred[Y[1],Y[2]] <- 2  # mutual random effect
pred[Y[2],Y[1]] <- 2  # mutual random effect
pred[Y[3],Y[1:2]] <- 2  # random effect

### specify visit sequence
vis <- 1:3
names(vis) <- c("hgt.z","wgt.z","bmi.z")

### ready to go - TIME CONSUMING (30 MINUTES)
imp.1 <- mice(data2, intercept = FALSE, meth=meth, pred=pred, m=5, maxit=10, seed=52711)
imp.2 <- mice(data2, intercept = FALSE, meth=meth, pred=pred, m=5, maxit=10, seed=88348)
imp.1745 <- ibind(imp.1, imp.2)  ## create m=10 imputations

### Figure 9.6 diagnostic plot
### Note: The solution is different from that in the book
###       The sample size in the book is much larger (n=2604 instead of n=306)
###       The larger sample regularizes many of the imputed sticks
###       The plot for non-typical subject 7646 reveals that
###       with n=306, fit to the data is low for this subject

cd <- complete(imp.1745,"long")
sup <- cd[cd$typ=="sup",]
data3 <- data.frame(.imp=0, data2)
data3 <- rbind(data3, sup[,-2])
idx <- data3$id %in% six
pd  <- data3[idx,]
pd$id <- as.factor(pd$id)
pd$grp <- pd$.imp
pd$grp[pd$grp==0] <- NA
pd$grp[pd$typ=="obs"] <- 11
pd$grp <- reorder(pd$grp, as.numeric(pd$grp))
tbcimp <- xyplot(wgt.z~age2|id, data=pd, xlim=c(-1,maxage+1), ylim=c(-4,4),
       as.table = TRUE,
       scales = list(x=list(at=warped.knots,
              labels=c("0","8d","4m","1y","2y","6y","10y","18y","29y"))),
       xlab = "Age", ylab = "Weight SDS",
       groups=grp, layout = c(2,3),
       pch=c(rep(20,10), 19),
       type=c(rep("l",10),"p"), lwd=c(rep(1,10),1),
       col=c(rep(mdc(5),10),mdc(1)), distribute.type=TRUE,
       panel = function(...) {
         panel.abline(v=warped.knots, lty=2, col="grey80")
         panel.abline(h=c(-2,0,2), lty=1, col="grey80")
         panel.xyplot(...)
       },
       strip=strip.custom(bg="grey95")
       )
print(tbcimp)

### Figure 9.7

cd <- complete(imp.1745,1)
idx <- (cd$id) %in% sample
cd <- cd[idx,]
shingle <- cut(cd$age,breaks=c(brk,29.01),right=FALSE, inc=TRUE,
               labels=c("0d-8d","8d-4m","4m-1y","1y-2y","2y-6y",
                 "6y-10y","10y-18y","18y-29y","29y"))
tbchw <- xyplot(wgt.z~hgt.z|shingle,data=cd,
       xlim=c(-4,4), ylim=c(-4,4), type=c("p","g"),
       group=(typ=="sup"),pch=c(1,20),
       col=c(mdc(1:2)),
       xlab="Height SDS", ylab="Weight SDS",pty="s",
       strip=strip.custom(bg="grey95"))
print(tbchw)

### Section 9.3.7 Complete data analysis

### post-processing of imputed data

imp <- imp.1745
cd <- complete(imp,"long")
sup <- cd[cd$typ=="sup",]

sup$age <- round(sup$age,2)
sup$hgt.z <- round(sup$hgt.z,2)
sup$wgt.z <- round(sup$wgt.z,2)
sup$bmi.z <- round(sup$bmi.z,2)

lowi <- reshape(sup, idvar=c('id','.imp'), timevar = 'age',
                v.names=c('hgt.z','wgt.z','bmi.z'),
                direction="wide",
                drop=c(".id","occ","first","typ","hgt","wgt","bmi",
                  paste("x",1:9,sep=""),"age2"))
hsiz <- lowi[,c(".imp","id","nocc","sex",
                "hgt.z.0","hgt.z.0.02","hgt.z.0.33",
                "hgt.z.1","hgt.z.2","hgt.z.6","hgt.z.10","hgt.z.18",
                "hgt.z.29")]
wsiz <- lowi[,c(".imp","id","nocc","sex",
                "wgt.z.0","wgt.z.0.02","wgt.z.0.33",
                "wgt.z.1","wgt.z.2","wgt.z.6","wgt.z.10","wgt.z.18",
                "wgt.z.29")]
bsiz <- lowi[,c(".imp","id","nocc","sex",
                "bmi.z.0","bmi.z.0.02","bmi.z.0.33",
                "bmi.z.1","bmi.z.2","bmi.z.6","bmi.z.10","bmi.z.18",
                "bmi.z.29")]
# merge outcome data
bsiz <- merge(bsiz, tbc.target,all.x=TRUE)

hinc <- cbind(hsiz[,1:4],t(diff(t(hsiz[,-1:-4]))))
winc <- cbind(wsiz[,1:4],t(diff(t(wsiz[,-1:-4]))))
binc <- cbind(bsiz[,1:4],t(diff(t(bsiz[,-1:-4]))))

# merge outcome data
binc <- merge(binc, tbc.target,all.x=TRUE)

bmi.z.0.02 <- by(binc, binc$.imp, function(x) lm(bmi.z.0.02~ao, data = x, na.action='na.omit'))
bmi.z.0.33 <- by(binc, binc$.imp, function(x) lm(bmi.z.0.33~ao, data = x, na.action='na.omit'))
bmi.z.1    <- by(binc, binc$.imp, function(x) lm(bmi.z.1   ~ao, data = x, na.action='na.omit'))
bmi.z.2    <- by(binc, binc$.imp, function(x) lm(bmi.z.2   ~ao, data = x, na.action='na.omit'))
bmi.z.6    <- by(binc, binc$.imp, function(x) lm(bmi.z.6   ~ao, data = x, na.action='na.omit'))
bmi.z.10   <- by(binc, binc$.imp, function(x) lm(bmi.z.10  ~ao, data = x, na.action='na.omit'))
bmi.z.18   <- by(binc, binc$.imp, function(x) lm(bmi.z.18  ~ao, data = x, na.action='na.omit'))
bmi.z.29   <- by(binc, binc$.imp, function(x) lm(bmi.z.29  ~ao, data = x, na.action='na.omit'))

tab.bmi.z.0.02 <- summary(pool(as.mira(bmi.z.0.02)))
tab.bmi.z.0.33 <- summary(pool(as.mira(bmi.z.0.33)))
tab.bmi.z.1    <- summary(pool(as.mira(bmi.z.1)))
tab.bmi.z.2    <- summary(pool(as.mira(bmi.z.2)))
tab.bmi.z.6    <- summary(pool(as.mira(bmi.z.6)))
tab.bmi.z.10   <- summary(pool(as.mira(bmi.z.10)))
tab.bmi.z.18   <- summary(pool(as.mira(bmi.z.18)))
tab.bmi.z.29   <- summary(pool(as.mira(bmi.z.29)))

(bmi.z.0.33 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.0.33+bmi.z.0.02, data = x, na.action='na.omit'))))))
(bmi.z.1 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.1+bmi.z.0.33, data = x, na.action='na.omit'))))))
(bmi.z.2 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.2+bmi.z.1, data = x, na.action='na.omit'))))))
(bmi.z.6 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.6+bmi.z.2, data = x, na.action='na.omit'))))))
(bmi.z.10 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.10+bmi.z.6, data = x, na.action='na.omit'))))))
(bmi.z.18 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.18+bmi.z.10, data = x, na.action='na.omit'))))))

(bmi.z.6 <- summary(pool(as.mira(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.6+bmi.z.2, data = x, na.action='na.omit'))))))
summary(by(bsiz, bsiz$.imp, function(x) lm(bmi.z.jv~bmi.z.6+bmi.z.2, data = x, na.action='na.omit'))[[2]])

# repeat the analysis for the broken stick estimates

bsiz.bs <- as.data.frame(t(t(ranef(fit.bmi)[[1]]) + fixef(fit.bmi)))
dimnames(bsiz.bs)[[2]] <- names(bsiz)[5:13]
binc.bs <- t(diff(t(bsiz.bs)))
bsiz.bs <- data.frame(id=bsiz[bsiz$.imp==1,"id"], bsiz.bs)
binc.bs <- data.frame(id=bsiz[bsiz$.imp==1,"id"], binc.bs)
bsiz.bs <- merge(bsiz.bs, tbc.target)
binc.bs <- merge(binc.bs, tbc.target)

bmi.z.0.02.bs <- lm(bmi.z.0.02~ao, data = binc.bs, na.action='na.omit')
bmi.z.0.33.bs <- lm(bmi.z.0.33~ao, data = binc.bs, na.action='na.omit')
bmi.z.1.bs    <- lm(bmi.z.1   ~ao, data = binc.bs, na.action='na.omit')
bmi.z.2.bs    <- lm(bmi.z.2   ~ao, data = binc.bs, na.action='na.omit')
bmi.z.6.bs    <- lm(bmi.z.6   ~ao, data = binc.bs, na.action='na.omit')
bmi.z.10.bs   <- lm(bmi.z.10  ~ao, data = binc.bs, na.action='na.omit')
bmi.z.18.bs   <- lm(bmi.z.18  ~ao, data = binc.bs, na.action='na.omit')

summary(bmi.z.0.02.bs)
summary(bmi.z.0.33.bs)
summary(bmi.z.1.bs)
summary(bmi.z.2.bs)
summary(bmi.z.6.bs)
summary(bmi.z.10.bs)
summary(bmi.z.18.bs)

### Materials for Table 9.3

a<-round(cor(bsiz[,-(1:4)],use="pair"),2)
b<-round(cor(bsiz.bs,use="pair"),2)

a2<-round(cor(bsiz[!is.na(bsiz$ao),-(1:4)],use="complete.obs"),2)
b2<-round(cor(bsiz.bs[!is.na(bsiz$ao),],use="complete.obs"),2)


 DTGSJA Example Program Results

 Number of infinite generalized singular values (K)
     1
 Number of finite generalized singular values (L)
     2
  Effective Numerical rank of (A**T B**T)**T (K+L)
     3

 Finite generalized singular values
     1.3151E+00  8.0185E-02

 Orthogonal matrix U
              1           2           3           4
 1  -1.3484E-01  5.2524E-01 -2.0924E-01  8.1373E-01
 2   6.7420E-01 -5.2213E-01 -3.8886E-01  3.4874E-01
 3   2.6968E-01  5.2757E-01 -6.5782E-01 -4.6499E-01
 4   6.7420E-01  4.1615E-01  6.1014E-01  1.5127E-15

 Orthogonal matrix V
              1           2
 1   3.5539E-01 -9.3472E-01
 2   9.3472E-01  3.5539E-01

 Orthogonal matrix Q
              1           2           3
 1  -8.3205E-01 -9.4633E-02 -5.4657E-01
 2   5.5470E-01 -1.4195E-01 -8.1985E-01
 3   0.0000E+00 -9.8534E-01  1.7060E-01

 Nonsingular upper triangular matrix R
              1           2           3
 1  -2.0569E+00 -9.0121E+00 -9.3705E+00
 2              -1.0882E+01 -7.2688E+00
 3                          -6.0405E+00

 Number of cycles of the Kogbetliantz method
     2


#' Calculate link position vectors
#'
#' Counts the number of times each link in a network occurs in each unique link position within the motifs
#' @param M A numeric matrix representing interactions between two groups of nodes. Each row corresponds to a node in one level
#' and each column corresponds to a node in the other level. Elements of M are positive numbers if nodes interact, and 0
#' otherwise. Formally, M is a biadjacency matrix. When nodes i and j interact, m_ij > 0; if they do not interact, m_ij = 0.
#'
#' @param six_node Logical; should positions in six node motifs be counted? Defaults to FALSE.
#' @param weights Logical; Should weights of the links be taken into account?
#' @param normalisation Which normalisation should be used: 'none','sum', 'position', 'sizeclass', 'sizeclass_plus1', 'sizeclass_NAzero', 'levelsize', 'levelsize_plus1', 'levelsize_NAzero','motif', 'motif_plus1' or 'motif_NAzero'?  Defaults to "none". (see details)
#' @details Counts the number of times each link in a network occurs in each of the 29 (if \code{six_node} = FALSE) or 106 (if \code{six_node} = TRUE) unique link positions within motifs (to quantify a link's structural role).
#' If \code{six_node} = FALSE, link positions in all motifs containing between 2 and 5 nodes are counted. If \code{six_node} = TRUE, link positions in all motifs containing between 2 and 6 nodes are counted. Analyses where \code{six_node} = FALSE are substantially faster
#' than when \code{six_node} = TRUE, especially for large networks. For large networks, counting six node motifs is also memory intensive. In some cases, R can crash if there is not enough memory.
#'
#' If interactions are weighted (non-zero matrix elements take values other than 1), these can be incorporated by setting \code{weights = TRUE}.
#' If \code{weights = TRUE}, the function will return the number of times each link occurs in each position,
#' multiplied by the weight of the link, following Mora et al. (2018).
#'
#' Links between nodes with more interactions will tend to appear in more positions. Normalisation helps control for this effect. bmotif include four types of normalisation:
#' \itemize{
#'  \item{\strong{"none"}: performs no normalisation and will return the raw position counts.}
#'  \item{\strong{"sum"}: divides the position measure for each link by the total number of times that link appears in any position (divides each element in a row by the row sum).}
#'  \item{\strong{"position"}: divides the position measure for each link by the total number of times any link occurs in that link position (divides each element in a column by the column sum). This gives a measure of how often a link occurs in a position relative to the other links in the network.}
#'  \item{\strong{Size class normalisation}
#'  \itemize{
#'  \item{\strong{"sizeclass"}: divides the position measure for each link by the total number of times that link appears in any position within the same motif size class (the number of nodes a motif contains).}
#'  \item{\strong{"sizeclass_plus1"}: same as 'sizeclass' but adds one to all position measure values. If a link does not occur in any motifs in a given size class, 'sizeclass' normalisation
#'  will return NAs. 'sizeclass_plus1' avoids this by adding one to all counts.}
#'  \item{\strong{"sizeclass_NAzero"}: same as 'sizeclass' but replaces all NA values with 0. If a link does not occur in any motifs in a given size class, 'sizeclass' normalisation
#'  will return NAs. 'sizeclass_NAzero' avoids this by replacing NAs with zero.}
#'  }
#'  }
#'  \item{\strong{Levelsize normalisation}
#'  \itemize{
#'  \item{\strong{"levelsize"}: divides the position measure for each link by the total number of times that link appears in any position within a motif with a given number of nodes in the top level and the bottom level.
#'  For example, the relative frequencies of all position measures in motifs with three nodes in the top level and two nodes in the bottom level will sum to one, as will the relative frequency of all position measures in motifs with 2 nodes in the top level and
#'  two nodes in the bottom level, and so on.}
#'  \item{\strong{"levelsize_plus1"}: same as 'levelsize' but adds one to all position measure values. If a link does not occur in any motifs with a given number of nodes in the top level and the bottom level, 'levelsize' normalisation
#'  will return NAs. 'levelsize_plus1' avoids this by adding one to all counts.}
#'  \item{\strong{"levelsize_NAzero"}: same as 'levelsize' but replaces all NA values with 0. If a link does not occur in any motifs with a given number of nodes in the top level and the bottom level, 'levelsize' normalisation
#'  will return NAs. 'levelsize_NAzero' avoids this by replacing NAs with zero.}
#'  }
#'  }
#'  \item{\strong{Motif normalisation}
#'  \itemize{
#'  \item{\strong{"motif"}: divides the position measure for each link by the total number of times that link appears in any position within the same motif.
#'  For example, the relative frequencies of all position measures in motif 5 will sum to one, as will the relative frequency of all position measures in motif 10, and so on.}
#'  \item{\strong{"motif_plus1"}: same as 'motif' but adds one to all position measure values. If a link does not occur in a particular motif, 'motif' normalisation
#'  will return NAs. 'motif_plus1' avoids this by adding one to all counts.}
#'  \item{\strong{"motif_NAzero"}: same as 'motif' but replaces all NA values with 0. If a link does not occur in a particular motif, 'levelsize' normalisation
#'  will return NAs. 'motif_NAzero' avoids this by replacing NAs with zero.}
#'  }
#'  }
#'  }
#'
#' If a matrix is provided without row or column names, default names will be assigned: the first row will be called called 'r1', the second row will be called 'r2' and so on. Similarly, the first column will be called 'c1', the second column will be called 'c2' and so on.
#'
#' @return  Returns a data frame with one column for each link position: 29 columns if \code{six_node} is FALSE, and 106 columns if \code{six_node} is TRUE.
#' Columns names are given as "lpx" where x is the ID of the position as described in the motif dictionary. \strong{To view the 'motif dictionary' showing
#' which link position a given ID corresponds to, enter \code{vignette("bmotif-dictionary")}.}
#'
#' Each row corresponds to one link in the network. Row names are gives as "x -- y", where x is the species in the first level (rows) and y is the species in the second level (columns).
#'
#' By default, the elements of the data frame will be the raw link position counts.
#' If \code{weight = TRUE}, link position counts will be multiplied by the link weight.
#' If \code{normalisation} is set to "sum", "sizeclass" or "position", the elements will be
#' normalised position counts as described above.
#'
#' @export
#' @references
#' Mora, B.B., Cirtwill, A.R. and Stouffer, D.B., 2018. pymfinder: a tool for the motif analysis of binary and quantitative complex networks. bioRxiv, 364703.
#'
#' Simmons, B. I., Sweering, M. J. M., Dicks, L. V., Sutherland, W. J. and Di Clemente, R. bmotif: a package for counting motifs in bipartite networks. bioRxiv. doi: 10.1101/302356
#' @examples
#' set.seed(123)
#' row <- 10
#' col <- 10
#'
#' # link positions in a binary network
#' m <- matrix(sample(0:1, row*col, replace=TRUE), row, col)
#' link_positions(M = m, six_node = TRUE, weights = FALSE, normalisation = "none")
#'
#' # link positions in a weighted network
#' m[m>0] <- stats::runif(sum(m), 0, 100)
#' link_positions(M = m, six_node = TRUE, weights = TRUE, normalisation = "none")

link_positions <- function(M, six_node = FALSE, weights, normalisation = "none") {
  # M is the adjacence matrix

  # check inputs
  if(inherits(M, "matrix") != TRUE){stop("'M' must be an object of class 'matrix'")} # make sure M is a matrix
  if(!all(apply(M, 1:2, is.numeric))){stop("Elements of 'M' must be numeric")} # make sure all elements of M are numbers
  if(!all(apply(M, 1:2, function(x) length(x) > 0))){stop("Elements of 'M' cannot have length zero")} # make sure no elements of M have 0 length e.g. numeric(0)
  if(!all(apply(M, 1:2, function(x) x >= 0))){stop("Elements of 'M' must be greater than or equal to zero")} # make sure all elements of M are >= zero
  if(all(apply(M, 1:2, function(x) x == 0))){stop("The matrix has no links (all elements are zero)")} # make sure M has some links
  if(inherits(six_node,"logical") != TRUE){stop("'six_node' must be of class 'logical' i.e. TRUE or FALSE")} # make sure six_node is logical i.e. TRUE or FALSE
  if(inherits(normalisation,"character") != TRUE){stop("'normalisation' must be of class 'character'")} # make sure 'normalisation' is a character
  if(!normalisation %in% c("none","sum","sizeclass", "sizeclass_plus1", "sizeclass_NAzero","levelsize", "levelsize_plus1", "levelsize_NAzero", "motif", "motif_plus1", "motif_NAzero","position")){
    stop("'normalisation' must equal 'none','sum', 'position', 'sizeclass', 'sizeclass_plus1', 'sizeclass_NAzero', 'levelsize', 'levelsize_plus1', 'levelsize_NAzero','motif', 'motif_plus1' or 'motif_NAzero'")
  }
  if(any(duplicated(rownames(M))) | any(duplicated(colnames(M)))){stop("Input matrix must not have duplicate row or column names")}
  if((normalisation == "sum" | normalisation == "sizeclass" | normalisation == "levelsize" | normalisation == "motif") & weights == TRUE){warning("Please note that with 'sum', 'sizeclass', 'levelsize' or 'motif' normalisation, the results won't change if weights are taken into account. This is because when weights = TRUE, each row of the output is multiplied by a fixed factor (the link's weight) and therefore the relative proportions are the same as if weights were not considered. Consider setting normalisation to 'position' or 'none'.")}

  W <- M # store copy of weighted matrix
  M[M > 0] <- 1 # ensure M is binary

  # compute auxiliary stuff
  NZ = nrow(M) # this is just Z in the latex file, but Z is also a matrix later
  NP = ncol(M) # this is just P in the latex file, but P is also a matrix later
  jZ = as.matrix(rep(1, NZ))
  jP = as.matrix(rep(1, NP))

  J <- matrix(rep(1, NZ * NP), nrow = NZ, ncol = NP) # Z x P matrix of ones
  JP <- matrix(rep(1, NP * NP), nrow = NP, ncol = NP) # P x P matrix of ones
  JZ <- matrix(rep(1, NZ * NZ), nrow = NZ, ncol = NZ) # Z x Z matrix of ones

  JT <- t(J) # transpose of J, P x Z matrix of ones
  MT <- t(M) # transpose of M
  N <- J - M # complement of M
  NT <- t(N) # transpose of the complement of M

  dZ = M %*% jP
  dP = MT %*% jZ

  V = matrix(rep(dZ, each = NP), nrow = NP, ncol = NZ, byrow = FALSE)
  U = matrix(rep(dP, times = NZ), nrow = NP, ncol = NZ, byrow = FALSE)

  # matrices like in the formulae
  Z <- M %*% MT
  Y <- M %*% NT
  X <- N %*% MT

  P <- MT %*% M
  Q <- MT %*% N
  R <- NT %*% M

  ##########################

  # now compute positions
  # possibly put this in a separate function later

  # create a data frame
  # each row corresponds to one edge in the graph
  # order: (1,1), (1,2), (1,3), ...(1,P), (2,1), ..., (2,P), ..., (Z,1), ..., (Z,P)
  # each column corresponds to an edge position, i.e. column p3 gives count for link position 3

  # note: the as.vector() command is not strictly needed
  # since R even converts the matrix automatically into a vector if we assign a matrix to a column like below

  k = 29 + (six_node) * (106 - 29) # k depends on whether 6 species are wanted or not

  lp <- matrix(rep (-1, k* NZ * NP), ncol = k)

  ######

  names <- rep("", NZ * NP)

  if (is.null(rownames(M))) {
    class1_names <- paste0("r", 1:NZ)
  } else {
    class1_names <- rownames(M)
  }
  if (is.null(colnames(M))) {
    class2_names <- paste0("c", 1:NP)
  } else {
    class2_names <- colnames(M)
  }

  for (i in 1:NZ) {
    for (j in 1:NP) {
      names[(i - 1)* NP + j] <- name_edge(class1_names[i], class2_names[j])
    }
  }
  rownames(lp) <- names

  if (six_node) {
    colnames(lp) <- paste('lp', 1:106, sep = '')
  }
  if (!six_node) {
    colnames(lp) <- paste('lp', 1:29, sep = '')
  }

  # want to compute only for present edges
  # so make a list of all the edges present in the network
  # column one is just a numbering
  # column 2 is the index of the row vertex, column 3 the index of the column vertex
  # for later use, if weights = TRUE, also store the weight of the link in column 4
  ct <- 0
  E <- sum(M) # number of edges

  if (six_node & !weights) {
    edges <- matrix(rep(0,3*E), nrow = E)
    for (i in 1:NZ) {
      for (j in 1:NP) {
        if (M[i,j] == 1) {
          ct <- ct + 1
          edges[ct,1] <- ct
          edges[ct,2:3] <- c(i, j)
        }
      }
    }
  }

  if (weights) {
    edges <- matrix(rep(0,4*E), nrow = E)
    for (i in 1:NZ) {
      for (j in 1:NP) {
        if (M[i,j] == 1) {
          ct <- ct + 1
          edges[ct,1] <- ct
          edges[ct,2:3] <- c(i, j)
          edges[ct, 4] <- W[i,j]
        }
      }
    }
  }

  # ------- PUT IN VALUES
  # up to 4 species

  lp[,1] <- as.vector(t(M))
  lp[,2] <- as.vector(V - JT)

  lp[,3] <- as.vector(U - JT)
  lp[,4] <- 0.5*as.vector((U - JT)*(U-2*JT))
  lp[,5] <- as.vector(NT %*% Z)
  lp[,6] <- as.vector(MT %*% Y)
  lp[,7] <- as.vector(MT %*% X)
  lp[,8] <- as.vector(MT %*% (Z - JZ) - V + JT)
  lp[,9] <- 0.5*as.vector((V - JT)*(V - 2*JT))

  # 5 species

  lp[,10] <- 1/6 * as.vector((U - JT)*(U - 2*JT) * (U - 3*JT))
  lp[,11] <- 0.5 * as.vector(t((M %*% (Q * (Q - JP)))))
  lp[,12] <- as.vector(t(N %*% (P * (R - JP))))
  lp[,13] <- 0.5 * as.vector(t(M %*% (R * (R - JP))))
  lp[,14] <- as.vector(t(N %*% (P * Q)))
  lp[,15] <- as.vector(t(M %*% (Q * R)))
  lp[,16] <- as.vector(t(M %*% (Q * (P - JP))))
  lp[,17] <- 0.5 * as.vector(t(N %*% (P * (P - JP))))
  lp[,18] <- as.vector(t(M %*% (R * (P - JP))))
  lp[,19] <- 0.5 * as.vector(t(M %*% ( (P - JP) * (P - 2*JP))) - (U - JT) * (U - 2*JT) )
  lp[,20] <- as.vector(NT %*% (Z * (X - JZ)))
  lp[,21] <- 0.5 * as.vector(MT %*% (Y * (Y - JZ)))
  lp[,22] <- 0.5 * as.vector(MT %*% (X * (X - JZ)))
  lp[,23] <- as.vector(NT %*% (Z * Y))
  lp[,24] <- as.vector(MT %*% (X * Y))
  lp[,25] <- 0.5 * as.vector(NT %*% (Z * (Z - JZ)))
  lp[,26] <- as.vector(MT %*% (Y * (Z - JZ)))
  lp[,27] <- as.vector(MT %*% (X * (Z - JZ)))
  lp[,28] <- 0.5 * as.vector(MT %*% ( (Z - JZ) * (Z - 2*JZ) ) - (V - JT) * (V - 2*JT))
  lp[,29] <- 1/6 * as.vector((V - JT) * (V - 2*JT) * (V - 3*JT))

  # 6 species
  if (six_node) {
    lp[,30] <- 1/24 * as.vector((U - JT) * (U - 2*JT) * (U - 3*JT) * (U - 4*JT))
    lp[,31] <- 1/6 * as.vector(t(M %*% (Q * (Q - JP) * (Q - 2*JP))))
    lp[,32] <- 1/2 * as.vector(t(N %*% ((R - JP) * (R - 2*JP) * P)))
    lp[,33] <- 1/6 * as.vector(t(M %*% (R * (R - JP) * (R - 2*JP))))
    lp[,34] <- 1/2 * as.vector(t(N %*% (Q * (Q - JP) * P)))
    lp[,35] <- 1/2 * as.vector(t(M %*% (Q * (Q - JP) * R)))
    lp[,36] <- as.vector(t(N %*% (Q * P * (R - JP))))
    lp[,37] <- 1/2 * as.vector(t(M %*% (R * (R - JP) * Q)))
    lp[,38] <- 1/2 * as.vector(t(M %*% (Q * (Q - JP) * (P - JP))))
    lp[,39] <- 1/2 * as.vector(t(N %*% (P * (P - JP) * (R - JP))))
    lp[,40] <- 1/2 * as.vector(t(M %*% (R * (R - JP) * (P - JP))))
    lp[,41] <- 1/2 * as.vector(t(N %*% (P * (P - JP) * Q)))
    lp[,42] <- as.vector(t(M %*% ((P - JP) * Q * R)))
    lp[,43] <- 1/2 * as.vector(t(M %*% ((P- JP) * (P - 2*JP) * Q)))
    lp[,44] <- 1/6 * as.vector(t(N %*% (P* (P - JP) * (P - 2*JP))))
    lp[,45] <- 1/2 * as.vector(t(M %*% ((P- JP) * (P - 2*JP) * R)))
    lp[,46] <- 1/6 * as.vector(t(M %*% ((P- JP) * (P - 2*JP) * (P - 3* JP))) - ((U - JT) * (U - 2*JT) * (U - 3*JT)))

    lp[,90] <- 1/2 * as.vector(NT %*% ((X - JZ) * (X - 2*JZ) * Z))
    lp[,91] <- 1/6 * as.vector(MT %*% (Y * (Y - JZ) * (Y - 2*JZ)))
    lp[,92] <- 1/6 * as.vector(MT %*% (X * (X - JZ) * (X - 2*JZ)))
    lp[,93] <- 1/2 * as.vector(NT %*% (Z * Y * (Y - JZ)))
    lp[,94] <- as.vector(NT %*% (Z * Y * (X - JZ)))
    lp[,95] <- 0.5 * as.vector(MT %*% (X * Y * (Y - JZ)))
    lp[,96] <- 0.5 * as.vector(MT %*% (Y * X * (X - JZ)))
    lp[,97] <- 0.5 * as.vector(NT %*% ((X - JZ) * Z * (Z - JZ)))
    lp[,98] <- 0.5 * as.vector(MT %*% ((Z - JZ) * Y * (Y - JZ)))
    lp[,99] <- 0.5 * as.vector(MT %*% ((Z - JZ) * X * (X - JZ)))
    lp[,100] <- 0.5 * as.vector(NT %*% (Z * (Z - JZ) * Y))
    lp[,101] <- as.vector(MT %*% (X * Y * (Z - JZ)))
    lp[,102] <- 1/6 * as.vector(NT %*% (Z * (Z - JZ) * (Z - 2*JZ)))
    lp[,103] <- 0.5 * as.vector(MT %*% (Y * (Z - JZ) * (Z - 2*JZ)))
    lp[,104] <- 0.5 * as.vector(MT %*% (X * (Z - JZ) * (Z - 2*JZ)))
    lp[,105] <- 1/6 * as.vector(MT %*% ((Z - JZ) * (Z - 2*JZ) * (Z - 3*JZ)) - ((V - JT) * (V - 2*JT) * (V - 3*JT)))
    lp[,106] <- 1/24 * as.vector((V - JT) * (V - 2*JT) * (V - 3*JT) * (V - 4*JT))

    # ---------- TENSOR CASES

    # now clean rows where the edge is not present
    lp <- subset(lp, lp[,1] == 1)

    if (NP >= NZ) { # NP > NZ
      JZ3 <- array(rep(1, NZ * NZ * NZ), c(NZ, NZ, NZ)) # Z x Z x Z array of ones
      KZ3 <- JZ3 # Z x Z x Z matrix, 0 if two indices are equal, 1 otherwise
      for (i in 1 : NZ){
        for (j in 1 : NZ){
          KZ3[i,j,j] <- 0
          KZ3[j,i,j] <- 0
          KZ3[j,j,i] <- 0
        }
      }

      # auxiliary matrices to create A, ..., G
      AP <- maketensor(MT, MT)
      BP <- maketensor(MT, NT)
      CP <- maketensor(NT, MT)
      DP <- maketensor(NT, NT)

      # the following are like the matrices A to G in the Latex File
      # multiply with KZ3 accounts for the terms (1 - delta_{ij}) etc.
      A_ <- tensor::tensor(M, AP, 2, 1) * KZ3
      B_ <- tensor::tensor(M, BP, 2, 1) * KZ3
      C_ <- tensor::tensor(M, CP, 2, 1) * KZ3
      D_ <- tensor::tensor(M, DP, 2, 1) * KZ3
      E_ <- tensor::tensor(N, AP, 2, 1) * KZ3
      F_ <- tensor::tensor(N, BP, 2, 1) * KZ3
      G_ <- tensor::tensor(N, CP, 2, 1) * KZ3

      ####################

      lp[,47] <- 0.5 * large_tensor(DP, A_ * D_ - A_, edges)
      lp[,48] <- 0.5 * large_tensor(AP, G_ * (G_ - JZ3), edges)
      lp[,49] <- 0.25 * large_tensor(AP, D_ * (D_ - JZ3), edges)
      lp[,50] <- large_tensor(BP, D_ * C_, edges)
      lp[,51] <- large_tensor(BP, E_ * F_, edges)
      lp[,52] <- 0.5 * large_tensor(DP, B_ * C_, edges)
      lp[,53] <- large_tensor(AP, D_ * G_, edges)
      lp[,54] <- 0.5 * large_tensor(AP, F_ * G_, edges)
      lp[,55] <- large_tensor(DP, A_ * G_, edges)
      lp[,56] <- large_tensor(DP, B_ * E_, edges)
      lp[,57] <- large_tensor(BP, E_ * G_, edges)
      lp[,58] <- large_tensor(BP, C_ * G_, edges)
      lp[,59] <- large_tensor(BP, D_ * E_, edges)
      lp[,60] <- large_tensor(BP, C_ * F_, edges)
      lp[,61] <- large_tensor(DP, A_ * B_, edges)
      lp[,62] <- large_tensor(CP, A_ * G_, edges)
      lp[,63] <- large_tensor(BP, A_ * D_, edges)
      lp[,64] <- large_tensor(AP, C_ * G_, edges)
      lp[,65] <- large_tensor(AP, B_ * D_, edges)
      lp[,66] <- large_tensor(AP, E_ * G_, edges)
      lp[,67] <- 0.5 * large_tensor(DP, A_ * E_, edges)
      lp[,68] <- large_tensor(BP, A_ * G_, edges)
      lp[,69] <- 0.5 * large_tensor(AP, D_ * E_, edges)
      lp[,70] <- large_tensor(AP, C_ * F_, edges)
      lp[,71] <- 0.25 * large_tensor(DP, A_ * (A_ - JZ3), edges)
      lp[,72] <- large_tensor(AP, (A_ - JZ3) * F_, edges)
      lp[,73] <- 0.5 * large_tensor(AP, (A_ - JZ3) * D_, edges)
      lp[,74] <- 0.5 * large_tensor(BP, E_ * (E_ - JZ3), edges)
      lp[,75] <- 0.5 * large_tensor(BP, C_ * (C_ - JZ3), edges)
      lp[,76] <- large_tensor(BP, (B_ - JZ3) * E_, edges)
      lp[,77] <- large_tensor(BP, (B_ - JZ3) * C_, edges)
      lp[,78] <- large_tensor(BP, A_ * (B_ - JZ3), edges)
      lp[,79] <- 0.25 * large_tensor(AP, E_ * (E_ - JZ3), edges)
      lp[,80] <- 0.5 * large_tensor(AP, B_ * (B_ - JZ3), edges)
      lp[,81] <- large_tensor(BP, C_ * E_, edges)
      lp[,82] <- large_tensor(BP, A_ * E_, edges)
      lp[,83] <- large_tensor(BP, A_ * C_, edges)
      lp[,84] <- large_tensor(AP, B_ * E_, edges)
      lp[,85] <- 0.5 * large_tensor(AP, B_ * C_, edges)
      lp[,86] <- 0.5 * large_tensor(BP, A_ * (A_ - JZ3), edges)
      lp[,87] <- 0.5 * large_tensor(AP, (A_ - JZ3) * E_, edges)
      lp[,88] <- large_tensor(AP, (A_ - JZ3) * B_, edges)
      lp[,89] <- 0.25 * large_tensor(AP, (A_ - JZ3) * (A_ - 2*JZ3) * KZ3, edges)

    }
    if (NP < NZ) { # NP <= NZ
      JP3 <- array(rep(1, NP * NP * NP), c(NP, NP, NP)) # P x P x P array of ones
      KP3 <- JP3 # P x P x P matrix, 0 if two indices are equal, 1 otherwise
      for (i in 1 : NP){
        for (j in 1 : NP){
          KP3[i,j,j] <- 0
          KP3[j,i,j] <- 0
          KP3[j,j,i] <- 0
        }
      }

      ####
      # those are needed if z > p
      AZ <- maketensor(M, M)
      BZ <- maketensor(M, N)
      CZ <- maketensor(N, M)
      DZ <- maketensor(N, N)

      # the following are like the matrices A' to G' in the Latex File
      # multiply with KP3 accounts for the terms (1 - delta_{ij}) etc.
      A_2 <- tensor::tensor(M, AZ, 1, 1) * KP3
      B_2 <- tensor::tensor(M, BZ, 1, 1) * KP3
      C_2 <- tensor::tensor(M, CZ, 1, 1) * KP3
      D_2 <- tensor::tensor(M, DZ, 1, 1) * KP3
      E_2 <- tensor::tensor(N, AZ, 1, 1) * KP3
      F_2 <- tensor::tensor(N, BZ, 1, 1) * KP3
      G_2 <- tensor::tensor(N, CZ, 1, 1) * KP3

      ####
      lp[,47] <- 0.5 * large_tensor(F_2 * (F_2 - JP3), AZ, edges)
      lp[,48] <- 0.5 * large_tensor(A_2 * (D_2 - JP3), DZ, edges)
      lp[,49] <- 0.25 * large_tensor(D_2 * (D_2 - JP3), AZ, edges)
      lp[,50] <- large_tensor(D_2 * G_2, AZ, edges)
      lp[,51] <- large_tensor(A_2 * G_2, DZ, edges)
      lp[,52] <- 0.5 * large_tensor(F_2 * G_2, AZ, edges)
      lp[,53] <- large_tensor(C_2 * D_2, BZ, edges)
      lp[,54] <- 0.5 * large_tensor(B_2 * C_2, DZ, edges)
      lp[,55] <- large_tensor(E_2 * F_2, BZ, edges)
      lp[,56] <- large_tensor(E_2 * G_2, BZ, edges)
      lp[,57] <- large_tensor(B_2 * E_2, DZ, edges)
      lp[,58] <- large_tensor(D_2 * E_2, BZ, edges)
      lp[,59] <- large_tensor(C_2 * G_2, BZ, edges)
      lp[,60] <- large_tensor(C_2 * F_2, BZ, edges)
      lp[,61] <- large_tensor(E_2 * G_2, AZ, edges)
      lp[,62] <- large_tensor(A_2 * F_2, BZ, edges)
      lp[,63] <- large_tensor(C_2 * G_2, AZ, edges)
      lp[,64] <- large_tensor(A_2 * D_2, BZ, edges)
      lp[,65] <- large_tensor(B_2 * D_2, AZ, edges)
      lp[,66] <- large_tensor(A_2 * B_2, DZ, edges)
      lp[,67] <- 0.5 * large_tensor(E_2 * (E_2 - JP3), BZ, edges)
      lp[,68] <- large_tensor((B_2 - JP3) * E_2, BZ, edges)
      lp[,69] <- 0.5 * large_tensor(C_2 * (C_2 - JP3), BZ, edges)
      lp[,70] <- large_tensor(B_2 * (C_2 - JP3), CZ, edges)
      lp[,71] <- 0.25* large_tensor(E_2 * (E_2 - JP3), AZ, edges)
      lp[,72] <- large_tensor(A_2 * (C_2 - JP3), CZ, edges)
      lp[,73] <- 0.5 * large_tensor(C_2 * (C_2 - JP3), AZ, edges)
      lp[,74] <- 0.5 * large_tensor(A_2 * E_2, DZ, edges)
      lp[,75] <- 0.5 * large_tensor(D_2 * E_2, AZ, edges)
      lp[,76] <- large_tensor(A_2 * G_2, BZ, edges)
      lp[,77] <- large_tensor(B_2 * G_2, AZ, edges)
      lp[,78] <- large_tensor((A_2 - JP3) * G_2, AZ, edges)
      lp[,79] <- 0.25 * large_tensor(A_2* (A_2 - JP3), DZ, edges)
      lp[,80] <- 0.5 * large_tensor((A_2 - JP3) * D_2, AZ, edges)
      lp[,81] <- large_tensor(C_2 * E_2, BZ, edges)
      lp[,82] <- large_tensor(A_2 * E_2, BZ, edges)
      lp[,83] <- large_tensor(B_2 * E_2, AZ, edges)
      lp[,84] <- large_tensor(A_2 * C_2, BZ, edges)
      lp[,85] <- 0.5 * large_tensor(B_2 * C_2, AZ, edges)
      lp[,86] <- 0.5 * large_tensor((A_2 - JP3) * E_2, AZ, edges)
      lp[,87] <- 0.5 * large_tensor(A_2 * (A_2 - JP3), BZ, edges)
      lp[,88] <- large_tensor((A_2 - JP3) * C_2, AZ, edges)
      lp[,89] <- 0.25 * large_tensor((A_2 - JP3) * (A_2 - 2*JP3) * KP3, AZ, edges)
    }
  }

  #######################
  # now clean rows where the edge is not present
  lp <- subset(lp, lp[,1] == 1)

  # --------------- CONSIDER WEIGHTS -------------------------------------
  if (weights) {
    for (k in 1:nrow(lp)) {
      lp[k,] <- lp[k,] * edges[k,4]
    }
  }

  # --------------- NORMALISE --------------------------------------------

  if (normalisation == "none") {
    return(as.data.frame(lp))
  } else {
    norm_lp <- normalise_link_positions(lp, type = normalisation, six_node = six_node)
    return(as.data.frame(norm_lp))
  }

}


Call:
lm(formula = y ~ x)

Residuals:
1       2       3       4       5       6
3.3333 -0.6667 -2.6667 -2.6667 -0.6667  3.3333

Coefficients:
            Estimate Std. Error t value Pr(>|t|)
(Intercept)  -9.3333     2.8441  -3.282 0.030453 *
x             7.0000     0.7303   9.585 0.000662 ***
---
Signif. codes:  0 ‘***’ 0.001 ‘**’ 0.01 ‘*’ 0.05 ‘.’ 0.1 ‘ ’ 1

Residual standard error: 3.055 on 4 degrees of freedom
Multiple R-squared: 0.9583,     Adjusted R-squared: 0.9478
F-statistic: 91.88 on 1 and 4 DF,  p-value: 0.000662

> par(mfrow=c(2, 2))     # Request 2x2 plot layout
> plot(lm_1)             # Diagnostic plot of regression model


# This code computes the IR/IFR rates
# for COVID-19 pandemic in Saint Petersburg, Russia
# paper with Bayesian evidence synthesis
#
# Code builds on https://github.com/harlanhappydog/BayesianSeroMetaAnalysis
# by Campbell and Gustafson (2021, https://www.medrxiv.org/content/10.1101/2021.05.12.21256975v2)
library(data.table)
library(stringi)
library(stringr)
library(lubridate)
library(binom)
library(rjags)
library(hBayesDM)

# Declare working directory beforehand in an environment variable
# SPB_COVID_STUDY_PATH = "path_to_your_folder"
# with the aid of usethis::edit_r_environ()
# Restart R session for the changes to make effect
setwd(Sys.getenv('SPB_COVID_STUDY_PATH'))

#######
# Load and prepare the estimated prevalence across waves
# (see the estimation code in analysis/wave*/estimate_prevalence_by_{model,variable}.r)
waves <- c(1, 1.5, 2, 3)

# Init an object to store the results
# of prevalence computations
perwave_prevalence_by_model_list <- list()

# Load the per-wave estimates and store them in one object
for(w in waves) {

	load(paste0("estimates/wave", w, "/prevalence_by_model.rdata"))

	perwave_prevalence_by_model_list[[ as.character(w) ]] <- prevalence_by_model

	rm(prevalence_by_model)

}

# Aggregate the results
perwave_prevalence_by_model <- rbindlist(perwave_prevalence_by_model_list, idcol = "wave")

# We are interested in per-wave 95% confidence intervals for ELISA Coronapass
# adjusted for non-response and test characteristics
model_data <- perwave_prevalence_by_model[ serotest == "IgA_or_G_or_M_testC" & prevalence_type == "univariate" & surveyweight == "no" & sensitivity != 1 & scenario == "basic", c("wave", "lowerbound", "upperbound")]
setnames(model_data, "wave", "k")

# Manually add wave start and end dates
model_data[ k == 1, start_date := ymd("2020-05-25") ]
model_data[ k == 1, end_date := ymd("2020-06-28") ]
model_data[ k == 1.5, start_date := ymd("2020-07-20") ]
model_data[ k == 1.5, end_date := ymd("2020-08-16") ]
model_data[ k == 2, start_date := ymd("2020-10-12") ]
model_data[ k == 2, end_date := ymd("2020-12-06") ]
model_data[ k == 3, start_date := ymd("2021-02-15") ]
model_data[ k == 3, end_date := ymd("2021-04-04") ]

# Manually add populations
## Adult population
model_data[ k %in% c(1, 1.5), p_k_adult := 4451025 ]
model_data[ k %in% c(2, 3), p_k_adult := 4439897 ]
#
model_data[ k %in% c(1, 1.5), p_k_all := 5398064 ]
model_data[ k %in% c(2, 3), p_k_all := 5384342 ]

#######
# Find the parameters of the binomial distribution that
# correspond to a given CI for seroprevalence (with given tolerance)

# Compute CIs space with all combinations
# of cc_k and t_k and their CI
ci_space <- CJ(cc_k = 1:1000, t_k = 2:1000)
ci_space <- ci_space[ cc_k < t_k]
# NB: the below takes time and memory
ci_space[, c("lower", "upper") := binom::binom.bayes(x = cc_k, n = t_k, type = "central", conf.level = 0.95, maxit = 200)[, c("lower", "upper")]]

search_binomial_parameters <- function(search_space, lowerbound, upperbound) {

	# Euclidean distance from the target
	search_space[, eucl_dist := sqrt((lower - lowerbound)^2 + (upper - upperbound)^2)]
	# Sort from smallest distance
	setorderv(search_space, "eucl_dist", 1)
	# Return the parameters producing the smallest distance
	smallest_distance_params <- unlist(search_space[1, c("cc_k", "t_k")])
	return(smallest_distance_params)

}

# Run and assign the corresponding values
model_data[, c("cc_k", "t_k") := NA_real_]

for(i in 1:nrow(model_data) ) {

	search_results <- search_binomial_parameters(search_space = ci_space,
												lowerbound = model_data[i]$lowerbound/100,
												upperbound = model_data[i]$upperbound/100)

	set(model_data, i = i, j = "cc_k", value = search_results["cc_k"])
	set(model_data, i = i, j = "t_k", value = search_results["t_k"])

}

# Clean up
rm(ci_space)
gc()

#model_data[, c("k", "lowerbound", "upperbound", "t_k", "cc_k")]
#     k lowerbound upperbound t_k cc_k
#1:   1        7.7       11.7 827   79
#2: 1.5        9.9       16.6 385   50
#3:   2       20.3       25.5 999  228
#4:   3       39.7       48.0 550  241

#######
# Get the data on COVID or excess deaths

# Daily covid deaths from stopcoronavirus.rf
stopcoronavirus_deaths <- fread("https://raw.githubusercontent.com/alexei-kouprianov/COVID-19.SPb/main/data/primary/covid.SPb.stopkoronavirus.rf.txt")
stopcoronavirus_deaths <- stopcoronavirus_deaths[, c("TIME", "d")]
setnames(stopcoronavirus_deaths, c("TIME", "d"), c("day", "covid_deaths"))
setorderv(stopcoronavirus_deaths, "day")

# Monthly excess deaths from https://doi.org/10.1111/1740-9713.01486
excess_deaths <- fread("https://raw.githubusercontent.com/dkobak/excess-mortality/main/russia_excess_deaths.csv", encoding = "UTF-8")[ V1 == "Санкт-Петербург"]
excess_deaths <- melt(excess_deaths, id.vars = "V1", variable.name = "month", value.name = "excess_deaths", variable.factor = F)
excess_deaths[, V1 := NULL]
excess_deaths[, month := gsub("Январь ", "01.01.", month)]
excess_deaths[, month := gsub("Февраль ", "01.02.", month)]
excess_deaths[, month := gsub("Март ", "01.03.", month)]
excess_deaths[, month := gsub("Апрель ", "01.04.", month)]
excess_deaths[, month := gsub("Май ", "01.05.", month)]
excess_deaths[, month := gsub("Июнь ", "01.06.", month)]
excess_deaths[, month := gsub("Июль ", "01.07.", month)]
excess_deaths[, month := gsub("Август ", "01.08.", month)]
excess_deaths[, month := gsub("Сентябрь ", "01.09.", month)]
excess_deaths[, month := gsub("Октябрь ", "01.10.", month)]
excess_deaths[, month := gsub("Ноябрь ", "01.11.", month)]
excess_deaths[, month := gsub("Декабрь ", "01.12.", month)]
excess_deaths[, month := dmy(month)]
setorderv(excess_deaths, "month")

# Compute the respective D^upper/lower_k for each wave
for(i in 1:nrow(model_data) ) {

	wave_start <- model_data$start_date[i]
	wave_end <- model_data$end_date[i]

	d_deaths_lower <- sum(stopcoronavirus_deaths[ day < wave_start + 14 ]$covid_deaths, na.rm = T)
	d_deaths_upper <- sum(stopcoronavirus_deaths[ day < wave_end + 14 ]$covid_deaths, na.rm = T)

	d_excess_lower <- sum(excess_deaths[ month <= floor_date(wave_start, "month") ]$excess_deaths, na.rm = T)
	d_excess_upper <- sum(excess_deaths[ month <= floor_date(wave_end, "month") ]$excess_deaths, na.rm = T)

	set(model_data, i = i, j = "d_deaths_lower", value = d_deaths_lower)
	set(model_data, i = i, j = "d_deaths_upper", value = d_deaths_upper)

	set(model_data, i = i, j = "d_excess_lower", value = d_excess_lower)
	set(model_data, i = i, j = "d_excess_upper", value = d_excess_upper)

}

#######
# Build and estimate the evidence synthesis model
# (borrows from https://github.com/harlanhappydog/BayesianSeroMetaAnalysis)

# Number of periods in the data
K <- nrow(model_data)

# JAGS model likelihood and summary
model_likelihood <- "

# Likelihood:

for(k in 1:K){
	cc[k] ~ dbin(ir[k], tests[k]);
	censor.index[k] ~ dinterval(deaths[k], c(deaths_lower[k], deaths_upper[k]))
	deaths[k] ~ dbin(ifr[k]*ir[k], pop[k]);
	cloglog(ir[k]) <- cloglog_ir[k];
	cloglog(ifr[k]) <- cloglog_ifr[k];

	cloglog_ir[k] ~ dnorm(beta, inv.var_sig);
	cloglog_ifr[k] ~ dnorm(theta, inv.var_tau);

}

# Summary:

g_IFR = theta;
IFR <- 1 - exp(-exp(g_IFR))
g_IR = beta;
IR <- 1 - exp(-exp(g_IR))
"

# Two types of priors
model_weakly_informative_priors <- "

# Weakly informative priors:
	icloglog_theta ~ dbeta(0.3, 30);
	icloglog_beta  ~ dbeta(1, 30);
	theta <- log(-log(1-icloglog_theta));
	beta  <- log(-log(1-icloglog_beta));
	inv.var_sig   <- (1/sigma)^2 ;
	inv.var_tau   <- (1/tau)^2 ;
	sigma     ~ dnorm(0, 1/10) T(0,);
	tau     ~ dnorm(0, 1/10) T(0,);

"

model_non_informative_priors <- "

# Non-informative priors:
	icloglog_theta ~ dunif(0, 1);
	icloglog_beta  ~ dunif(0, 1);
	theta <- log(-log(1-icloglog_theta));
	beta  <- log(-log(1-icloglog_beta));
	inv.var_sig   <- (1/sigma)^2 ;
	inv.var_tau   <- (1/tau)^2 ;
	sigma     ~ dnorm(0, 1/100) T(0,);
	tau     ~ dnorm(0, 1/100) T(0,);

"

# Gather into the full JAGS model statements
model_statement_weakly_priors <- paste0("model {", model_weakly_informative_priors, model_likelihood, "\n}")
model_statement_noninformative_priors <- paste0("model {", model_non_informative_priors, model_likelihood, "\n}")

# Build the models
## Weakly informative priors, official deaths
jags.model_weakly_priors_official_deaths <- jags.model(textConnection(model_statement_weakly_priors),
	data = list(
		K = K,
		tests = model_data$t_k,
		cc = model_data$cc_k,
		pop = model_data$p_k_adult,
		deaths_lower = model_data$d_deaths_lower,
		deaths_upper = model_data$d_deaths_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data$d_deaths_lower, model_data$d_deaths_upper), 1, mean))
	)
)

## Weakly informative priors, excess deaths
jags.model_weakly_priors_excess_deaths <- jags.model(textConnection(model_statement_weakly_priors),
	data = list(
		K = K,
		tests = model_data$t_k,
		cc = model_data$cc_k,
		pop = model_data$p_k_adult,
		deaths_lower = model_data$d_excess_lower,
		deaths_upper = model_data$d_excess_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data$d_excess_lower, model_data$d_excess_upper), 1, mean))
	)
)

# Non-informative priors, official deaths
jags.model_noninformative_priors_official_deaths <- jags.model(textConnection(model_statement_noninformative_priors),
	data = list(
		K = K,
		tests = model_data$t_k,
		cc = model_data$cc_k,
		pop = model_data$p_k_adult,
		deaths_lower = model_data$d_deaths_lower,
		deaths_upper = model_data$d_deaths_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data$d_deaths_lower, model_data$d_deaths_upper), 1, mean))
	)
)

# Non-informative priors, excess deaths
jags.model_noninformative_priors_excess_deaths <- jags.model(textConnection(model_statement_noninformative_priors),
	data = list(
		K = K,
		tests = model_data$t_k,
		cc = model_data$cc_k,
		pop = model_data$p_k_adult,
		deaths_lower = model_data$d_excess_lower,
		deaths_upper = model_data$d_excess_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data$d_excess_lower, model_data$d_excess_upper), 1, mean))
	)
)

## Weakly informative priors, excess deaths, all population
jags.model_weakly_priors_excess_deaths_all_pop <- jags.model(textConnection(model_statement_weakly_priors),
	data = list(
		K = K,
		tests = model_data$t_k,
		cc = model_data$cc_k,
		pop = model_data$p_k_all,
		deaths_lower = model_data$d_excess_lower,
		deaths_upper = model_data$d_excess_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data$d_excess_lower, model_data$d_excess_upper), 1, mean))
	)
)

params <- c("theta", "ir", "ifr", "tau", "sigma", "IFR", "IR")

# Set the number of iterations
MCMCn <- 2e6

# Sample from the models

sample_weakly_priors_official_deaths <- coda.samples(jags.model_weakly_priors_official_deaths, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

sample_weakly_priors_excess_deaths <- coda.samples(jags.model_weakly_priors_excess_deaths, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

sample_noninformative_priors_official_deaths <- coda.samples(jags.model_noninformative_priors_official_deaths, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

sample_noninformative_priors_excess_deaths <- coda.samples(jags.model_noninformative_priors_excess_deaths, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

sample_weakly_priors_excess_deaths_all_pop <- coda.samples(jags.model_weakly_priors_excess_deaths_all_pop, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

# Save the estimation results
save(sample_weakly_priors_official_deaths, file = "estimates/pandemic_course_paper/sample_weakly_priors_official_deaths.rdata", compress = "gzip")
save(sample_weakly_priors_excess_deaths, file = "estimates/pandemic_course_paper/sample_weakly_priors_excess_deaths.rdata", compress = "gzip")
save(sample_noninformative_priors_official_deaths, file = "estimates/pandemic_course_paper/sample_noninformative_priors_official_deaths.rdata", compress = "gzip")
save(sample_noninformative_priors_excess_deaths, file = "estimates/pandemic_course_paper/sample_noninformative_priors_excess_deaths.rdata", compress = "gzip")
save(sample_weakly_priors_excess_deaths_all_pop, file = "estimates/pandemic_course_paper/sample_weakly_priors_excess_deaths_all_pop.rdata", compress = "gzip")

#######
# Report the results
load("estimates/pandemic_course_paper/sample_weakly_priors_official_deaths.rdata")
load("estimates/pandemic_course_paper/sample_weakly_priors_excess_deaths.rdata")
load("estimates/pandemic_course_paper/sample_noninformative_priors_official_deaths.rdata")
load("estimates/pandemic_course_paper/sample_noninformative_priors_excess_deaths.rdata")
load("estimates/pandemic_course_paper/sample_weakly_priors_excess_deaths_all_pop.rdata")

# Model summary function
summary_md_ci <- function(xx, samps) {
	c(	md = summary(samps)$quantiles[xx, "50%"],
		lower = HDIofMCMC(unlist(((samps[[1]])[,xx])), credMass = 0.95)[1],
		higher = HDIofMCMC(unlist(((samps[[1]])[,xx])), credMass = 0.95)[2]
	)
}

# Create the confidence intervals
ir_irf_results_object <- CJ(k = 1:K, type = c("ir", "ifr"), deaths = c("official", "excess"), priors = c("weakly", "noninformative"), population = c("", "_all_pop"))
ir_irf_results_object[, c("pointest", "lower", "upper", "ir_overall", "ir_overall_lower", "ir_overall_upper", "ifr_overall", "ifr_overall_lower", "ifr_overall_upper") := NA_real_ ]

for(i in 1:nrow(ir_irf_results_object)) {

	# Compute the confidence intervals programmatically
	results_object_name <- paste0("sample_", ir_irf_results_object$priors[i], "_priors_", ir_irf_results_object$deaths[i], "_deaths", ir_irf_results_object$population[i])

	if( exists(results_object_name) ) {

		# Example: summary_md_ci('ir[1]', sample_weakly_priors_official_deaths)
		# Per-wave IR/IFR
		eval(parse(text = paste0("res <- summary_md_ci('", ir_irf_results_object$type[i], "[", ir_irf_results_object$k[i], "]', ", results_object_name, ")")))
		set(ir_irf_results_object, i = i, j = "pointest", value = res["md"])
		set(ir_irf_results_object, i = i, j = "lower", value = res["lower"])
		set(ir_irf_results_object, i = i, j = "upper", value = res["higher"])
		# Overall IR
		eval(parse(text = paste0("ir_overall <- summary_md_ci('IR', ", results_object_name, ")")))
		set(ir_irf_results_object, i = i, j = "ir_overall", value = ir_overall["md"])
		set(ir_irf_results_object, i = i, j = "ir_overall_lower", value = ir_overall["lower"])
		set(ir_irf_results_object, i = i, j = "ir_overall_upper", value = ir_overall["higher"])
		# Overall IFR
		eval(parse(text = paste0("ifr_overall <- summary_md_ci('IFR', ", results_object_name, ")")))
		set(ir_irf_results_object, i = i, j = "ifr_overall", value = ifr_overall["md"])
		set(ir_irf_results_object, i = i, j = "ifr_overall_lower", value = ifr_overall["lower"])
		set(ir_irf_results_object, i = i, j = "ifr_overall_upper", value = ifr_overall["higher"])

	}
}

# Fix the waves
ir_irf_results_object[, k := as.numeric(k)]
ir_irf_results_object[ k == 2, k := 1.5 ]
ir_irf_results_object[ k == 3, k := 2 ]
ir_irf_results_object[ k == 4, k := 3 ]

# To percentages and round
ir_irf_results_object[, c("pointest", "lower", "upper", "ir_overall", "ir_overall_lower", "ir_overall_upper", "ifr_overall", "ifr_overall_lower", "ifr_overall_upper") := lapply(.SD, function(x) round(100*x, 2)), .SDcols = c("pointest", "lower", "upper", "ir_overall", "ir_overall_lower", "ir_overall_upper", "ifr_overall", "ifr_overall_lower", "ifr_overall_upper")]

# Remove empty
ir_irf_results_object <- ir_irf_results_object[!is.na(pointest)]

# Save point
save(ir_irf_results_object, file = "estimates/pandemic_course_paper/ir_irf_results_object.rdata", compress = "gzip")

############################################################################
# Per agegroup-sex analysis

# Load the per-age group/sex seroprevalence estimates
# created by estimate_prevalence_by_agegroup_sex.r
load("estimates/pandemic_course_paper/prevalence_by_agegroup_sex.rdata")
model_data_agesex <- prevalence_by_agegroup_sex[ prevalence_type == "univariate" ]

# Create new CI search space
# with all combinations
# of cc_k and t_k and their CI
ci_space <- CJ(cc_k = 1:1000, t_k = 2:1000)
ci_space <- ci_space[ cc_k < t_k]
# NB: the below takes time and memory
ci_space[, c("lower", "upper") := binom::binom.bayes(x = cc_k, n = t_k, type = "central", conf.level = 0.95, maxit = 200)[, c("lower", "upper")]]

# Run and assign the corresponding values
model_data_agesex[, c("cc_k", "t_k") := NA_real_]

for(i in 1:nrow(model_data_agesex) ) {

	search_results <- search_binomial_parameters(search_space = ci_space,
												lowerbound = model_data_agesex[i]$lowerbound/100,
												upperbound = model_data_agesex[i]$upperbound/100)

	set(model_data_agesex, i = i, j = "cc_k", value = search_results["cc_k"])
	set(model_data_agesex, i = i, j = "t_k", value = search_results["t_k"])

}

# Add population data
model_data_agesex[ agegroup == "18-29" & male == 0, p_k := 419516]
model_data_agesex[ agegroup == "30-39" & male == 0, p_k := 543892]
model_data_agesex[ agegroup == "40-49" & male == 0, p_k := 432308]
model_data_agesex[ agegroup == "50-59" & male == 0, p_k := 418371]
model_data_agesex[ agegroup == "60-69" & male == 0, p_k := 437588]
model_data_agesex[ agegroup == "70+" & male == 0, p_k := 482856]

model_data_agesex[ agegroup == "18-29" & male == 1, p_k := 413098]
model_data_agesex[ agegroup == "30-39" & male == 1, p_k := 527930]
model_data_agesex[ agegroup == "40-49" & male == 1, p_k := 404853]
model_data_agesex[ agegroup == "50-59" & male == 1, p_k := 342243]
model_data_agesex[ agegroup == "60-69" & male == 1, p_k := 281789]
model_data_agesex[ agegroup == "70+" & male == 1, p_k := 241605]

# Add excess deaths data
model_data_agesex[ agegroup == "18-29" & male == 0, d_excess_lower := 0]
model_data_agesex[ agegroup == "18-29" & male == 0, d_excess_upper := 36]
model_data_agesex[ agegroup == "30-39" & male == 0, d_excess_lower := 57]
model_data_agesex[ agegroup == "30-39" & male == 0, d_excess_upper := 182]
model_data_agesex[ agegroup == "40-49" & male == 0, d_excess_lower := 167]
model_data_agesex[ agegroup == "40-49" & male == 0, d_excess_upper := 341]
model_data_agesex[ agegroup == "50-59" & male == 0, d_excess_lower := 367]
model_data_agesex[ agegroup == "50-59" & male == 0, d_excess_upper := 596]
model_data_agesex[ agegroup == "60-69" & male == 0, d_excess_lower := 1412]
model_data_agesex[ agegroup == "60-69" & male == 0, d_excess_upper := 1755]
model_data_agesex[ agegroup == "70+" & male == 0, d_excess_lower := 7679]
model_data_agesex[ agegroup == "70+" & male == 0, d_excess_upper := 8509]

model_data_agesex[ agegroup == "18-29" & male == 1, d_excess_lower := 0]
model_data_agesex[ agegroup == "18-29" & male == 1, d_excess_upper := 105]
model_data_agesex[ agegroup == "30-39" & male == 1, d_excess_lower := 161]
model_data_agesex[ agegroup == "30-39" & male == 1, d_excess_upper := 352]
model_data_agesex[ agegroup == "40-49" & male == 1, d_excess_lower := 528]
model_data_agesex[ agegroup == "40-49" & male == 1, d_excess_upper := 791]
model_data_agesex[ agegroup == "50-59" & male == 1, d_excess_lower := 856]
model_data_agesex[ agegroup == "50-59" & male == 1, d_excess_upper := 1167]
model_data_agesex[ agegroup == "60-69" & male == 1, d_excess_lower := 2111]
model_data_agesex[ agegroup == "60-69" & male == 1, d_excess_upper := 2527]
model_data_agesex[ agegroup == "70+" & male == 1, d_excess_lower := 5330]
model_data_agesex[ agegroup == "70+" & male == 1, d_excess_upper := 5928]

# Add the row index
model_data_agesex[, k := 1:.N ]

# Define the models
K <- nrow(model_data_agesex)

# Set the number of iterations
MCMCn <- 2e6

# Model parameters
params <- c("theta", "ir", "ifr", "tau", "sigma", "IFR", "IR")

## Weakly informative priors, excess deaths
jags.model_weakly_priors_excess_deaths_agesex <- jags.model(textConnection(model_statement_weakly_priors),
	data = list(
		K = K,
		tests = model_data_agesex$t_k,
		cc = model_data_agesex$cc_k,
		pop = model_data_agesex$p_k,
		deaths_lower = model_data_agesex$d_excess_lower,
		deaths_upper = model_data_agesex$d_excess_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data_agesex$d_excess_lower, model_data_agesex$d_excess_upper), 1, mean))
	)
)

# Non-informative priors, excess deaths
jags.model_noninformative_priors_excess_deaths_agesex <- jags.model(textConnection(model_statement_noninformative_priors),
	data = list(
		K = K,
		tests = model_data_agesex$t_k,
		cc = model_data_agesex$cc_k,
		pop = model_data_agesex$p_k,
		deaths_lower = model_data_agesex$d_excess_lower,
		deaths_upper = model_data_agesex$d_excess_upper,
		deaths = rep(NA, K),
		censor.index = rep(1, K)
	),
	n.chains = 5,
	n.adapt = 5000,
	inits = list(
		deaths = round(apply(cbind(model_data_agesex$d_excess_lower, model_data_agesex$d_excess_upper), 1, mean))
	)
)

# Sample from the models
sample_weakly_priors_excess_deaths_agesex <- coda.samples(jags.model_weakly_priors_excess_deaths_agesex, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

sample_noninformative_priors_excess_deaths_agesex <- coda.samples(jags.model_noninformative_priors_excess_deaths_agesex, params,
				n.iter = MCMCn, thin = 100, n.adapt = round(MCMCn*0.20))

# Save the estimation results
save(sample_weakly_priors_excess_deaths_agesex, file = "estimates/pandemic_course_paper/sample_weakly_priors_excess_deaths_agesex.rdata", compress = "gzip")
save(sample_noninformative_priors_excess_deaths_agesex, file = "estimates/pandemic_course_paper/sample_noninformative_priors_excess_deaths_agesex.rdata", compress = "gzip")

#######
# Report the results
load("estimates/pandemic_course_paper/sample_weakly_priors_excess_deaths_agesex.rdata")
load("estimates/pandemic_course_paper/sample_noninformative_priors_excess_deaths_agesex.rdata")

# Model summary function
summary_md_ci <- function(xx, samps) {
	c(	md = summary(samps)$quantiles[xx, "50%"],
		lower = HDIofMCMC(unlist(((samps[[1]])[,xx])), credMass = 0.95)[1],
		higher = HDIofMCMC(unlist(((samps[[1]])[,xx])), credMass = 0.95)[2]
	)
}

# Create the confidence intervals
ir_irf_agesex_results_object <- CJ(k = 1:K, type = c("ir", "ifr"), deaths = c("excess"), priors = c("weakly", "noninformative"))
ir_irf_agesex_results_object[, c("pointest", "lower", "upper") := NA_real_ ]

# Add agegroup and sex
ir_irf_agesex_results_object <- merge(ir_irf_agesex_results_object, model_data_agesex[, c("k", "agegroup", "male"), with = F], by = "k", all.x = T, all.y = F)

for(i in 1:nrow(ir_irf_agesex_results_object)) {

	# Compute the confidence intervals programmatically
	results_object_name <- paste0("sample_", ir_irf_agesex_results_object$priors[i], "_priors_", ir_irf_agesex_results_object$deaths[i], "_deaths_agesex")

	if( exists(results_object_name) ) {

		# Example: summary_md_ci('ir[1]', sample_weakly_priors_official_deaths)
		# Per-wave IR/IFR
		eval(parse(text = paste0("res <- summary_md_ci('", ir_irf_agesex_results_object$type[i], "[", ir_irf_agesex_results_object$k[i], "]', ", results_object_name, ")")))
		set(ir_irf_agesex_results_object, i = i, j = "pointest", value = res["md"])
		set(ir_irf_agesex_results_object, i = i, j = "lower", value = res["lower"])
		set(ir_irf_agesex_results_object, i = i, j = "upper", value = res["higher"])

	}
}

# To percentages and round
ir_irf_agesex_results_object[, c("pointest", "lower", "upper") := lapply(.SD, function(x) round(100*x, 2)), .SDcols = c("pointest", "lower", "upper")]

# Remove empty
ir_irf_agesex_results_object <- ir_irf_agesex_results_object[!is.na(pointest)]

# Save point
save(ir_irf_agesex_results_object, file = "estimates/pandemic_course_paper/ir_irf_agesex_results_object.rdata", compress = "gzip")


#'@title calcHMWakeFraction
#'
#'@description Calculate wake fraction (\code{wakeFraction}) (dimensionless)
#'using the Holtrop & Mennen method.
#'
#'@param breadth Moulded breadth (vector of numericals, m)
#'@param wettedSA Wetted hull surface area (vector of numericals, m^2) (see
#' \code{\link{calcHMWettedSA}})
#'@param maxDraft Maximum summer load line draft (vector of numericals, m)
#'@param nProp Number of propellers (vector of numericals, see
#' \code{\link{calcPropNum}})
#'@param lwl Waterline Length (vector of numericals, m) (see \code{\link{calclwl}})
#'@param propDiam Propeller diameter (vector of numericals, m) (see
#' \code{\link{calcPropDia}})
#'@param formFactor Form factor (1+k) (vector of numericals, dimensionless) (see
#' \code{\link{calcHMFormFactor}})
#'@param Cf Frictional resistance coefficient (vector of numericals, dimensionless)
#' (see \code{\link{calcCf}})
#'@param Ca Incremental hull (roughness) resistance coefficient (vector of
#' numericals, dimensionless) (see \code{\link{calcHMCa}})
#'@param Cbw Waterline block coefficient (vector of numericals, dimensionless)
#' (see \code{\link{calcCbw}})
#'@param Cp Prismatic coefficient (vector of numericals, dimensionless) (see
#' \code{\link{calcCp}})
#'@param Cm Midship section coefficient (vector of numericals, dimensionless)
#' (see \code{\link{calcCm}})
#'@param aftDraft Aft draft (deviation from actual draft indicates trim)
#' (vector of numericals, m)
#'@param Cstern Afterbody form coefficient: \itemize{\item V-shaped Hull = -10
#'         \item U-Shaped Hull =  10
#'         \item Normal Hull = 0 (default)}
#' Can supply either a vector of numericals, a single number, or rely on the default
#'@param lcb Longitudinal position of center of buoyancy (vector of numericals,
#' see \code{\link{calclcb}})
#'
#'@details
#' "The speed of advance of the propeller relative to the water in which it is working is
#'  lower than the observed speed of the vessel. This difference in speed,
#'  expressed as a percentage of the ship speed, is known as the wake fraction coefficient".
#'  \url{https://www.wartsila.com/encyclopedia/term/wake-fraction-coefficient}
#'
#'Wake fraction is a component of hull efficiency as well as a component of
#'propeller efficiency.
#'
#'Actual draft is typically obtained from sources such as AIS messages or ship
#'records.
#'
#'Note: In "A Statistical Re-Analysis of Resistance and Propulsion Data", the
#'authors re-analyze with the inclusion of Series 64 hull forms for a total of
#'334 models included in the analysis. They suggest an update of the single
#'screw wake fraction equation but that the original equations should be used
#'for twin screw ships.
#'
#'Viscous resistance coefficient: \code{Cv = (1+k) * Cf + Ca}
#'
#'We are assuming here that \code{1+k = 1+k_1} from \code{\link{calcHMFormFactor}}.
#'
#'Additionally, we are assuming conventional stern for all single screw ships.
#'Holtrop & Mennen also include an estimation for wake fraction for single screw
#'ships with open stern for fast sailing ships, but that is not included here.
#'
#'@return \code{wakeFraction} (vector of numericals, dimensionless)
#'
#'@references
#'Holtrop, J. and Mennen, G. G. J. 1982. "An approximate power prediction
#'method." International Shipbuilding Progress 29.
#'
#'Holtrop, J. and Mennen, G. G. J. 1984. "A Statistical Re-Analysis of Resistance
#'and Propulsion Data'.
#'
#'@seealso \itemize{
#'\item \code{\link{calcPropNum}}
#'\item \code{\link{calclwl}}
#'\item \code{\link{calcPropDia}}
#'\item \code{\link{calcCf}}
#'\item \code{\link{calcCbw}}
#'\item \code{\link{calcCp}}
#'\item \code{\link{calcCm}}
#'\item \code{\link{calclcb}}
#'}
#'
#'@family Holtrop-Mennen Calculations
#'
#'@examples
#' calcHMWakeFraction(c(32.25,32.20),
#'                    c(10746.28,8669.7),
#'                    c(13.57,11.49),
#'                    c(1,1),
#'                    c(218.75,209.25),
#'                    c(6.7,7),
#'                    c(1.27,1.18),
#'                    c(0.0015,0.0014),
#'                    c(0.00033,0.00035),
#'                    c(0.81,0.65),
#'                    c(0.81,0.67),
#'                    c(0.99,0.98),
#'                    c(13.57,11.49))
#'@export

calcHMWakeFraction<- function(breadth, wettedSA,maxDraft,nProp,
                              lwl,propDiam,formFactor,Cf,Ca,Cbw,
                    Cp,Cm,aftDraft=maxDraft,Cstern=0,lcb=0){

    c8<-ifelse(breadth/aftDraft<5,#case 1
      breadth*wettedSA/(lwl*propDiam*aftDraft)
    ,#case 2
      wettedSA*(7*breadth/aftDraft-25)/(lwl*propDiam*(breadth/aftDraft-3))
    )
#================
    c9<- ifelse(c8<28,#case 1
    c8,#case 2
      32-(16/(c8-24 ))
    )
#================
    c11<- ifelse(aftDraft/propDiam<2,#case 1
      aftDraft/propDiam
   , #case 2
    0.0833333*((aftDraft/propDiam)^3)+1.33333
    )
#================
c19<-ifelse(Cp<0.7,
            0.12997/(0.95-Cbw)-0.11056/(0.95-Cp),
            0.18567/(1.3571-Cm)-0.71276+0.38648*Cp
            )
#================

# wakeFraction<- ifelse(nProp==1,#begin single screw estimation using [2] method
#   c9*(#c20
#     1+0.015*Cstern
#   )*
#     (#Cv
#       formFactor*Cf+Ca
#     )*(lwl/aftDraft)*(0.050776+0.93405*c11*(
#       (#Cv
#         formFactor*Cf+Ca
#       )/(1-(#Cp1
#         1.45*Cp-0.315-0.0225*lcb
#       ))
#     )
#     )+0.27915*(#c20
#       1+0.015*Cstern
#     )*sqrt(
#       breadth/(lwl*(1-(#Cp1
#         1.45*Cp-0.315-0.0225*lcb
#       )))
#     )+c19*
#     (#c20
#       1+0.015*Cstern
#     )
# #end single screw estimation
# ,#begin twin screw estimation using [1] method
#  0.3095*Cbw+10*
#     (#Cv
#       formFactor*Cf+Ca
#     )*Cbw-(0.23*propDiam/sqrt(breadth*maxDraft))
# )#end twin screw estimation

wakeFraction<-Cp

#single propeller and will not create error by square rooting negative values
index<-which(nProp==1 & (breadth/(lwl*(1-(1.45*Cp-0.315-0.0225*lcb))))>0)
wakeFraction[index]<- c9[index]*(#c20
                1+0.015*Cstern
              )*
              (#Cv
                formFactor[index]*Cf[index]+Ca[index]
              )*(lwl[index]/aftDraft[index])*(0.050776+0.93405*c11[index]*(
                (#Cv
                  formFactor[index]*Cf[index]+Ca[index]
                )/(1-(#Cp1
                  1.45*Cp[index]-0.315-0.0225*lcb[index]
                ))
              )
              )+0.27915*(#c20
                1+0.015*Cstern
              )*sqrt(
                breadth[index]/(lwl[index]*(1-(#Cp1
                  1.45*Cp[index]-0.315-0.0225*lcb[index]
                )))
              )+c19[index]*
              (#c20
                1+0.015*Cstern
              )

#twin propeller or those that will create errors with single propeller method
index<-which(nProp==2 | (breadth/(lwl*(1-(1.45*Cp-0.315-0.0225*lcb))))<0)
wakeFraction[index]<- 0.3095*Cbw[index]+10*
            (#Cv
              formFactor[index]*Cf[index]+Ca[index]
            )*Cbw[index]-(0.23*propDiam[index]/sqrt(breadth[index]*maxDraft[index]))

return(wakeFraction)
}