rm(list=ls())
library(survival)


PRIB1=function(PFS, STATUS_PFS, ETIME, STATUS_TTR, time.max){
  taumax=time.max
  ETIME[STATUS_TTR==0]=Inf
  
  y1=pmin(PFS, ETIME)
  d1=1*(STATUS_TTR+STATUS_PFS>0)
  
  y2=PFS
  d2=STATUS_PFS
  
  fit1=survfit(Surv(y1, d1)~1)
  n1=length(fit1$time)
  
  fit2=survfit(Surv(y2, d2)~1)
  n2=length(fit2$time)
  
  tau.grd=sort(unique(c(y1[d1==1], y2[d2==1])))
  tau.grd=c(tau.grd[tau.grd<taumax], taumax)
  m.tau.grd=length(tau.grd)
  n=length(PFS)
  
  surv1.tot=surv2.tot=rep(NA, m.tau.grd)
  atrisk1.tot=atrisk2.tot=rep(NA, m.tau.grd)
  hazard1.tot=hazard2.tot=rep(NA, m.tau.grd)
  
  for(i in 1:m.tau.grd)
  {t0=tau.grd[i]
  id1=max((1:(n1+1))[c(0, fit1$time)<=t0])
  id2=max((1:(n2+1))[c(0, fit2$time)<=t0])
  
  surv1.tot[i]=c(1, fit1$surv)[id1]
  surv2.tot[i]=c(1, fit2$surv)[id2]
  
  atrisk1.tot[i]=sum(y1>=t0)
  atrisk2.tot[i]=sum(y2>=t0)
  
  hazard1.tot[i]=sum((y1==t0)*d1)/sum(y1>=t0)
  hazard2.tot[i]=sum((y2==t0)*d2)/sum(y2>=t0)
  }
  dsurv.tot=surv2.tot-surv1.tot
  
  
  dsd.tot=rep(NA, m.tau.grd)
  for(i in 1:m.tau.grd)
  {t0=tau.grd[i]
  tau1=tau2=rep(0, n)
  for(j in 1:n)
    tau1[j]=(y1[j]<=t0)*d1[j]/sum(y1>=y1[j])-sum((hazard1.tot/atrisk1.tot)[tau.grd<=min(t0,y1[j])])
  for(j in 1:n)
    tau2[j]=(y2[j]<=t0)*d2[j]/sum(y2>=y2[j])-sum((hazard2.tot/atrisk2.tot)[tau.grd<=min(t0,y2[j])])
  tau1=surv1.tot[i]*tau1
  tau2=surv2.tot[i]*tau2
  dsd.tot[i]=sqrt(sum((tau1-tau2)^2))  
  }           
  
  logd.tot=log(dsurv.tot/(1-dsurv.tot))
  logsd.tot=dsd.tot/(dsurv.tot*(1-dsurv.tot))
  
  cilogd=cbind(logd.tot-1.96*logsd.tot, logd.tot+1.96*logsd.tot)
  cid=exp(cilogd)/(1+exp(cilogd))
  
  res.tot=cbind(tau.grd, dsurv.tot, dsd.tot, cid)
  colnames(res.tot)=c("time", "PRIB", "std", "ci-low", "ci-up")
  
  res.tot=data.frame(res.tot)
  
  return(res.tot)
}

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


PRIB2=function(PFS, STATUS_PFS, ETIME, STATUS_TTR, TRT, time.max){
  
  res1=PRIB1(PFS[TRT==1], STATUS_PFS[TRT==1], ETIME[TRT==1], STATUS_TTR[TRT==1], time.max)
  res0=PRIB1(PFS[TRT==0], STATUS_PFS[TRT==0], ETIME[TRT==0], STATUS_TTR[TRT==0], time.max)
  
  n1=length(res1[,1])
  n0=length(res0[,1])
  
  time0=sort(unique(c(res1$time, res0$time)))
  n.time0=length(time0)
  
  diff=rep(NA, n.time0)
  sd.diff=rep(NA, n.time0)
  ci1.diff=rep(NA, n.time0)
  ci2.diff=rep(NA, n.time0)
  
  for(i in 1:n.time0)
  {id1=max((1:(n1+1))[c(0, res1$time)<=time0[i]])
   id0=max((1:(n0+1))[c(0, res0$time)<=time0[i]])
   diff[i]=c(0,res1$PRIB)[id1]-c(0, res0$PRIB)[id0]
   sd.diff[i]=sqrt(c(0,res1$std)[id1]^2+c(0, res0$std)[id0]^2)
  }
  
  diff.tr=log((1-diff)/(1+diff))
  sd.difftr=2*sd.diff/(1-diff^2)  
  
  ci.tr=cbind(diff.tr-1.96*sd.difftr, diff.tr+1.96*sd.difftr)
  ci=(1-exp(ci.tr))/(1+exp(ci.tr))
  
  res=cbind(time0, diff, sd.diff, ci[,2], ci[,1])
  colnames(res)=c("time", "diff in PRIB", "std", "ci-low", "ci-up")
  
  res=data.frame(res)
  
  return(res)
}

###########################################################################################
##### INPUT ###############################################################################

# ETIME: time to response (month)
# STATUS_TTR: status for response (=1 for responders; =0 for non-responders)
# PFS: progression-free survival time (month)
# STATUS_PFS: 1=event; 0=censored
# TRT:  treatment indicator
# time.max:  truncation time for defining the maximum time point of interest
set.seed(123)
n=200
TRT=trt=rbinom(n, 1, 0.5)

error=rnorm(n)
tr=exp(rnorm(n)+error-trt*0.5+0.5)
tp=exp(rnorm(n)+error+trt*0.25)
tr[tp<tr]=Inf
tc=runif(n, 3, 8.5)


t2response=pmin(tr, tc)
delta_response=1*(tr<tc)
t2progression=pmin(tp, tc)
delta_progression=1*(tp<tc)

par(mfrow=c(1,2))

fit=PRIB2(PFS=t2progression, STATUS_PFS=delta_progression, ETIME=t2response, STATUS_TTR=delta_response, TRT=trt, time.max=8)

plot(range(fit$time), range(c(fit$ci.low, fit$ci.up)), xlab="time", ylab="difference in PRIB", lwd=2, type="n")
lines(fit$time, fit$diff, lwd=2)
lines(fit$time, fit$ci.low, col=2)
lines(fit$time, fit$ci.up, col=2)

fit1=PRIB1(PFS=t2progression[trt==1], STATUS_PFS=delta_progression[trt==1], t2response[trt==1], delta_response[trt==1], time.max=8)
fit0=PRIB1(PFS=t2progression[trt==0], STATUS_PFS=delta_progression[trt==0], t2response[trt==0], delta_response[trt==0], time.max=8)


plot(fit1$time, fit1$PRIB, type="l")
lines(fit0$time, fit0$PRIB, col=2)