#' Create a basic survival object for doubly truncated data.
#'
#' @description
#' Creates a survival object to contain the event times, left truncation times,
#' and right truncation times from a doubly truncated sample, with one observation
#' per subject.
#'
#' Typically used as the input for nonparametric analysis methods or as the
#' response variable in a regression model formula.
#'
#' @param time A vector of follow up times.
#' @param ltrunc A vector of left truncation times.
#' @param rtrunc A vector of right truncation times.
#' @param x Any R object.
#' @param i,j Row and columns index vectors. The result is a `Survdt` object if
#'   `j` is missing, otherwise `x` is treated as a matrix.
#' @param sort Whether to sort the data rows by follow up time (and `colorvar`,
#'   if provided) before plotting (default `TRUE`).
#' @param colorvar An optional vector used to color the dots in the plot. Will
#'   be coerced to factor before using.
#' @param ... Additional arguments passed to other methods.
#'
#' @details
#' For a given event time of interest \eqn{T}, a sample is doubly truncated when
#' it only includes subjects with \eqn{L_i<T_i \leq R_i} for some random
#' truncation times \eqn{L_i} and \eqn{R_i}.
#' Analyzing this type of data requires information on the \eqn{i}th subject's
#' event time \eqn{T_i}, left truncation time \eqn{L_i}, and right truncation
#' time \eqn{R_i}.
#'
#' [Survdt()] provides the basic data structure for organizing such data and
#' checking its validity in the `survdt` package.
#' Attempting to run [Survdt()] with inputs that do not satisfy \eqn{L<T \leq R}
#' will result in an error.
#'
#' @returns
#' * [Survdt()] returns an object of class `Survdt`, with columns `time`,
#'   `ltrunc`, and `rtrunc`.
#'   Indexing a `Survdt` object by rows will return another `Survdt` object,
#'   otherwise it is treated as a matrix.
#'   The length of a `Survdt` object is its sample size, and `is.na()` returns
#'   a logical index vector for the non-complete cases.
#'   Calling `as.data.frame()` on a `Survdt` object will return a data frame
#'   stripped of the `Survdt` class attribute.
#' * [is.Survdt()] returns `TRUE` if `x` inherits from the class `Survdt` and
#'   `FALSE` otherwise.
#' * The plot method for `Survdt` objects returns a simple visualization from
#'   [ggplot2::ggplot()] with x-axis time, y-axis subject, rectangles for the
#'   truncation intervals, and dots for the event times.
#'   The dots can optionally be colored by a factor variable supplied in
#'   `colorvar`.
#'   Incomplete cases are dropped from the plot (with a warning).
#'
#' @export
#'
#'
#' @seealso
#' * [npmle()], [test_samplediff_indtrunc()]: nonparametric estimation and
#'   analysis for doubly truncated data.
#' * [coxph_indtrunc()]: inverse probability weighted Cox regression.
#' * [Survdt2()]: to encode time-varying covariates; also [time_split()] for
#'   converting data from `Survdt` structure to `Survdt2`.
#'
#' @examples
#' y <- with(aids, Survdt(incu, ltrunc, rtrunc))
#' plot(y)
#' plot(y, colorvar = aids$age_gp)
#' npmle(y)
#'
Survdt <- function(time, ltrunc, rtrunc){
  # check input
  if (inherits(time, "difftime")) time <- as.numeric(time)
  if (inherits(ltrunc, "difftime")) ltrunc <- as.numeric(ltrunc)
  if (inherits(rtrunc, "difftime")) rtrunc <- as.numeric(rtrunc)

  stopifnot(is.numeric(time), is.vector(time), !any(is.infinite(time)),
            is.numeric(ltrunc), is.vector(ltrunc), !any(is.infinite(ltrunc)),
            is.numeric(rtrunc), is.vector(rtrunc), !any(is.infinite(rtrunc)))
  stopifnot(length(time) > 0, length(time) == length(ltrunc),
            length(time) == length(rtrunc))

  na_ltrunc <- is.na(ltrunc)
  na_rtrunc <- is.na(rtrunc)
  na_trunc <- na_ltrunc | na_rtrunc
  na_time <- is.na(time)
  if (any(!na_trunc) && any((ltrunc >= rtrunc)[!na_trunc]))
    stop("ltrunc must by < rtrunc")
  if (any(!na_ltrunc & !na_time) && any((ltrunc >= time)[!na_ltrunc & !na_time]))
    stop("Each event time must be >ltrunc")
  if (any(!na_rtrunc & !na_time) && any((time > rtrunc)[!na_rtrunc & !na_time]))
    stop("Each event time must be <=rtrunc")

  x <- cbind(time = time, ltrunc = ltrunc, rtrunc = rtrunc)
  class(x) <- "Survdt"
  x
}

