#' Plot method for a MarkovTest object
#' 
#' Plot method for an object of class 'MarkovTest'. It plots the trace of the
#' log-rank statistics provided by \code{\link{MarkovTest}}.
#' 
#' 
#' @param x Object of class 'MarkovTest'
#' @param y The grid at which \code{MarkovTest} was calculated
#' @param what Choose "states" for plotting state-specific traces, and
#' "overall" for the overall chi-squared trace
#' @param idx Vector of indices of wild bootstrap traces to plot
#' @param quantiles Boolean whether or not to plot the 2.5 and 97.5 percent
#' quantiles, default is \code{TRUE}
#' @param qsup The index of the function in either \code{fn} (when plotting
#' state-specific) or \code{fn2} (when plotting overall) to plot along with the
#' traces; when missing this line is not included
#' @param states Number of the qualifying state(s) to plot trace for
#' @param xlab Text for x-axis label
#' @param ylab Text for y-axis label
#' @param main Text for title (main)
#' @param \dots Further arguments to plot
#' @return No return value
#' @author Hein Putter \email{H.Putter@@lumc.nl}
#' @seealso \code{\link{MarkovTest}}
#' @keywords hplot
#' @examples
#' 
#' \dontrun{
#' # Example provided by the prothrombin data
#' data("prothr")
#' # Apply Markov test to grid of monthly time points over the first 7.5 years
#' year <- 365.25
#' month <- year / 12
#' grid <- month * (1:90)
#' # Markov test for transition 1 (wild bootstrap based on 100 replications)
#' MT <- MarkovTest(prothr, id = "id", transition = 1,
#'                  grid = grid, B = 100)
#' 
#' plot(MT, grid, what="states", idx=1:50, states=rownames(attr(prothr, "trans")),
#'      xlab="Days since randomisation", ylab="Log-rank test statistic",
#'      main="Transition Normal -> Low")
#' 
#' plot(MT, grid,what="overall", idx=1:50,
#'      xlab="Days since randomisation", ylab="Chi-square test statistic",
#'      main="Transition Normal -> Low")
#' 
#' plot(MT, grid, what="states", quantiles=FALSE) # only trace
#' plot(MT, grid, what="states") # trace plus quantiles (default)
#' plot(MT, grid, what="states", idx=1:10) # trace plus quantiles, plus first 10 bootstrap traces
#'  
#' plot(MT, grid, what="overall", quantiles=FALSE) # only trace
#' plot(MT, grid, what="overall") # trace plus quantiles (default)
#' plot(MT, grid, what="overall", idx=1:10) # trace plus quantiles, plus first 10 bootstrap traces
#' 
#' }
#' 
#' @export
plot.MarkovTest <- function(x, y, what=c("states", "overall"), idx=NULL, quantiles=TRUE, qsup, states,
                            xlab, ylab, main, ...)
{
  ct <- NULL # To remove NOTE generated by R CMD CHECK about no visible binding
  # Not ready for publication, see definition of q95
  what <- match.arg(what)
  B <- dim(x$n_wb_trace)[1]
  ny <- length(y)
  if (missing(xlab)) xlab <- "Time"
  if (missing(ylab)) ylab <- "Test statistic"
  if (missing(main)) main <- ""
  if (what=="states") {
    qualset <- x$qualset
    J <- length(qualset)
    dfr <- data.frame(time=rep(y, J), zbar=as.numeric(x$zbar), qualstate=rep(qualset, each=ny), ct=0)
    lwd <- 2
    lty <- 1
    col <- 1
    if (quantiles) {
      dfrl1 <- data.frame(time=rep(y, J), zbar=as.numeric(x$est_quant[1, , ]), qualstate=rep(qualset, each=ny), ct=1)
      dfru1 <- data.frame(time=rep(y, J), zbar=as.numeric(x$est_quant[2, , ]), qualstate=rep(qualset, each=ny), ct=3)
      dfr <- rbind(dfr, dfrl1, dfru1)
      lwd <- c(lwd, 2, 2)
      lty <- c(lty, 3, 3)
      col <- c(col, 1, 1)
    }
    if (!missing(qsup)) {
      if (qsup %in% 1:dim(x$b_stat_wb)[2]) {
        q95 <- apply(x$b_stat_wb[, qsup, ], 2, quantile, 0.95)
        dfrl2 <- data.frame(time=rep(y, J), zbar=rep(-q95, each=ny), qualstate=rep(qualset, each=ny), ct=2)
        dfru2 <- data.frame(time=rep(y, J), zbar=rep(q95, each=ny), qualstate=rep(qualset, each=ny), ct=4)
        dfr <- rbind(dfr, dfrl2, dfru2)
        lwd <- c(lwd, 2, 2)
        lty <- c(lty, 3, 3)
        col <- c(col, 1, 1)
      }
    }
    if (!is.null(idx)) {
      idx <- intersect(1:B, idx)
      nB <- length(idx)
      if (nB > 0) {
        dfrb <- data.frame(time=rep(rep(y, J), each=nB),
                           zbar=as.numeric(x$n_wb_trace[idx, , ]),
                           qualstate=rep(qualset, each=ny*nB),
                           ct=rep(-idx, ny*J))
        dfr <- rbind(dfrb, dfr)
        lwd <- c(rep(0.5, nB), lwd)
        lty <- c(rep(1, nB), lty)
        col <- c(rep(8, nB), col)
      }
    }
    # print(dim(dfr))
    if (missing(states)) dfr$qualstate <- factor(dfr$qualstate)
    else dfr$qualstate <- factor(dfr$qualstate, levels=qualset, labels=states[qualset])
    xyplot(zbar ~ time | qualstate, data=dfr, groups=ct, lwd=lwd, type="l", col=col, lty=lty,
           xlab=xlab, ylab=ylab, main=main)
  }
  else if (what=="overall") {
    dfr <- data.frame(time=y, zbar=as.numeric(x$obs_chisq_trace), ct=0)
    lwd <- 2
    lty <- 1
    col <- 1
    if (quantiles) {
      
      dfru <- data.frame(time=y, zbar=apply(x$nch_wb_trace, 2, quantile, probs=0.95), ct=-1)
      dfr <- rbind(dfr, dfru)
      lwd <- c(2, lwd)
      lty <- c(3, lty)
      col <- c(1, col)
    }
    if (!is.null(idx)) {
      idx <- intersect(1:B, idx)
      nB <- length(idx)
      if (nB > 0) {
        dfrb <- data.frame(time=rep(y, each=nB),
                           zbar=as.numeric(x$nch_wb_trace[idx, ]),
                           ct=rep(idx, ny))
        dfr <- rbind(dfrb, dfr)
        lwd <- c(lwd, rep(0.5, nB))
        lty <- c(lty, rep(1, nB))
        col <- c(col, rep(8, nB))
      }
    }
    xyplot(zbar ~ time, data=dfr, groups=ct, lwd=lwd, type="l", col=col, lty=lty,
           xlab=xlab, ylab=ylab, main=main)
  }
}

