library(survival)
source("http://myweb.uiowa.edu/pbreheny/7210/f19/notes/fun.R")

# Basic KM
fit <- survfit(Surv(time/365.25, status) ~ trt, data=veteran)
plot(fit, mark.time=FALSE, col=pal(2), lwd=3, las=1, bty='n', xlab='Time (years)', ylab='Survival')
toplegend(legend=c('Standard', 'Test'), lwd=3, col=pal(2))

# Basic Cox
fit <- coxph(Surv(time/365.25, status) ~ trt + karno + diagtime + age + prior + celltype, data=veteran)
summary(fit)

# Treatment, v1
fit <- coxph(Surv(time/365.25, status) ~ strata(trt) + karno + diagtime + age + prior + celltype, data=veteran)
sfit <- survfit(fit)
plot(sfit, mark.time=FALSE, col=pal(2), lwd=3, fun=function(x) log(-log(x)), xmax=1.5, bty='n', las=1,
     xlab='Time (years)', ylab=expression(log(Lambda(t))))
toplegend(legend=c('Standard', 'Test'), lwd=3, col=pal(2))

# Treatment, v2
fit0 <- coxph(Surv(time/365.25, status) ~ trt + karno + diagtime + age + prior + celltype, data=veteran)
cumHaz <- function(sfit, t) {
  K <- length(sfit$strata)
  s <- c(0, cumsum(sfit$strata))
  H <- matrix(NA, nrow=length(t), ncol=K)
  for (i in 1:K) {
    ind1 <- (s[i] + 1):s[i + 1]
    H[,i] <- approxfun(c(0,sfit$time[ind1]), c(0,sfit$cumhaz[ind1]), method='constant')(t)
  }
  H
}
Time <- seq(0, 1.5, len=99)
plot(Time, apply(log(cumHaz(sfit, Time)), 1, diff), type='s', bty='n', las=1, lwd=2,
     xlab="Time (years)", ylab=expression(log*Lambda[01](t)-log*Lambda[02](t)))
abline(h=coef(fit0)[1], col='red', lwd=2)

# Treatment, v3
plot(cumHaz(sfit, Time), type='s', bty='n', las=1, lwd=2,
     xlab=expression(Lambda[1](t)), ylab=expression(Lambda[2](t)))
abline(0, exp(coef(fit0)[1]), col='red', lwd=2)

# Cell type
fit <- coxph(Surv(time/365.25, status) ~ trt + karno + diagtime + age + prior + strata(celltype), data=veteran)
sfit <- survfit(fit)
plot(sfit, mark.time=FALSE, col=pal(4), lwd=3, fun=function(x) log(-log(x)), xmax=1.5, bty='n', las=1,
     xlab='Time (years)', ylab=expression(log(Lambda(t))))
toplegend(legend=levels(veteran$celltype), lwd=3, col=pal(4))

# Cell type, v2
Time <- seq(0, 1.5, len=99)
LCH <- log(cumHaz(sfit, Time))
matplot(Time, LCH[,1:3]-LCH[,4], type='s', bty='n', col=pal(4)[1:3], lwd=2, lty=1, las=1,
        xlab="Time (years)", ylab=expression(log*Lambda(t)-log*Lambda[large](t)))
toplegend(legend=levels(veteran$celltype)[1:3], lwd=3, col=pal(4)[1:3])

# Karno
veteran$cKarno <- cut(veteran$karno, c(0, 30, 60, 100))
fit <- coxph(Surv(time/365.25, status) ~ trt + strata(cKarno) + diagtime + age + prior + celltype, data=veteran)
sfit <- survfit(fit)
plot(sfit, mark.time=FALSE, col=pal(3), lwd=3, fun=function(x) log(-log(x)), xmax=1.5, bty='n', las=1,
     xlab='Time (years)', ylab=expression(log(Lambda(t))))
toplegend(legend=levels(veteran$cKarno), lwd=3, col=pal(3))

# Karno, v2
Time <- seq(0, 1.5, len=99)
LCH <- log(cumHaz(sfit, Time))
matplot(Time, LCH[,1:2]-LCH[,3], type='s', bty='n', col=pal(3)[1:2], lwd=2, lty=1, las=1,
        xlab="Time (years)", ylab=expression(log*Lambda(t)-log*Lambda[60-100](t)))
toplegend(legend=levels(veteran$cKarno)[1:2], lwd=3, col=pal(3)[1:2])

# Prediction
fit <- coxph(Surv(time/365.25, status) ~ trt + karno + diagtime + age + prior + strata(celltype), data=veteran)
newdata <- data.frame(trt=1, celltype=c('squamous', 'squamous', 'large'), karno=c(80, 60, 60), diagtime=12, age=40, prior=0)
pred <- survfit(fit, newdata=newdata)
plot(pred, col=pal(3), mark.time=FALSE, xmax=1.5, lwd=3, bty='n', las=1,
     xlab='Time (years)', ylab='Survival')
leg <- paste(substr(newdata$celltype, 1, 3), ': krn=', newdata$karno, sep='')
toplegend(legend=leg, col=pal(3), lwd=3)