#' @rdname Survdt
#' @export
`[.Survdt` <- function(x, i, j, ...){
  if (missing(j) && !missing(i)){
    x <- unclass(x)[i, , drop = FALSE]
    class(x) <- "Survdt"
    x
  }else{
    class(x) <- "matrix"
    NextMethod("[")
  }
}

#' @rdname Survdt
#' @export
is.na.Survdt <- function(x){
  as.vector(rowSums(is.na(unclass(x))) > 0)
}

#' @rdname Survdt
#' @export
is.Survdt <- function(x){
  inherits(x, "Survdt")
}

#' @rdname Survdt
#' @export
length.Survdt <- function(x){
  nrow(x)
}

#' @rdname Survdt
#' @export
as.data.frame.Survdt <- function(x, ...){
  as.data.frame(unclass(x), ...)
}

#' @rdname Survdt
#' @export
plot.Survdt <- function(x, colorvar = NULL, sort = TRUE, ...){
  stopifnot(is.logical(sort), is.vector(sort), all(!is.na(sort)))
  stopifnot(length(sort) == 1)
  hasc <- !is.null(colorvar)
  if (hasc) {
    colorlabel <- format(substitute(colorvar))
    colorvar <- as.factor(colorvar)
    stopifnot(length(colorvar) == length(x))
  }else
    colorlabel <- ""

  n <- length(x)
  cc_indx <- !is.na(x)
  x <- x[cc_indx, ]
  n_miss <- n - length(x)
  if (n_miss > 0){
    if (n_miss == n){
      stop("No complete observations available to plot.")
    }
    warning(paste0("omitted ", n_miss, " observations with missing values."))
  }
  if (hasc)
    colorvar <- colorvar[cc_indx]
  else
    colorvar <- rep(1, sum(cc_indx))


  if (sort == TRUE){
    sort_indx <- order(colorvar, x[, "time", drop = TRUE])

  }else{
    sort_indx <- seq_along(x[, "time", drop = TRUE])
  }

  plot_data <- cbind(rownum = seq_along(x[, "time", drop = TRUE]),
                     colorvar = colorvar[sort_indx],
                     as.data.frame(x)[sort_indx, , drop = FALSE])
  plt <- with(plot_data,
              ggplot2::ggplot(plot_data,
                              ggplot2::aes(x = time, y = rownum - 0.5)) +
                ggplot2::geom_rect(ggplot2::aes(xmin = ltrunc, xmax = rtrunc,
                                                ymin = rownum-1, ymax = rownum),
                                   fill = "gray"))
  if (hasc)
    plt <- with(plot_data,
                plt +
                  ggplot2::geom_point(ggplot2::aes(color = colorvar)) +
                  ggplot2::scale_color_viridis_d(begin = 0, end = 0.8, guide =
                          ggplot2::guide_legend(override.aes = list(size = 3))))
  else
    plt <- with(plot_data,
                plt +
                  ggplot2::geom_point())

  with(plot_data,
       plt +
         ggplot2::labs(x = "Time", y = "Subject", color = colorlabel) +
         ggplot2::coord_cartesian(expand = FALSE) +
         ggplot2::theme_classic() +
         ggplot2::theme(legend.position = "bottom"))
}

