# ===========================================================
# C.2 Code for the onset of arthritis mutilans in PsA
# R code in C.2.3 Inverse intensity weighted nonparametric 
# estimation
#
# December 11, 2018
# ===========================================================


library(parallel)
library(survival)


iivw.nonpara.f <- function(intt, indata, r, cores) {
  indata$egap <- round(indata$egap, 10)

  fitWGT.GAP1 <- coxph(Surv(egap, status) ~ sex.female + age.psa + b27 + c3,
                       data=indata[indata$enum == 1,], method="breslow", x=TRUE)

  fitWGT.GAP <- coxph(Surv(egap, status) ~ sex.female + age.psa + b27 + c3 + factor(from),
                      data=indata[indata$enum > 1,], method="breslow", x=TRUE)

  pid <- sort(unique(indata$id))

  wFt <- Ft <- NULL
  for (tt in intt) {
    pp <- mclapply(pid, function(id, inr, intt, indata, infit1, infit) {
      datai <- indata[indata$id == id,]
      datai <- datai[order(datai$estart, datai$estop),]

      nlen <- length(datai$id)

      Iij <- ifelse((datai$estop >  (intt - 0.5)) &
                    (datai$estop <= (intt + 0.5)), 1, 0)
      Yir <- ifelse(datai$to == inr, 1, 0)

      if ( sum(Iij) == 0 ) {
        top <- bot <- wtop <- wbot <- 0
      }
      else {
        top <- bot <- wtop <- wbot <- 0

        first.Yir <- 1
        for (j in 1:nlen) {
          if (Iij[j] != 0) {
            if (datai$enum[j] == 1) {
              km <- survfit(infit1, newdata=data.frame(sex.female=datai$sex.female[j],
                                                       age.psa=datai$age.psa[j],
                                                       b27=datai$b27[j],
                                                       c3=datai$c3[j]),
                            type="aalen")
            }
            else {
              km <- survfit(infit, newdata=data.frame(sex.female=datai$sex.female[j],
                                                      age.psa=datai$age.psa[j],
                                                      b27=datai$b27[j],
                                                      c3=datai$c3[j],
                                                      from=datai$from[j]),
                            type="aalen")
            }
            km <- data.frame(tt=c(0, km$time), St=c(1, km$surv))

            lo.tt <- max( c(0, intt - 0.5 - datai$estart[j]) )
            up.tt <- intt + 0.5 - datai$estart[j]
    
            lo <- km$St[sum(km$tt <= lo.tt)]
            up <- km$St[sum(km$tt <= up.tt)]
      
            Wij <- Iij[j]/(lo - up)
            Wij <- ifelse((lo - up) == 0, 0, Wij)

            bot  <- bot + Iij[j]
            wbot <- wbot + Wij
      
            top <- top + (Iij[j]*Yir[j])
            wtop <- wtop + (Wij*Yir[j])
          }
        }
      }
      return( data.frame(top, bot, wtop, wbot) )
    }, inr=r, intt=tt, indata=indata[indata$status == 1,],
       infit1=fitWGT.GAP1, infit=fitWGT.GAP, mc.cores=cores)
    pp <- do.call("rbind", pp)
            
    Ft  <- c(Ft,  sum(pp$top)/sum(pp$bot))
    wFt <- c(wFt, sum(pp$wtop)/sum(pp$wbot))
  }
                                   
  outdata <- data.frame(tt=intt, Ft, wFt)
  return(outdata)
}


mutilans <- read.table("mutilans.dat", header=T)

mutilansMS <- mutilans
mutilansMS$estop  <- mutilansMS$times
mutilansMS$estart <- c(NA, mutilansMS$times[-nrow(mutilansMS)])
mutilansMS$to     <- mutilansMS$state
mutilansMS$from   <- c(NA, mutilansMS$state[-nrow(mutilansMS)])
mutilansMS <- mutilansMS[mutilansMS$enum > 1,]
mutilansMS$enum <- mutilansMS$enum - 1
mutilansMS$egap <- mutilansMS$estop - mutilansMS$estart
mutilansMS$from <- mutilansMS$from - 1
mutilansMS$to   <- mutilansMS$to - 1
mutilansMS <- mutilansMS[,c("id","sex.female","age.psa","b27","c3",
                            "enum","from","to","estart","estop","egap","status")]
mutilansMS[mutilansMS$id == 9,]

# id sex.female age.psa b27 c3 enum from  to estart  estop   egap status
#  9          1      26   0  0    1    0   1  0.000  9.489  9.489      1
#  9          1      26   0  0    2    1   1  9.489 11.488  1.999      1
#  9          1      26   0  0    3    1   1 11.488 13.528  2.040      1
#  9          1      26   0  0    4    1   3 13.528 20.627  7.099      1
#  9          1      26   0  0    5    3   5 20.627 23.072  2.445      1
#  9          1      26   0  0    6    5 998 23.072 45.024 21.952      0


nonpara <- iivw.nonpara.f(intt=seq(1,30,by=1), indata=mutilansMS, r=5, cores=1)

nonpara.isoreg <- isoreg(x=c(0,nonpara$tt), y=c(0,nonpara$wFt))

plot(nonpara.isoreg$x, nonpara.isoreg$yf, type="s")