#' Create a survival object for use with time-varying covariates
#'
#' @description
#' Creates a survival object for doubly truncated data with the follow
#' up time for each subject split across multiple observation windows using
#' (start, stop] encoding.
#'
#' @param start A vector of start times for the observation windows.
#' @param stop A vector of stop times for the observation windows.
#' @param event A vector of 0/1 event indicators (1=event).
#' @param ltrunc A vector of left truncation times.
#' @param rtrunc A vector of right truncation times.
#' @param id A vector of subject ids.
#' @param x Any R object.
#' @param i,j Row and columns index vectors. The result is a `Survdt2` object if
#'   `j` is missing, otherwise `x` is treated as a matrix.
#' @param ... Additional arguments passed to other methods.
#'
#' @details
#' The \eqn{i}th subject's survival data consists of the left and right truncation
#' times \eqn{L_i} and \eqn{R_i} as well as a series of (start, stop] observation
#' windows \eqn{(t_{i0}}, \eqn{t_{i1}]}, \eqn{(t_{i1}}, \eqn{t_{i2}]},
#' \eqn{\ldots}, \eqn{(t_{i,j-1}}, \eqn{t_{ij}}\eqn{]} for some \eqn{j}.
#' Furthermore, each of these observation windows has an associated event indicator
#' that is equal to one if the event time is equal to the stop time, and zero
#' otherwise.
#' Each subject should have an observation window with an event, and multiple
#' events per subject are not allowed.
#' Attempting to run [Survdt2()] with inputs that do not follow this specification
#' will result in an error.
#'
#' This data structure is used for Cox regression with time-varying covariates,
#' since the subject's covariate values can vary across observation windows.
#'
#'
#' @returns
#' * [Survdt2()] returns an object of class `Survdt2`, with columns `start`,
#'   `stop`, `event`, `ltrunc`, `rtrunc`, and `id`.
#'   Indexing a `Survdt2` object by rows will return another `Survdt2` object,
#'   otherwise it is treated as a matrix.
#'   The length of a `Survdt2` object is the total number of observations across
#'   all subjects, and `is.na()` returns a logical index vector for the non-complete
#'   rows.
#'   Calling `as.data.frame` on a `Survdt2` object will return a data frame
#'   stripped of the `Survdt2` class attribute.
#' * [is.Survdt2()] returns `TRUE` if `x` inherits from the class `Survdt2` and
#'   `FALSE` otherwise.
#' * [survdt2to1()] takes a `Survdt2` object as input and returns the result from
#'   calling [Survdt()] on collapsed data with one row per `id`.
#'   The subjects in the collapsed data are sorted in the order they first appear
#'   in the `Survdt2` object.
#'   If all event indicators for a subject are zero and/or `NA`, the event time
#'   in the collapsed data is set to `NA`.
#'   Lastly, the `id` values are preserved within the row names of the `Survdt`
#'   object.
#' * The plot method for `Survdt2` first collapses the data into `Survdt`
#'   format using [survdt2to1()] and then calls [plot.Survdt()].
#'
#' @export
#'
#' @seealso
#' * [coxph_indtrunc()]: inverse probability weighted Cox regression with
#'   time-varying covariates.
#' * [Survdt()]: the more basic survival object with one row per subject.
#'
#' @examples
#' # first create (start, stop] data frame;
#' # could then apply time transformations
#' # to covariates
#' splitdat <- time_split(aids, cut = c(20, 40),
#'                        stop_name = "incu",
#'                        id_name = "id",
#'                        origin = -1)
#'
#' y1 <- with(splitdat,
#'            Survdt2(start, incu, event, ltrunc, rtrunc, id))
#' plot(y1)
#'
#' y2 <- with(aids,
#'            Survdt(incu, ltrunc, rtrunc))
#' max(abs(unclass(survdt2to1(y1)) - unclass(y2)))
#'
#'
#' coxph_indtrunc(Survdt2(start, incu, event, ltrunc, rtrunc, id) ~ age_gp,
#'                data = splitdat)
#'
Survdt2 <- function(start, stop, event, ltrunc, rtrunc, id){
  # check input
  if (inherits(start, "difftime")) start <- as.numeric(start)
  if (inherits(stop, "difftime")) stop <- as.numeric(stop)
  if (inherits(ltrunc, "difftime")) ltrunc <- as.numeric(ltrunc)
  if (inherits(rtrunc, "difftime")) rtrunc <- as.numeric(rtrunc)
  if (is.logical(event)) event <- as.numeric(event)

  stopifnot(is.numeric(start), is.vector(start), !any(is.infinite(start)),
            is.numeric(stop), is.vector(stop), !any(is.infinite(stop)),
            is.numeric(event), is.vector(event),
            is.numeric(ltrunc), is.vector(ltrunc), !any(is.infinite(ltrunc)),
            is.numeric(rtrunc), is.vector(rtrunc), !any(is.infinite(rtrunc)),
            is.factor(id) || is.vector(id), !is.list(id))
  if (!is.factor(id)) id <- as.factor(id)

  stopifnot(length(start) > 0, length(start) == length(stop),
            length(start) == length(event), length(start) == length(ltrunc),
            length(start) == length(rtrunc), length(start) == length(id),
            !any(is.na(id)), all(is.na(event) | (event %in% c(0, 1))))

  if (any(stats::aggregate(event ~ id,
                           FUN = function(z) {sum(is.na(z) | (z == 1)) < 1})[, 2]))
    stop("cannot have an id with all event=0")
  if (any(stats::aggregate(ltrunc ~ id,
                           FUN = function(z) {length(unique(z)) > 1})[, 2]))
    stop("ltrunc must be constant within id")
  if (any(stats::aggregate(rtrunc ~ id,
                           FUN = function(z) {length(unique(z)) > 1})[, 2]))
    stop("rtrunc must be constant within id")

  na_ltrunc <- is.na(ltrunc)
  na_rtrunc <- is.na(rtrunc)
  na_trunc <- na_ltrunc | na_rtrunc
  na_start <- is.na(start)
  na_stop <- is.na(stop)
  event_indx <- !is.na(event) & (event == 1)
  if (any(event_indx) && any(duplicated(id[event_indx])))
    stop("cannot have multiple events per id")
  if (any(!na_start & !na_stop) && any((start >= stop)[!na_start & !na_stop]))
    stop("must have start<stop")
  if (any(!na_trunc) && any((ltrunc >= rtrunc)[!na_trunc]))
    stop("ltrunc must be <rtrunc")
  if (any(event_indx) && any(!na_stop[event_indx] & !na_ltrunc[event_indx]) &&
      any((ltrunc[event_indx] >= stop[event_indx])[!na_stop[event_indx] &
                                                   !na_ltrunc[event_indx]]))
    stop("each event time must be >ltrunc")
  if (any(event_indx) && any(!na_stop[event_indx] & !na_rtrunc[event_indx]) &&
      any((rtrunc[event_indx] < stop[event_indx])[!na_stop[event_indx] &
                                                  !na_rtrunc[event_indx]]))
    stop("each event time must be <=rtrunc")

  id_levels <- levels(id)
  id <- as.numeric(id)
  x <- cbind(start = start, stop = stop, event = event, ltrunc = ltrunc,
             rtrunc = rtrunc, id = id)
  attr(x, "id_levels") <- id_levels

  complete_id_indx <- sapply(unique(id),
                             FUN = function(z) {!any(is.na(x[id == z, ]))})
  if (any(complete_id_indx)){
    tmpx <- x[complete_id_indx, ]
    tmpx <- tmpx[order(tmpx[, "id", drop = TRUE], tmpx[, "start", drop = TRUE]), ]
    if (any((tmpx[-nrow(tmpx), "stop", drop = TRUE] !=
             tmpx[-1, "start", drop = TRUE]) &
            (tmpx[-nrow(tmpx), "id", drop = TRUE] ==
             tmpx[-1, "id", drop = TRUE])))
      stop("(start, stop] windows must be adjacent without overlap or gaps")
  }

  class(x) <- "Survdt2"
  x
}

#' @rdname Survdt2
#' @export
survdt2to1 <- function(x){
  stopifnot(is.Survdt2(x))
  x <- x[order(factor(x[, "id", drop = TRUE],
                      levels = unique(x[, "id", drop = TRUE])),
               x[, "stop", drop = TRUE],
               na.last = FALSE), ]

  id_indx <- which(!duplicated(x[, "id", drop = TRUE], fromLast = TRUE))
  event_indx <- which(!is.na(x[, "event", drop = TRUE]) &
                        (x[, "event", drop = TRUE] == 1))

  unique_ids <- x[id_indx, "id", drop = TRUE]
  id_has_event <- unique_ids %in% x[event_indx, "id", drop = TRUE]
  id_indx[id_has_event] <- event_indx
  event_times <- ifelse(id_has_event, x[id_indx, "stop", drop = TRUE], NA)
  names(event_times) <- as.character(x[id_indx, "id", drop = TRUE])
  Survdt(event_times,
         x[id_indx, "ltrunc", drop = TRUE],
         x[id_indx, "rtrunc", drop = TRUE])
}

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

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

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

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

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

#' @rdname Survdt2
#' @export
plot.Survdt2 <- function(x, ...){
  plot(survdt2to1(x), ...)
}
