# survplot.R
#
# 2009-02-20
# Aron Charles Eklund
# http://www.cbs.dtu.dk
#
# depends on "survival" package.
#

# FUNCTIONS IN THIS FILE ARE THE PART OF survplot PACKAGE (version 0.0.7)
# CREATED BY Aron Charles Eklund

nrisk <- function(x, times = pretty(x$time)) { 
  stopifnot(class(x) == 'survfit')
  if('strata' %in% names(x)) {
    ns <- length(x$strata)
    idx <- rep.int(1:ns, x$strata)
    str.n.risk <- split(x$n.risk, idx)
    str.times <- split(x$time, idx)
    m <- sapply(times, function(y) {
      sapply(1:ns, function(i) {
        w <- which(str.times[[i]] >= y)[1]
        ifelse(is.na(w), 0, str.n.risk[[i]][w])
      })
    })
    rownames(m) <- names(x$strata)
  } else {  # no strata
    m1 <- sapply(times, function(y) {
      w <- which(x$time >= y)[1]
      ifelse(is.na(w), 0, x$n.risk[w])
    })
    m <- matrix(m1, nrow = 1)
  }
  colnames(m) <- times
  m
}


addNrisk <- function(x, at = axTicks(1), 
                     line = 4, hadj = 0.5, 
                     title = 'Number at risk', title.adj = 0, 
                     labels, hoff = 5, col = 1) {
  m <- nrisk(x, times = at)
  ns <- nrow(m)  # number of strata
  if(missing(labels)) {
    if(ns > 1) { 
      labels <- names(x$strata)
    } else {
      labels <- NA
    }
  }    
  label.pad <- paste(rep(' ', hoff), collapse = '') 
  labels2 <- paste(labels, label.pad, sep = '')
  labels2[is.na(labels)] <- NA
  col <- rep(col, length.out = ns)
  hasTitle <- (!is.null(title)) && (!is.na(title))
  if(hasTitle) {
    title(xlab = title, line = line, adj = title.adj)
  }
  for (i in 1:ns) {
    axis(1, at = at, labels = m[i,], 
         line = line + i + hasTitle - 2, tick = FALSE, 
         col.axis = col[i], hadj = hadj)
    axis(1, at = par('usr')[1], labels = labels2[i],
         line = line + i + hasTitle - 2, tick = FALSE,
         col.axis = col[i], hadj = 1)
  }
  invisible(m)
}



survplot <- function(x, data = NULL, subset = NULL, 
                     snames, stitle, 
                     col, lty, lwd,
                     show.nrisk = TRUE, color.nrisk = TRUE,
                     hr.pos = 'topright', legend.pos = 'bottomleft', ...) {
  eval(bquote(s <- survfit(x, data = data, subset = .(substitute(subset)))))
  if('strata' %in% names(s)) {
    if(missing(stitle)) stitle <- strsplit(deparse(x), " ~ ")[[1]][2]
    if(missing(snames)) {
      snames <- names(s$strata)
      prefx <- paste(strsplit(deparse(x), " ~ ")[[1]][2], '=', sep = '')
      if(all(substr(snames, 1, nchar(prefx)) == prefx)) {
        snames <- substr(snames, nchar(prefx) + 1, 100)
      }
    }
    ns <- length(s$strata)
    stopifnot(length(snames) == ns)
  } else {    # no strata
    ns <- 1
    snames <- NA
    legend.pos <- NA
  }
  if(show.nrisk) {
    mar <- par('mar')
    mar[1] <- mar[1] + ns + 0.5
    opar <- par(mar = mar)
    on.exit(par(opar))
  }
  if(missing(col)) col <- 1:ns
  if(missing(lty)) lty <- 1
  if(missing(lwd)) lwd <- par('lwd')
  plot(s, col = col, lty = lty, lwd = lwd, ...)
  if(length(legend.pos) > 1 || !is.na(legend.pos)) {
    legend(legend.pos, legend = snames, title = stitle,
           col = col, lty = lty, lwd = lwd, bty = 'n')
  }
  if(ns == 2) {
    eval(bquote(cox <- summary(coxph(x, data = data, subset = .(substitute(subset))))))
    hr <- format(cox$conf.int[1, c(1, 3, 4)], digits = 2)
    p <- format(cox$sctest[3], digits = 2)
    txt1 <- paste('HR = ', hr[1], ' (', hr[2], ' - ', hr[3], ')', sep = '')
    txt2 <- paste('logrank P =', p)
    legend(hr.pos, legend = c(txt1, txt2), bty = 'n')
  }
  if(show.nrisk) {
    addNrisk(s, labels = snames,
             col = if(color.nrisk) col else 1)
  }
  if(ns == 2) return(invisible(c(txt1, txt2)))
}