#-------------------------
# Libraries
#-------------------------

library(survival)
library(mice)
library(mitools)


#-------------------------
# Data sets
#-------------------------

# full cohort data
data.cohort<-read.table("tb.cohort.txt",header=T)

# case-cohort data
data.caco<-read.table("tb.cascoh.txt",header=T)

# nested case-control data
data.ncc<-read.table("tb.ncc.txt",header=T)


#--------------------------------------------------------
# Correlation between dose and nof (on transformed scale)
#--------------------------------------------------------

# generate log2(1+nof/100) [log is base 2]
data.cohort$l2nof<-log2(1+data.cohort$nof/100)
                        
# generate log2(dose+1) [log is base 2]
data.cohort$l2d<-log2(data.cohort$dose+1)

cor(data.cohort$l2d,data.cohort$l2nof)
plot(data.cohort$l2d,data.cohort$l2nof)


#-----------------------------------
# Correlation between dose and ageRx
#-----------------------------------

cor(data.cohort$l2d,data.cohort$age.rx)
plot(data.cohort$l2d,data.cohort$age.rx)


#-----------------------
# Full cohort analysis
#-----------------------

# generate log2(dose+1) [log is base 2]
data.cohort$l2d<-log2(data.cohort$dose+1)

# 1. not using robust SEs

model.cohort<-coxph(Surv(age.in,age.out,bc)~l2d+I(age.rx/10),data=data.cohort)
summary(model.cohort)

# 2. using robust SEs

model.cohort.robust<-coxph(Surv(age.in,age.out,bc)~l2d+I(age.rx/10)+cluster(id),data=data.cohort)
summary(model.cohort.robust)


#----------------------
# Case-cohort analysis
#----------------------

# 1. Prentics's estimator with robust SEs (cf. sections 16.4.2 and 20.3)

# generate new entry time for cases outside the subcohort
data.caco$age.in.new<-ifelse(data.caco$bc==1 & data.caco$subcoh==0,data.caco$age.out-0.0001,data.caco$age.in)

# fit model
model.caco<-coxph(Surv(age.in.new,age.out,bc)~l2d+I(age.rx/10)+cluster(id),data=data.caco)
summary(model.caco)

#2. IPW estimator (cf. sections 16.4.3, 17.4.2, and 17.5.4) 

model.caco.cch<-cch(Surv(age.in,age.out,bc)~l2d+I(age.rx/10),data=data.caco,subcoh=~subcoh,id=~id,cohort.size=1720,method="LinYing",robust=T)
print(model.caco.cch)
summary(model.caco.cch)


#------------------------------
# Nested case-control analysis
#-----------------------------

# generate log2(dose+1) [log is base 2]
data.ncc$l2d<-log2(data.ncc$dose+1)

# fit model
model.ncc<-coxph(Surv(age.in,age.out,case)~l2d+I(age.rx/10)+strata(setno),data=data.ncc)
summary(model.ncc)


#-------------------------
# MI - case-cohort analysis
#-------------------------

# read data again because we have created some variables above that we don't need for this part
data.cohort<-read.table("tb.cohort.txt",header=T)
data.caco<-read.table("tb.cascoh.txt",header=T)

# merge case-cohort data with full cohort data
data.cohort.caco<-merge(data.cohort,data.caco[,c("id","l2d","subcoh","incc")],by="id",all.x=T,all.y=T)

# generate log2(1+nof/100) [log is base 2]
data.cohort.caco$l2nof<-log2(1+data.cohort.caco$nof/100) 

# generate Nelson-Aalen estimate of cumulative hazard at each event/censoring time and at each left-truncation time
nafit<-survfit(Surv(age.in,age.out,bc)~1,data=data.cohort.caco,type="fleming-harrington")
chaz<- -log(nafit$surv)
sorted.times<-sort(unique(c(data.cohort.caco$age.in,data.cohort.caco$age.out)))
chaz.time<-data.frame(cbind(sorted.times,chaz))

data.cohort.caco$chaz.age.in<-chaz.time$chaz[match(data.cohort.caco$age.in,chaz.time$sorted.times)] #cumhaz at each left-truncation time
data.cohort.caco$chaz.age.out<-chaz.time$chaz[match(data.cohort.caco$age.out,chaz.time$sorted.times)] #cumhaz at each event/censoring time  

data.cohort.caco$chaz.diff<-data.cohort.caco$chaz.age.out-data.cohort.caco$chaz.age.in #difference in cumulative hazards  
       

# 1. perform MI - using mice - 100 imputations

pred.mat<-matrix(0,nrow=dim(data.cohort.caco)[2],ncol=dim(data.cohort.caco)[2])
colnames(pred.mat)=names(data.cohort.caco)
rownames(pred.mat)=names(data.cohort.caco)
pred.mat["l2d",c("bc","age.rx","chaz.diff")]<-1

method=rep("",dim(data.cohort.caco)[2])
method[which(names(data.cohort.caco)=="l2d")]="norm"

