# km.r - Kaplan-Meier survival analysis wrapper
# USAGE
# time   numeric, follow up time, i.e. event times or censored times.
# event  numeric or logical, censoring indicator, 0 or FALSE for events
#        NOTE: This is the OPPOSITE of what Surv() expects (0=alive, 1=dead)
# grp    factor, variable defining groups to be compared.
# tit    text, main title for graph and output.
# shorttit text, shortened title for y-axis.
# jpg    logical, TRUE=save plot as JPEG file, FALSE=do not save (default)
# fname  name of JPEG file
# colv   vector of colors to be used in plot. 
# timev  vector of times at which to compute estimated survival for each group,
# ...    optional graphical parameters to be passed on to plot()

   library(survival)
   data(lung)
   # km(nn$PATYEARS, nn$PATCEN, nn$SBNN, "Patient Survival", "death")
   km <- function(time, event, grp, tit="Kaplan-Meier Survival", shorttit="death",
                  jpg=FALSE, fname, colv=NULL, timev=sort(c(0:5, 0.25, 0.5)), ...) {
     levg <- levels(grp)
     ng <- length(levg)
     if (is.null(colv)) colv <- 1:ng+1
     mktitle(paste(tit))
     so <- Surv(time, event==0)
     print(km <- survfit(so ~ grp))
     print(kmt <- survdiff(so ~ grp))
     print(summary(km, times=timev))
     if (jpg) jpeg(fname, quality=99, width=640)
     plot(km, mark.time=TRUE, las=1, xlab="time (years)",
          ylab=paste("fraction free of", shorttit),
          main=paste(tit, " (P=", format(1-pchisq(kmt$chisq, ng-1), digits=3),
            ")", sep=""), col=colv,...)
     legend(0.5, 0.45, lty=1, col=colv, legend=levg, bty="n", cex=0.8)
     if (jpg) dev.off()
     # pair-wise comparisons
     if (ng > 2) {
        mktitle("Pair-wise Comparisons")
        for (i in 1:(ng-1)) {
           for (j in (i+1):ng) {
              kmt2 <- survdiff(so ~ grp, subset = (grp %in% levg[c(i,j)]))
              cat(paste(levg[i], "vs.", levg[j]), ": ",
                  format(1-pchisq(kmt2$chisq, 1), digits=3), "(Chisq =",
                  format(kmt2$chisq, digits=3), "on 1 d.f.)\n")
           }
        }
     }
  }

# mktitle - enclose a string within a box, with optional page break
mktitle <- function(string, pgbrk=FALSE, dt=FALSE, boxch="-") {
  strlen <- function(string) length(unlist(strsplit(string, "")))
  if (!is.character(string))
    stop("mktitle: string must be of character class.\n")
  lstr <- strlen(string)
  if (pgbrk) cat("")
  if (boxch=="-") {
    topbot <- paste("+-", paste(rep("-", lstr), collapse=""), "-+\n", sep="")
    mid <- paste("| ", string, " |\n", sep="", collapse="")
  } else {
    topbot <- paste(paste(rep(boxch, lstr+4), collapse=""), "\n", sep="")
    mid <- paste(boxch, string, boxch, "\n", collapse="")
  }
  if (dt) {
    date.txt <- if(dt) format(Sys.time()) else format(Sys.time(), '%Y-%m-%d')
    outw <- options()$width
    ldate <- strlen(date.txt)
    date.txt <- paste(paste(rep(" ", outw-(ldate+3)), collapse=""),
                      date.txt, "\n", sep="")
    cat(date.txt)
  }
  cat(topbot)
  cat(mid)
  cat(topbot)
  cat("\n")
  invisible()
  # return(NULL)
}
km(vas$survyrs, vas$status=="alive", vas$surgtype, colv=2:3)