#' Split rows in a data frame into multiple observations at specified time
#' intervals
#'
#' @description
#' Splits each record in the data frame into a series of observation windows
#' based on the supplied cut times.
#' Can be used to reshape data with time-varying covariates to the appropriate
#' format for [Survdt2()].
#' The main workhorse is a modified call to [survival::survSplit()].
#'
#' @param data A data frame to be split.
#' @param cut A numeric vector of cut times.
#' @param start_name The name for the start time column to be created in the split
#'  data. If it matches a column in `data` (meaning the data was previously split)
#'  then `event_name` must also match a column in `data`.
#' @param stop_name The name of the existing event time or stop time column in
#'  `data`.
#' @param event_name The name for the event indicator column to be created in the
#'  split data. See `start_name` for special requirements.
#' @param id_name An optional name to create an id column in the split data. If
#'  this column already exists in `data` its values will be used as ids.
#'  If the data is already in (start, stop] format, do not use this argument.
#' @param origin The time that the follow up starts at. Defaults to zero.
#'
#' @details
#' A basic use case of [time_split()] is to take a data frame with one row per
#' subject and transform it into a data frame with multiple rows per subject,
#' each corresponding to a specific observation time window.
#' The event time information observed in each window is encoded by four
#' potentially new variables:
#' * `start`: the starting time for the observation window.
#' * `stop`: the end time for the observation window.
#' * `event`: equal to one when the subject's event occurred at `stop`, and zero
#'   otherwise.
#' * `id`: the label that identifies which subject the data corresponds to.
#'
#' In this use case, the name of the event time column in `data` should always be
#' provided, as well as the name for the new `id` column if one does not already
#' exist.
#' After `cut` is sorted, the \eqn{i}th subject's observation windows will be
#' \eqn{(}`origin`, `cut[1]`\eqn{]}, \eqn{(}`cut[1]`, `cut[2]`\eqn{]}, \eqn{\ldots},
#' \eqn{(}`cut[j]`, \eqn{T_i}\eqn{]} for some
#' `j`\eqn{\leq}`length(cut)`.
#'
#' Another use case of [time_split()] is to introduce additional splits in a
#' previously split data frame.
#' The names for the existing `start`, `stop`, and `event` columns should be
#' provided to [time_split()] in this case.
#' Requesting a new `id` variable for previously cut data will return an error,
#' since the new ids would not respect the prior groupings.
#'
#' @returns
#' A split data frame with (start, stop] records.
#' The new column names are taken from the input arguments of `time_split()`.
#'
#' @export
#'
#' @seealso
#' * [Survdt2()]: create survival objects with split data.
#' * [coxph_indtrunc()]: Cox regression with time-varying covariates.
#' * [survdt2to1()]: collapse data back to one row per subject.
#'
#' @examples
#' splitdat1 <- time_split(aids, cut = c(20, 40),
#'                         stop_name = "incu",
#'                         id_name = "id",
#'                         origin = -1)
#' tail(splitdat1)
#'
#' splitdat2 <- time_split(splitdat1, cut = c(10, 30),
#'                         start_name = "start",
#'                         stop_name = "incu",
#'                         event_name = "event",
#'                         id_name = "id",
#'                         origin = -1)
#' tail(splitdat2)
time_split <- function(data, cut, stop_name, start_name = "start",
                       event_name = "event", id_name = NULL, origin = 0){
  # check input
  if (is.matrix(data)) data <- as.data.frame(data)
  stopifnot(is.data.frame(data), is.numeric(cut), is.vector(cut), !any(is.na(cut)),
            is.character(start_name), is.vector(start_name), !any(is.na(start_name)),
            is.character(stop_name), is.vector(stop_name), !any(is.na(stop_name)),
            is.character(event_name), is.vector(event_name), !any(is.na(event_name)),
            is.character(start_name), is.vector(start_name), !any(is.na(start_name)),
            is.numeric(origin), is.vector(origin), !any(is.na(origin)))
  stopifnot(length(cut) > 0, length(start_name) == 1, length(stop_name) == 1,
            length(event_name) == 1, nrow(data) > 0, ncol(data) > 0)
  if (!is.null(id_name)){
    stopifnot(is.character(id_name), is.vector(id_name), !any(is.na(id_name)))
    stopifnot(length(id_name) == 1)
  }
  if (any(cut <= origin)) stop("All cut times must be >origin")

  colindx <- match(c(start_name, stop_name, event_name), colnames(data))
  if (is.na(colindx[2]))
    stop("No column in data matches stop_name (should be the stop or event time)")
  if (is.na(colindx)[1] != is.na(colindx[3]))
    stop(paste0("Must have (1) both start_name and event_name exist in data,\n",
                "or (2) neither exist in data"))
  min_time <- min(data[, stop_name, drop = TRUE], na.rm = TRUE)
  if (is.na(min_time)) stop("All stop times are NA")
  if (origin >= min_time)
    stop("Time origin must be strictly less than all stop times")

  # construct observation windows
  if (is.na(colindx[1])){
    form <- stats::reformulate(".",
                               response = paste0("Surv(",
                                                 stop_name,
                                                 ", ",
                                                 "rep(1, nrow(data))",
                                                 ")"))
  }else{
    form <- stats::reformulate(".",
                               response = paste0("Surv(",
                                                 start_name,
                                                 ", ",
                                                 stop_name,
                                                 ", ",
                                                 event_name,
                                                 ")"))
  }

  # survsplit
  Surv <- survival::Surv
  if (is.null(id_name) || (id_name %in% colnames(data))){
    res <- survival::survSplit(form, data = data, cut = cut, start = start_name,
                               zero = origin, end = stop_name, event = event_name)
  }else{
    if (!is.na(colindx[1]))
      stop("Cannot preserve ids in (start, stop] data; id_name not in data")

    res <- survival::survSplit(form, data = data, cut = cut, start = start_name,
                               zero = origin, end = stop_name, event = event_name,
                               id = id_name)


  }
  res
}