set.seed(475987)
imp<-mice(data.cohort.caco, m = 100, method = method, 
          predictorMatrix = pred.mat, 
          maxit = 5, diagnostics = FALSE, printFlag = FALSE)

imp.combine<-with(imp,coxph(Surv(age.in,age.out,bc)~l2d+I(age.rx/10))) 

# pool estimates using MIcombine
betas=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$coefficients))
vars=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$var))
summary(MIcombine(betas,vars))

# pool estimates using pool()
summary(pool(imp.combine))


# 2. perform MI - using mice - 100 imputations - including 'nof' as auxilliary variable
pred.mat<-matrix(0,nrow=dim(data.cohort.caco)[2],ncol=dim(data.cohort.caco)[2])
colnames(pred.mat)=names(data.cohort.caco)
rownames(pred.mat)=names(data.cohort.caco)
pred.mat["l2d",c("bc","age.rx","chaz.diff","l2nof")]<-1

set.seed(475987)
imp<-mice(data.cohort.caco, m = 100, method = method, 
          predictorMatrix = pred.mat, 
          maxit = 5, diagnostics = FALSE, printFlag = FALSE)

imp.combine<-with(imp,coxph(Surv(age.in,age.out,bc)~l2d+I(age.rx/10))) 

betas=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$coefficients))
vars=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$var))
summary(MIcombine(betas,vars))

summary(pool(imp.combine))


#----------------------------------
# MI - nested case-control analysis
#----------------------------------

# read data again because we have created some variables above that we don't need for this part

data.cohort<-read.table("tb.cohort.txt",header=T)
data.ncc<-read.table("tb.ncc.txt",header=T)
data.ncc$l2d<-log2(data.ncc$dose+1)

# merge case-cohort data with full cohort data
data.cohort.ncc<-merge(data.cohort,data.ncc[,c("id","l2d","setno","case")],by="id",all.x=T,all.y=T)

# generate log2(1+nof/100) [log is base 2]
data.cohort.ncc$l2nof<-log2(1+data.cohort.ncc$nof/100) 

#drop duplicated individuals (this occurs because some individuals appear more than oncein the ncc data)
data.cohort.ncc<-data.cohort.ncc[duplicated(data.cohort.ncc$id)==F,]

# generate Nelson-Aalen estimate of cumulative hazard at each event/censoring time and at each left-truncation time
nafit<-survfit(Surv(age.in,age.out,bc)~1,data=data.cohort.ncc,type="fleming-harrington")
chaz<--log(nafit$surv)
sorted.times<-sort(unique(c(data.cohort.ncc$age.in,data.cohort.ncc$age.out)))
chaz.time<-data.frame(cbind(sorted.times,chaz))

data.cohort.ncc$chaz.age.in<-chaz.time$chaz[match(data.cohort.ncc$age.in,chaz.time$sorted.times)] #cumhaz at each left-truncation time
data.cohort.ncc$chaz.age.out<-chaz.time$chaz[match(data.cohort.ncc$age.out,chaz.time$sorted.times)] #cumhaz at each event/censoring time 
       
data.cohort.ncc$chaz.diff<-data.cohort.ncc$chaz.age.out-data.cohort.ncc$chaz.age.in #difference in cumulative hazards  

# 1. perform MI - using mice - 100 imputations
pred.mat<-matrix(0,nrow=dim(data.cohort.ncc)[2],ncol=dim(data.cohort.ncc)[2])
colnames(pred.mat)=names(data.cohort.ncc)
rownames(pred.mat)=names(data.cohort.ncc)
pred.mat["l2d",c("bc","age.rx","chaz.diff")]<-1

method=rep("",dim(data.cohort.ncc)[2])
method[which(names(data.cohort.ncc)=="l2d")]="norm"

set.seed(475987)
imp<-mice(data.cohort.ncc, m = 100, method = method, 
          predictorMatrix = pred.mat, 
          maxit = 5, diagnostics = FALSE, printFlag = FALSE)

imp.combine<-with(imp,coxph(Surv(age.in,age.out,bc)~l2d+I(age.rx/10))) 

betas=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$coefficients))
vars=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$var))
summary(MIcombine(betas,vars))

summary(pool(imp.combine))


# 2. perform MI - using mice - 100 imputations - including 'nof' as auxilliary variable
pred.mat<-matrix(0,nrow=dim(data.cohort.ncc)[2],ncol=dim(data.cohort.ncc)[2])
colnames(pred.mat)=names(data.cohort.ncc)
rownames(pred.mat)=names(data.cohort.ncc)
pred.mat["l2d",c("bc","age.rx","chaz.diff","l2nof")]<-1

set.seed(475987)
imp<-mice(data.cohort.ncc, m = 100, method = method, 
          predictorMatrix = pred.mat, 
          maxit = 5, diagnostics = T, printFlag = FALSE)

imp.combine<-with(imp,coxph(Surv(age.in,age.out,bc)~l2d+I(age.rx/10))) 

betas=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$coefficients))
vars=sapply(1:100,FUN=function(x)list(imp.combine$analyses[[x]]$var))
summary(MIcombine(betas,vars))

summary(pool(imp.combine))