#' Cox regression under quasi-independent double truncation.
#'
#' @description
#' Computes an inverse probability weighted Cox regression estimator using
#' nonparametric weights (see [npmle()]).
#' Includes options for robust estimators with time-varying weights, which
#' introduce an event time stabilization weight (see below).
#'
#' @param formula A model formula with a `Survdt` or `Survdt2` response.
#'   Subjects with any missing values are discarded.
#'   See [Survdt()] and [Survdt2()].
#' @param data An optional data frame used to evaluate `formula` and
#' `strata_formula`.
#' @param strata_formula An optional one-sided formula with stratification
#'   variables on the right hand side, to fit a stratified Cox model.
#' @param ipw_type Type of inverse probability weights.
#'   1. `"W-asurv"` (the default): stabilized robust time-varying weights based
#'      on the NPMLE survival function.
#'   2. `"W-surv"`: robust time-varying weights based on the NPMLE survival function.
#'   3. `"W-1"`: standard weights with no time-varying component (Rennert and Xie, 2018)
#'   4. `"W-a"`: stabilized weights (Mandel et al., 2018)
#'   5. For custom time-varying weights, provide a function that takes a numerical
#'      vector of times as input and returns a numerical vector of positive weights.
#' @param se Whether to calculate standard errors (default `TRUE`).
#' @param robust_se Whether to calculate standard errors that are robust to
#'   model misspecification (default `TRUE`).
#'   Ignored if `se`=`FALSE`.
#' @param basehaz Whether to return the baseline hazard estimate
#'   (default `FALSE`).
#' @param restr_times An optional numerical vector of length two giving a time
#'   window to restrict the score function.
#'   Values of `NA` are replaced with the minimum and maximum observed event
#'   time respectively, i.e. no restriction.
#'   The default setting is no restriction.
#' @param npmle_fit An optional pre-computed `npmle` object fit on the complete
#'   cases.
#'   If it was fit with `ainv`=`FALSE` (or `infl`=`FALSE` when `se`=`TRUE`) then
#'   it is ignored and a new NPMLE is fit.
#'   See [npmle()].
#' @param robust_infl A square matrix of influence values for the event weight,
#'   with one column per individual and one row per follow up time, both in the
#'   order that the subject appears in the response variable.
#'   Only used if `ipw_type` is a function and `robust_se`=`TRUE`.
#' @param simple_labels Whether to drop the factor names from the stratum labels
#'   (Default `FALSE`).
#' @param keep_x Whether to include the model matrix in the `cox_fit` component
#'   of the return value (Default `FALSE`).
#' @param ... Additional arguments passed to other methods.
#' @param object,x A `coxph_indtrunc` object.
#' @param parm A numeric or character vector specifying a set of model
#'   parameters.
#'   Defaults to all identifiable (non-missing) coefficients.
#' @param level Confidence level for confidence intervals.
#' @param type The type of residuals to return.
#' @param per_subject Whether to return per-subject residuals as opposed to
#'   per-observation residuals, for models with `Survdt2` data.
#'   Defaults to `TRUE` for deviance residuals and `FALSE` otherwise.
#'
#' @details
#' Given a sample of doubly truncated data with event times \eqn{T_i}, truncation
#' times \eqn{L_i} and \eqn{R_i}, and potentially time-varying covariates
#' \eqn{X_i(t)}, let \eqn{N_i(t)=1(T_i \leq t)} be the usual counting process
#' and \eqn{Y_i(t)=1(T_i \geq t)} be the at-risk process.
#'
#' [coxph_indtrunc()] fits a Cox model by maximizing an inverse probability
#' weighted partial likelihood, where the double truncation bias is accounted
#' for through nonparametric estimates of
#' \eqn{a(t)^{-1}=P(L<T \leq R)/P(L<T \leq R|T=t)} (see [npmle()]).
#' Assuming `strata_formula=NULL` for sake of illustration, setting
#' `ipw_type="W-1"` specifies non-time-varying weights, so the model is fitted by
#' finding the root of the weighted score function
#' \deqn{\sum_{i}\int \frac{1}{\widehat{a}(T_i)}
#'       [X_i(t) - \overline{X}_{\widehat{a}}(t,\beta)]dN_i(t)}
#' \deqn{\text{where } \overline{X}_{\widehat{a}}(t,\beta)
#'       = \frac{\sum_j X_j(t)\widehat{a}(T_j)^{-1}Y_j(t)\exp(X_j(t)^T\beta)}
#'              {\sum_j \widehat{a}(T_j)^{-1}Y_j(t)\exp(X_j(t)^T\beta)}.}
#'
#' More generally, [coxph_indtrunc()] allows for time-varying weights with a
#' stabilization component \eqn{\widehat{w}(t)} that determines the relative
#' importance of each event time.
#' This can improve robustness and numerical convergence without introducing any
#' bias, in contrast to weight truncation, for example.
#' The weighted score function becomes
#' \deqn{\sum_{i}\int \frac{\widehat{w}(t)}{\widehat{a}(T_i)}
#'       [X_i(t) - \overline{X}_{\widehat{a}}(t,\beta)]dN_i(t).}
#' Several choices of time-varying weights are available:
#' * `ipw_type="W-surv"`: \eqn{\widehat{w}(t)=\widehat{S}(t-)}, where
#'    \eqn{\widehat{S}(t)} is the marginal survival function estimated by
#'    [npmle()].
#'    Can increase robustness to outliers.
#' * `ipw_type="W-a"`: \eqn{\widehat{w}(t)=\widehat{a}(t)}.
#'    Can reduce the influence of extreme inverse probability weights.
#' * `ipw_type="W-asurv"`: \eqn{\widehat{w}(t)=\widehat{a}(t)\widehat{S}(t-)}.
#'    Combines the benefits of both weights.
#' * Custom stabilization weights \eqn{\widehat{w}(t)} can be specified by
#'   supplying a function as the input to `ipw_type`.
#'   If `robust_se=TRUE`, this also requires the influence values for
#'   \eqn{\widehat{w}(t)} to be supplied in `robust_infl`, in order to account
#'   for its variability in the standard errors.
#'
#' @section Stratified Cox model:
#'
#' A stratified Cox model is specified through `strata_formula`.
#' The formula should look like `~svar1 + svar2`... where `svar1` and `svar2` are
#' factor variables defining the strata.
#' This allows the baseline hazard to differ across strata.
#'
#' @section Restricted time window in score function:
#' The user can choose to restrict the time range of the score function during
#' estimation.
#' This may be useful when extreme event times have very large weights.
#' Setting the input argument `restr_times=c(t1,t2)` will restrict the range of
#' the integrals in the score function to the interval \eqn{[}`t1`, `t2`\eqn{]}
#' instead of the entire event time range.
#' Notably, this time restriction will not introduce any bias in the fitted model.
#'
#' @returns
#' [coxph_indtrunc()] returns an object of class `coxph_indtrunc`, i.e. a list
#' with the following components:
#' \item{n}{The complete case sample size.}
#' \item{n_miss}{The number of individuals removed due to missing data.}
#' \item{n_event}{The number of observed events in the score function.}
#' \item{coef}{A four-column matrix containing the estimated coefficients,
#'   standard errors, z-statistics, and p-values. Coefficient values of `NA`
#'   indicate collinearity.}
#' \item{var_matrix}{The estimated variance matrix for the coefficient estimates.
#'   Values of `NA` indicate collinearity.}
#' \item{robust_se}{The input value for `robust_se`.}
#' \item{offset}{The offsets from the model formula for complete case observations.}
#' \item{ipw}{The estimated inverse probability of non-truncation for each
#'   complete case observation.}
#' \item{time_wt}{The time-varying component of the inverse probability weights
#'   for each complete case observation.}
#' \item{cc_indx}{A logical index vector for the complete cases in the original
#'   data.}
#' \item{id_indx}{A vector of ids matching each observation with a subject among
#'   the complete cases.}
#' \item{sort_indx}{A numerical vector to sort the complete cases into the order
#'   they are used in `cox_fit` (see below).}
#' \item{s}{The stratum value for each complete case observation if a stratified
#'   Cox model was fit.
#'   Otherwise `NULL`.}
#' \item{martingale_resid}{The weighted martingale residuals for the complete
#'   cases, with one entry per observation.}
#' \item{deviance_resid}{The weighted deviance residuals for the complete cases,
#'   with one entry per subject.}
#' \item{basehaz}{For input `basehaz`=`TRUE`, a list with first element `s`
#'   containing the stratum values and one additional element for each stratum,
#'   in the order of `s`.
#'   Each contains a matrix with following columns:
#'   `time`, `cumhaz` with the cumulative baseline hazard value (baseline at x=0),
#'   `var_logcumhaz` with the estimated variance of the log(cumhaz), and
#'   additional columns, labeled by predictor names, with the estimated covariance
#'   between the log(cumhaz) and the regression coefficients.
#'   For input `basehaz`=`FALSE`, `NULL`.}
#' \item{timerange}{A vector with the minimum and maximum times used in the
#' score function.}
#' \item{eventrange}{A vector with the unrestricted range of observed event times.}
#' \item{wt_name}{A string giving the name of the weights used in `ipw_type`.}
#' \item{cox_fit}{The `coxph` object used to obtain the estimates.
#'   Note that the data was sorted by strata and time before fitting the model,
#'   and the offset may include an additional term of `-log(time_wt)`.}
#' \item{strata_summary}{Some summary statistics for each stratum, computed
#'   prior to applying `restr_times`.}
#' \item{call_}{The call to [coxph_indtrunc()].}
#'
#' `coxph_indtrunc` objects have methods for `vcov` to extract the variance matrix
#' of estimated coefficients, `coef` to extract the estimated coefficients,
#' `confint` to compute confidence intervals for coefficients, `residuals` for
#' model diagnostics, and `print` to display a summary of the fitted model.
#'
#' `residuals` returns per-subject residuals in the order that the subject first
#' appeared in the original data.
#' The deviance residuals are computed as
#' \eqn{\text{sign}(\widehat{M}_{wi})\sqrt{
#' -2[\widehat{M}_{wi} +
#' \delta_{wi}\log(1 - \widehat{M}_{wi}/\delta_{wi})]}}
#' where \eqn{\widehat{M}_{wi}} is the weighted martingale residual and
#' \eqn{\delta_{wi}} is the weighted event indicator, following
#' Therneau et al. (1990).
#'
#' @export
#'
#' @seealso
#' * [plot_coxsurv()]: compute/plot estimates and pointwise confidence intervals
#'   for the conditional survival function or cumulative hazard from a fitted Cox
#'   model.
#' * [positivity_sens_indtrunc()]: sensitivity analysis for positivity violations
#'   in the Cox model.
#' * [test_quasiindep_covariates()]: test for dependence between truncation times
#'   and the event time or covariates.
#'   The nonparametric weights used by [coxph_indtrunc()] requires quasi-independence.
#' * `vignette("cox-regression-quasiindep", package = "survdt")`: an
#'   illustration of Cox regression and related diagnostics applied to the
#'   [aids] data, if vignettes were built when installing [survdt].
#'
#' @examples
#' coxph_indtrunc(Survdt(incu, ltrunc, rtrunc) ~ age_gp,
#'                data = aids)
#'
#' @references
#'
#' Mandel, M., de Uña-Álvarez, J., Simon, D. K., & Betensky, R. A. (2018).
#' Inverse probability weighted Cox regression for doubly truncated data.
#' *Biometrics*, **74**(2), 481--487.
#'
#' Rennert, L., & Xie, S. X. (2018). Cox regression model with doubly truncated
#' data. *Biometrics*, **74**(2), 725--733.
#'
#' Therneau, T. M., Grambsch, P. M., & Fleming, T. R. (1990). Martingale-based
#' residuals for survival models. *Biometrika*, **77**(1), 147--160.
#'
#' Vazquez, O., & Xie, S. X. (2025). Robust inverse probability weighted
#' estimators for doubly truncated Cox regression with closed-form standard
#' errors. *Lifetime Data Analysis*, **31**(2), 364–-393.
coxph_indtrunc <- function(formula, data = NULL, strata_formula = NULL,
                           ipw_type = c("W-asurv", "W-surv", "W-1", "W-a"),
                           se = TRUE, robust_se = TRUE, basehaz = FALSE,
                           restr_times = NULL, npmle_fit = NULL,
                           robust_infl = NULL, simple_labels = FALSE,
                           keep_x = FALSE){
  # check input
  if (missing(formula))
    stop('\n  argument "formula" is missing, with no default')

  fn_call <- match.call()
  vals <- prep_data_for_fit(formula, data, strata_formula, simple_labels)
  if (is.Survdt2(vals$y))
    y_npmle <- survdt2to1(vals$y[vals$complete_cases, ])
  else
    y_npmle <- vals$y[vals$complete_cases, ]
  n <- length(y_npmle)
  s <- NULL
  if (!is.null(vals$s)) s <- vals$s[vals$sort_indx]

  if (!is.null(s)){
    strata_vals <- unique(s)
    n_strata <- length(strata_vals)
    if (n_strata < 2)
      stop(paste0("Only one stratum in complete cases.\n ",
                  "Either update the strata formula or drop it altogether"))

    if (is.Survdt2(vals$y)){
      event_rows <- vals$sorted_cc[vals$y[vals$sorted_cc, "event",
                                          drop = TRUE] == 1]
    }else{
      event_rows <- vals$sorted_cc
    }

    n_unique_strata <- stats::aggregate(vals$y[event_rows, vals$time_col,
                                               drop = TRUE] ~ s,
                                        FUN = function(x) c(length(unique(x))))
    small_strata <- n_unique_strata[n_unique_strata[, 2, drop = TRUE] < 2, 1,
                                    drop = TRUE]
    if (length(n_unique_strata) > 5)
      n_unique_strata <- c(as.character(n_unique_strata[1:5]), "...")
    if (length(small_strata) > 0)
      warning(paste0("Less than two unique event times found in stratum: ",
                     paste0(small_strata, collapse = "; ")))
    strata_summary <- data.frame(strata = n_unique_strata[, 1, drop = TRUE],
                                 strata_size =
                                   sapply(n_unique_strata[, 1, drop = TRUE],
                                          FUN = function(x) sum(s == x)),
                                 unique_event_times =
                                   n_unique_strata[, 2, drop = TRUE])
  }else{
    strata_vals <- NULL
    n_strata <- 1
    strata_summary <- NULL
  }

  if (is.null(restr_times)) restr_times <- c(NA_real_, NA_real_)
  stopifnot((is.vector(ipw_type) && is.character(ipw_type)) ||
              is.function(ipw_type),
            is.logical(se), is.vector(se), length(se) == 1, all(!is.na(se)),
            is.logical(robust_se), is.vector(robust_se), length(robust_se) == 1,
            all(!is.na(robust_se)), is.logical(basehaz), is.vector(basehaz),
            length(basehaz) == 1, all(!is.na(basehaz)), is.numeric(restr_times),
            is.vector(restr_times), length(restr_times) == 2,
            is.null(npmle_fit) || is.npmle(npmle_fit),
            is.null(robust_infl) || is.matrix(robust_infl),
            is.vector(keep_x), is.logical(keep_x), length(keep_x) == 1,
            all(!is.na(keep_x)))
  robust_se <- ifelse(se == TRUE, robust_se, FALSE)

  # time window
  eventrange <- range(y_npmle[, "time", drop = TRUE])
  time_lims <- eventrange
  if (!is.na(restr_times[1]) && (restr_times[1] >= time_lims[2]))
    stop("restr_times[1] must be less than the maximum observed event time")
  if (!is.na(restr_times[2]) && (restr_times[2] <= time_lims[1]))
    stop("restr_times[2] must be greater than the minimum observed event time")
  if (all(!is.na(restr_times)) && (restr_times[1] >= restr_times[2]))
    stop("Must have restr_times[1]<restr_times[2]")
  test_window <- ifelse(is.na(restr_times), time_lims,
                        c(max(restr_times[1], time_lims[1]),
                          min(restr_times[2], time_lims[2])))
  ltrunc_indx <-
    vals$y[vals$sorted_cc, vals$time_col, drop = TRUE] < test_window[1]
  rcens_indx <-
    test_window[2] < vals$y[vals$sorted_cc, vals$time_col, drop = TRUE]
  n_events <- sum(!ltrunc_indx[vals$event_indx] & !rcens_indx[vals$event_indx])
  if (n_events == 0) stop("No observed event times in restricted window")

  if (is.character(ipw_type)){
    ipw_type <- match.arg(ipw_type)
    wt_name <- ipw_type
  }else{
    if ((robust_se == TRUE) && is.null(robust_infl))
      stop(paste0("Supplying a function for ipw_type with ",
                  "robust_se=TRUE requires a matrix of influence values in ",
                  "robust_infl"))
    time_wt <- ipw_type(y_npmle[, "time", drop = TRUE])
    if (!is.numeric(time_wt) || !is.vector(time_wt))
      stop("Function ipw_type did not return a numerical vector")
    if (length(time_wt) != length(y_npmle))
      stop(paste0("Function ipw_type should return ",
                  "one value per time input. Output length is ", length(time_wt),
                  " but ", length(y_npmle), " times requested"))

    if (any(is.na(time_wt)))
      stop("Missing values in output of function ipw_type")
    if (any(time_wt < 0))
      stop("Negative values in output of function ipw_type")
    if (robust_se == TRUE){
      stopifnot(nrow(robust_infl) == ncol(robust_infl))
      if (nrow(robust_infl) != vals$n_init)
        stop(paste0("Matrix robust_infl has dimension ", nrow(robust_infl),
                    ", does not match number of subjects ", vals$n_init))
      robust_infl <- robust_infl[vals$complete_id_nums, vals$complete_id_nums,
                                 drop = FALSE]
      stopifnot(!any(is.na(robust_infl)), !any(is.infinite(robust_infl)))
    }
    wt_name <- paste0("user-supplied function ", format(substitute(ipw_type)))
  }

  # npmle
  if (!is.null(npmle_fit)){
    stopifnot(is.npmle(npmle_fit))
    if (nrow(npmle_fit[["masses"]]) != vals$n)
      stop(paste0("npmle_fit sample size (", nrow(npmle_fit[["masses"]]), ") ",
                  "does not match the number of complete cases in the model ",
                  "frame (", vals$n, ")"))
    if (is.null(npmle_fit[["ainv"]]) ||
        (is.null(npmle_fit[["infl_ainv"]]) && (se == TRUE)))
      npmle_fit <- NULL
    else{
      if (any(abs(y_npmle[, "time", drop = TRUE] -
                  npmle_fit[["ainv"]][, "time", drop = TRUE]) > 1e-9))
        stop("Follow up times in formula response and npmle_fit do not match")
    }
  }
  if (is.null(npmle_fit)){
    npmle_fit <- npmle(y_npmle, ainv = TRUE, se = se, infl = se)
    if (!npmle_fit[["converged"]]) stop("Fitted npmle did not converge")
  }
  ipw <- npmle_fit[["ainv"]][, "ainv", drop = TRUE]

  # time-varying weights
  if (is.character(ipw_type)){
    calculate_infl <- (se == TRUE) && (robust_se == TRUE)
    if (ipw_type == "W-1"){
      time_wt <- rep(1, n)
    }else if (ipw_type == "W-a"){
      time_wt <- 1 / ipw
      if (calculate_infl == TRUE){
        robust_infl <- apply(npmle_fit[["infl_ainv"]], MARGIN = 2,
                             FUN = function(x) {-x / (ipw^2)})
      }
    }else{
      cdf_time_match <- pmin(findInterval(y_npmle[, "time", drop = TRUE],
                                          npmle_fit[["cdf"]][,
                                                             "time",
                                                             drop = TRUE],
                                          left.open = TRUE) + 1,
                             nrow(npmle_fit[["cdf"]]))
      lcontsurv <- c(1, 1 - npmle_fit[["cdf"]][,
                                               "cdf",
                                               drop = TRUE])[cdf_time_match]

      if (ipw_type == "W-surv"){
        time_wt <- lcontsurv
        if (calculate_infl == TRUE){
          robust_infl <- -rbind(0, npmle_fit[["infl_cdf"]])[cdf_time_match, ,
                                                            drop = FALSE]
        }
      }else if (ipw_type == "W-asurv"){
        time_wt <- lcontsurv / ipw
        if (calculate_infl == TRUE){
          robust_infl <- sapply(seq_len(n), FUN = function(i){
            (-npmle_fit[["infl_ainv"]][, i, drop = TRUE] *
               lcontsurv / (ipw^2)) -
              (c(0,
                 npmle_fit[["infl_cdf"]][, i, drop = TRUE])[cdf_time_match] /
                 ipw)
          })
        }
      }
    }
  }
  time_wt <- time_wt + 1e-7
  # time_wt and ipw are in original data order, one entry per id

  X <- vals$X[vals$sorted_cc, , drop = FALSE]
  ofst <- vals$ofst[vals$sorted_cc]
  n_miss <- vals$n_init - vals$n
  expanded_time_wt <- time_wt[vals$match_id_indx[vals$sorted_cc]]
  ofst_wt <- ofst - log(expanded_time_wt)
  expanded_ipw <- ipw[vals$match_id_indx[vals$sorted_cc]]
  outer_wt <- expanded_ipw * expanded_time_wt

  if (is.Survdt2(vals$y)){
    y_coxph <- survival::Surv(vals$y[vals$sorted_cc, "start", drop = TRUE],
                              vals$y[vals$sorted_cc, "stop", drop = TRUE],
                              as.numeric((vals$y[vals$sorted_cc, "event",
                                                 drop = TRUE] == 1) &
                                           !ltrunc_indx & !rcens_indx))
  }else{
    y_coxph <- survival::Surv(vals$y[vals$sorted_cc, "time", drop = TRUE],
                              as.numeric(!ltrunc_indx & !rcens_indx))
  }
  n_event <- sum(y_coxph[, ncol(y_coxph), drop = TRUE])


  # fit cox model
  offset <- stats::offset
  strata <- survival::strata
  if (is.null(s)){
    cox_fit <- survival::coxph(y_coxph ~ X + offset(ofst_wt), weights = outer_wt,
                               na.action = stats::na.fail, ties = "breslow",
                               robust = FALSE, x = keep_x)
  }else{
    cox_fit <- survival::coxph(y_coxph ~ X + offset(ofst_wt) + strata(s),
                               weights = outer_wt, na.action = stats::na.fail,
                               ties = "breslow", robust = FALSE, x = keep_x)
  }

  # std errors
  p <- length(cox_fit[["coefficients"]])
  coef_na <- is.na(cox_fit[["coefficients"]])
  effective_p <- sum(!coef_na)
  if (effective_p == 0)
    stop("All coefficient estimates are NA. Check predictors for collinearity.")
  var_matrix <- matrix(NA, nrow = p, ncol = p)
  if (se == TRUE){
    # effective_p x n
    var_components <- t(stats::residuals(cox_fit, "score", weighted = TRUE,
                                         collapse = vals$id[vals$sorted_cc]))
    var_components <- var_components[!coef_na, , drop = FALSE]
    var_components <- var_components +
      (var_components %*%
         diag(1 / ipw, length(ipw)) %*%
         npmle_fit[["infl_ainv"]] / n)
    time_varying_weights <- is.function(ipw_type) ||
      (is.character(ipw_type) && (ipw_type != "W-1"))

    if ((robust_se == TRUE) && (time_varying_weights == TRUE)){
      delta_indx <- as.logical(y_coxph[, ncol(y_coxph), drop = TRUE])
      event_id_indx <- vals$complete_id_nums %in%
        (vals$id[vals$sorted_cc])[delta_indx]
      sch_resids <- t(stats::residuals(cox_fit, "schoenfeld",
                                       weighted = TRUE))[!coef_na, ,
                                                         drop = FALSE]
      var_components <- var_components +
        (sch_resids[, order(vals$sorted_cc[delta_indx]), drop = FALSE] %*%
           diag(1 / time_wt[event_id_indx], sum(event_id_indx)) %*%
           robust_infl[event_id_indx, , drop = FALSE] / n)
    }

    meat <- apply(var_components, MARGIN = 2, FUN = function(x) x %o% x)
    if (effective_p == 1){
      meat <- sum(meat)
    }else{
      meat <- rowSums(meat)
    }
    var_matrix[!coef_na, !coef_na] <-
      cox_fit[["var"]][!coef_na, !coef_na, drop = FALSE] %*%
      matrix(meat, nrow = effective_p) %*%
      cox_fit[["var"]][!coef_na, !coef_na, drop = FALSE]
  }

  ests <- cbind(coef = cox_fit[["coefficients"]], se = sqrt(diag(var_matrix)),
                zstat = cox_fit[["coefficients"]] / sqrt(diag(var_matrix)),
                pval = stats::pnorm(abs(cox_fit[["coefficients"]] /
                                          sqrt(diag(var_matrix))),
                                    lower.tail = FALSE) * 2)
  rownames(ests) <- colnames(X)
  rownames(var_matrix) <- colnames(X)
  colnames(var_matrix) <- colnames(X)

  # baseline hazard
  bhaz <- NULL
  if (basehaz == TRUE){
    if (is.null(s)) s <- factor(rep(1, nrow(y_coxph)))
    stratas <- unique(s)
    bhaz <- list(s = stratas)

    wexps <- exp(c(X[, !coef_na, drop = FALSE] %*%
                     cox_fit[["coefficients"]][!coef_na]) + ofst) *
      expanded_ipw

    for (i in seq_along(stratas)){
      ys <- unclass(cox_fit[["y"]])[s == stratas[i], , drop = FALSE]
      if (is.Survdt2(vals$y)){
        ys[, ncol(ys)] <- vals$y[vals$sorted_cc[s == stratas[i]],
                                 "event",
                                 drop = TRUE]
      }else{
        ys[, ncol(ys)] <- rep(1, nrow(ys))
      }

      strata_indx <- which(s == stratas[i])
      strata_ipw <- expanded_ipw[strata_indx]
      strata_ids <- (vals$id[vals$sorted_cc])[strata_indx]
      strata_uniqueids <- unique(strata_ids)
      strata_wexps <- wexps[strata_indx]

      etimes <- sort(unique(ys[ys[, ncol(ys), drop = TRUE] == 1,
                               ncol(ys) - 1, drop = TRUE]))
      n_etime_strata <- length(etimes)
      # row time, column observation window
      if (ncol(ys) == 3){
        atrisk <- apply(ys, MARGIN = 1,
                        FUN = function(z) (z[1] < etimes) & (etimes <= z[2]))
      }else{
        atrisk <- sapply(ys[, 1, drop = TRUE], FUN = function(z) z >= etimes)
      }
      s0 <- c(atrisk %*% strata_wexps)
      hazincr <- sapply(etimes, FUN = function(tj) {
        sum(strata_ipw[(ys[, ncol(ys), drop = TRUE] == 1) &
                         (ys[, ncol(ys) - 1, drop = TRUE] == tj)])
      }) / s0

      bhaz[[i+1]] <- matrix(NA, ncol = 3 + effective_p, nrow = length(etimes))
      colnames(bhaz[[i+1]]) <- c("time", "cumhaz", "var_logcumhaz",
                                 colnames(X)[!coef_na])
      bhaz[[i+1]][, "time"] <- etimes
      bhaz[[i+1]][, "cumhaz"] <- cumsum(hazincr)

      if (se == TRUE){
        xavg <- diag(1 / s0, length(s0)) %*%
          (atrisk %*% diag(strata_wexps, length(strata_wexps)) %*%
             X[strata_indx, !coef_na, drop = FALSE])

        match_etimes <- match(ys[, ncol(ys) - 1, drop = TRUE],
                              etimes,
                              nomatch = NA_integer_)
        dM_infl <- sapply(strata_uniqueids, FUN = function(j){
          id_indx <- which(strata_ids == j)
          match_id_times <- c(apply(atrisk[, id_indx, drop = FALSE], MARGIN = 1,
                                    FUN = function(z) {
                                      colindx <- utils::tail(which(z), 1)
                                      ifelse(length(colindx) == 0, NA, colindx)
                                      }))
          nomatch <- is.na(match_id_times)
          x <- rep(0, length(etimes))
          if (any(!nomatch)){
            event_id <- id_indx[which(ys[id_indx,
                                         ncol(ys),
                                         drop = TRUE] == 1)[1]]
            if (!is.na(event_id)){
              x[match_etimes[event_id]] <- (strata_ipw[event_id])[1]
            }
            x <- x - ifelse(nomatch,
                            0,
                            (strata_wexps[id_indx[match_id_times]] * hazincr))
            x <- n * x / s0
          }
          cumsum(x)
        })

        # row time, col individual
        infl_cumhaz <- apply(xavg, MARGIN = 2,
                             FUN = function(z) cumsum(z * hazincr)) %*%
          cox_fit[["var"]][!coef_na, !coef_na, drop = FALSE] %*%
          var_components * (-n)

        strata_indivs <- match(strata_uniqueids, vals$complete_id_nums)
        infl_cumhaz[, strata_indivs] <- infl_cumhaz[, strata_indivs, drop = FALSE] +
          dM_infl
        infl_cumhaz <- infl_cumhaz +
          (dM_infl %*%
             diag(1 / (ipw[strata_indivs] * n), length(strata_indivs)) %*%
             npmle_fit[["infl_ainv"]][strata_indivs, , drop = FALSE])

        bhaz[[i+1]][, "var_logcumhaz"] <- rowMeans(infl_cumhaz^2) /
          (bhaz[[i+1]][, "cumhaz"]^2 * n)
        bhaz[[i+1]][, colnames(X[, !coef_na, drop = FALSE])] <-
          diag(1 / (bhaz[[i+1]][, "cumhaz", drop = TRUE] * n),
               length(bhaz[[i+1]][, "cumhaz", drop = TRUE])) %*%
          infl_cumhaz %*%
          t(var_components) %*%
          cox_fit[["var"]][!coef_na, !coef_na, drop = FALSE]
      }
    }
  }

  # residuals
  mresid <- stats::residuals(cox_fit, type = "martingale", weighted = TRUE)
  rsums <- rowsum(cbind(mresid,
                        y_coxph[, ncol(y_coxph), drop = TRUE] * outer_wt,
                        y_coxph[, ncol(y_coxph), drop = TRUE]),
                  vals$id[vals$sorted_cc]) # returned in original data order
  subj_mresid <- rsums[, 1, drop = TRUE]
  subj_wt <- ifelse(rsums[, 2, drop = TRUE] == 0, 1, rsums[, 2, drop = TRUE])
  subj_event <- rsums[, 3, drop = TRUE]
  subj_dresid <- sign(subj_mresid) *
    sqrt(-2 * (subj_mresid +
                 ifelse(subj_event > 0,
                        subj_wt * log(1 - subj_mresid / subj_wt),
                        0)))

  unsort_indx <- order(vals$sort_indx)
  res <- list(n = n, n_miss = n_miss, n_event = n_event, coef = ests,
              var_matrix = var_matrix, robust_se = robust_se,
              offset = ofst[unsort_indx], ipw = expanded_ipw[unsort_indx],
              time_wt = expanded_time_wt[unsort_indx],
              cc_indx = vals$complete_cases,
              id_indx = vals$id[vals$complete_cases],
              sort_indx = vals$sort_indx, s = s,
              martingale_resid = mresid[unsort_indx],
              deviance_resid = subj_dresid, basehaz = bhaz,
              timerange = test_window, eventrange = eventrange,
              wt_name = wt_name, cox_fit = cox_fit,
              strata_summary = strata_summary, call_ = fn_call)
  class(res) <- "coxph_indtrunc"
  res
}


#' Plot survival or hazard estimates from a Cox model fit
#'
#' @description
#' Given a fitted Cox model from [coxph_indtrunc()], plot the conditional survival
#' or cumulative hazard estimates along with pointwise confidence intervals.
#' The baseline hazard and its standard errors are computed by [plot_coxsurv()]
#' if they were not included in the `coxph_indtrunc` fit.
#'
#' @param object A `coxph_indtrunc` object. See [coxph_indtrunc()].
#' @param newdata A data frame with the covariate values for the plot.
#'   The original model formula is applied to extract predictors.
#'   Rows that are duplicated or have missing values are discarded.
#' @param ci Whether to plot pointwise confidence intervals (default `TRUE`).
#'   Will re-fit the model if `basehaz` or its standard errors were not computed.
#' @param level Confidence level for the pointwise confidence intervals
#'   (default 0.95).
#' @param ctype Type of confidence interval.
#'   * `"plain"`: linear confidence from a normal approximation.
#'   * `"log"`: back-transformed confidence intervals from the natural log scale.
#'     Guaranteed to be non-negative.
#'   * `"log-log"` (the default): for the survival function, back-transformed
#'     confidence intervals from the (natural) log minus log scale, which are
#'     guaranteed to be between zero and one.
#'     For the cumulative hazard, the same confidence intervals as `"log"` are
#'     used.
#' @param target Which quantity to plot.
#'   * `"survival"` (the default): the conditional survival function.
#'   * `"cumhaz"`: the conditional cumulative hazard.
#'
#' @details
#' In order to construct the linear predictors at the requested covariate values,
#' the design matrix obtained from evaluating the original model formula on
#' `newdata` should have the same column names as the design matrix for the
#' fitted model.
#' Issues can occur for categorical variables that were converted to factors by
#' [coxph_indtrunc()] for estimation.
#' If an error occurs, check that each categorical variable in `newdata` is a
#' factor variable that is not missing any levels used to fit the model.
#'
#' @returns
#' After displaying the plot, a list with two elements is returned invisibly.
#' The first element `ests` is a data frame with the columns `time`, `est`,
#' `lower`, and `upper`, containing the point estimates and their computed
#' confidence intervals.
#' The second element `plot` is the plotted object returned by [ggplot2::ggplot()].
#' Lines are colored based on the covariate value.
#'
#' @export
#'
#' @seealso [coxph_indtrunc()]
#'
#' @examples
#' fit <- coxph_indtrunc(Survdt(incu, ltrunc, rtrunc) ~ age_gp,
#'                       data = aids,
#'                       basehaz = TRUE)
#' plot_coxsurv(fit, newdata = aids[c(1,40), ])
#'
#' plot_coxsurv(fit, newdata = aids[c(1,40), ],
#'              target = "cumhaz")
plot_coxsurv <- function(object, newdata, ci = TRUE, level = 0.95,
                         ctype = c("log-log", "log", "plain"),
                         target = c("survival", "cumhaz")){
  # check input
  if (!inherits(object, "coxph_indtrunc"))
    stop("Model object must be of class coxph_indtrunc.")
  stopifnot(is.matrix(newdata) || is.data.frame(newdata))
  if (is.matrix(newdata)) newdata <- as.data.frame(newdata)
  ctype <- match.arg(ctype)
  target <- match.arg(target)
  stopifnot(is.logical(ci), is.vector(ci), all(!is.na(ci)),
            is.numeric(level), is.vector(level), all(!is.na(level)))
  stopifnot(nrow(newdata) > 0, length(ci) == 1, length(level) == 1)
  stopifnot(level > 0, level < 1)
  q <- stats::qnorm(1-(1-level)/2)

  # Survdt2 response not supported
  if (ncol(object[["cox_fit"]]$y) > 2)
    stop("plot_coxsurv does not currently support time-varying covariates")

  # check for basehaz
  if (is.null(object[["basehaz"]]) ||
      ((ci == TRUE) &&
       (!("var_logcumhaz" %in% colnames(object[["basehaz"]][[2]])) ||
        (any(is.na(object[["basehaz"]][[2]][, "var_logcumhaz"])))))){

    new_call <- object[["call_"]]
    new_call$basehaz <- TRUE
    new_call$se <- ci
    message(paste0("Obtaining the baseline hazard ",
                   ifelse(ci == TRUE, "with standard errors ", ""),
                   "by calling:\n ",
                   paste0(format(new_call), collapse = "\n")))
    object$basehaz <- eval.parent(new_call, n = 1)[["basehaz"]]
  }

  # evaluate model formulas with newdata
  newdata <- newdata[!duplicated(newdata), , drop = FALSE]
  formula <- eval.parent(object[["call_"]]$formula, n = 1)
  strata_formula <- eval.parent(object[["call_"]]$strata_formula, n = 1)
  simple_labels <- eval.parent(object[["call_"]]$simple_labels, n = 1)
  if (is.null(simple_labels))
    simple_labels <- formals(coxph_indtrunc)[["simple_labels"]]

  # get strata
  if (!is.null(strata_formula)){
    strata_vars <- stats::model.frame(stats::update(strata_formula, ~.-1),
                                      data = newdata, na.action = stats::na.pass)
    complete_cases_s <- rowSums(is.na(strata_vars)) == 0
    sdf_names <- c(apply(strata_vars[1, , drop = FALSE],
                         MARGIN = 1, FUN = names))
    if (simple_labels == TRUE)
      factor_names <- rep("", length(sdf_names))
    else
      factor_names <- sdf_names
    s <- apply(strata_vars, MARGIN = 1, FUN = function(x) {
      paste0(factor_names, x, sep = "", collapse = ", ")
    })
  }else{
    complete_cases_s <- TRUE
    s <- NULL
  }

  # main formula
  main_terms <-
    stats::delete.response(stats::terms(stats::update(formula, ~.+1)))
  mf <- stats::model.frame(main_terms,
                           data = newdata,
                           na.action = stats::na.pass)

  if (sum(sapply(mf, FUN = function(x) any(is.infinite(x)))) > 0)
    stop("Infinite values occurred when evaluating 'formula' in newdata")
  ofst <- stats::model.offset(mf)
  X <- stats::model.matrix(mf, data = newdata)
  X <- X[, -which(attr(X, "assign") == 0), drop = FALSE]
  coef_na <- is.na(object[["coef"]][, "coef", drop = TRUE])
  pred_match_indx <- match(rownames(object[["coef"]])[!coef_na],
                           colnames(X),
                           nomatch = NA)
  if (all(is.na(pred_match_indx)))
    stop("Names of predictors constructed from newdata did not match ",
         "any identifiable parameters")
  X <- X[, pred_match_indx[!is.na(pred_match_indx)], drop = FALSE]
  if (ncol(X) != sum(!coef_na))
    stop("Missing some parameter names in predictors constructed from newdata")

  # missing values
  complete_cases <- (rowSums(is.na(X)) == 0) & complete_cases_s
  if (!is.null(ofst)) complete_cases <- complete_cases & !is.na(ofst)
  n <- sum(complete_cases)
  if (n == 0) stop("All observations in newdata produced missing values")
  X <- X[complete_cases, , drop = FALSE]
  if (is.null(ofst))
    ofst <- rep(0, n)
  else
    ofst <- ofst[complete_cases]

  if (!is.null(s)) {
    s <- s[complete_cases]
    s <- factor(s, levels = levels(object[["s"]]))
    if (any(is.na(s)))
      stop("newdata contained stratum values not found in original fit")
    unique_indx <- !duplicated(as.data.frame(cbind(X, ofst, s)))
    s <- s[unique_indx]

    n_strata <- length(object[["basehaz"]][[1]])
    strata_sizes <- sapply(object[["basehaz"]][[1]],
                           FUN = function(sj) sum(s == sj))
    strata_etimes <- sapply(object[["basehaz"]][-1], FUN = nrow)
    total_times <- sum(strata_sizes * strata_etimes)
  } else{
    unique_indx <- !duplicated(as.data.frame(cbind(X, ofst)))
    s <- rep(object[["basehaz"]][[1]][1], sum(unique_indx))
    n_strata <- 1
    strata_etimes <- nrow(object[["basehaz"]][[2]])
    total_times <- sum(unique_indx) * strata_etimes
  }
  X <- X[unique_indx, , drop = FALSE]
  ofst <- ofst[unique_indx]
  n <- nrow(X)

  # get labels for plot
  rows <- which(complete_cases)[unique_indx]
  maindf_names <- c(apply(mf[1, , drop = FALSE], MARGIN = 1, FUN = names))
  pltlabels <- c(apply(mf[rows, , drop = FALSE], MARGIN = 1,
                       FUN = function(z) {
                         paste0(maindf_names, z, sep = "", collapse = ", ")
                       }))
  if (n_strata > 1){
    keep_sdfcol <- !(sdf_names %in% maindf_names)
    if (any(keep_sdfcol)){
      pltlabels <- paste0(pltlabels, ", ",
                          c(apply(strata_vars[rows, , drop = FALSE], MARGIN = 1,
                                  FUN = function(z) {
                                    paste0(sdf_names[keep_sdfcol], z[keep_sdfcol],
                                           sep = "", collapse = ", ")
                                  })))
    }
  }

  # get bhaz & std errors
  exps <- exp(c(X %*% object[["coef"]][!coef_na, "coef", drop = TRUE]) + ofst)
  event_time <- rep(NA, total_times + n)
  cumhaz <- rep(NA, total_times + n)
  se_logcumhaz <- rep(NA, total_times + n)
  subject <- rep(NA, total_times + n)
  label2 <- rep(NA, total_times + n)
  pos <- 0
  for (j in seq_len(n_strata)){
    strata_indx <- s == object[["basehaz"]][[1]][j]
    n_in_s <- sum(strata_indx)
    if (n_in_s == 0) next
    posindx <- pos + seq_len((strata_etimes[j] + 1) * n_in_s)
    pos <- pos + (strata_etimes[j] + 1) * n_in_s

    event_time[posindx] <- rep(c(object[["basehaz"]][[j + 1]][1,
                                                              "time",
                                                              drop = TRUE],
                                 object[["basehaz"]][[j + 1]][,
                                                              "time",
                                                              drop = TRUE]),
                               n_in_s)
    cumhaz[posindx] <- rep(c(0,
                             object[["basehaz"]][[j + 1]][,
                                                          "cumhaz",
                                                          drop = TRUE]),
                           n_in_s) *
      rep(exps[strata_indx], each = strata_etimes[j] + 1)
    subject[posindx] <- rep(pltlabels[strata_indx], each = strata_etimes[j] + 1)
    label2[posindx] <- rep(which(strata_indx), each = strata_etimes[j] + 1)

    if (ci == TRUE){
      var_matrix <- object[["var_matrix"]][!coef_na, !coef_na, drop = FALSE]
      se_logcumhaz[posindx] <-
        c(apply(X[strata_indx, , drop = FALSE],
                MARGIN = 1,
                FUN = function(z) {
                  sqrt(c(t(z) %*% var_matrix %*% z) +
                         c(NA,
                           object[["basehaz"]][[j + 1]][,
                                                        "var_logcumhaz",
                                                        drop = TRUE]) +
                         c(NA,
                           2 *
                             object[["basehaz"]][[j + 1]][,
                                                          colnames(X),
                                                          drop = FALSE] %*%
                             z))
                }))
    }
  }

  # generate plot
  ci_transform <- ifelse(ctype == "plain", "linear", ctype)
  if (target == "cumhaz"){
    ylab <- "Cumulative hazard"
    ests <- data.frame(x_number = label2,
                       x_label = subject,
                       time = event_time,
                       est = cumhaz,
                       se = se_logcumhaz * cumhaz,
                       lower = rep(NA, length(event_time)),
                       upper = rep(NA, length(event_time)))

    ci_transform <- ifelse(ctype == "log-log", "log", ci_transform)
  }else{
    ylab <- "Survival function"
    ests <- data.frame(x_number = label2,
                       x_label = subject,
                       time = event_time,
                       est = exp(-cumhaz),
                       se = se_logcumhaz * cumhaz * exp(-cumhaz),
                       lower = rep(NA, length(event_time)),
                       upper = rep(NA, length(event_time)))
  }


  if (ci == TRUE){
    if (ci_transform == "linear"){
      ests$lower <- ests$est - q * ests$se
      ests$upper <- ests$est + q * ests$se
    }else if (ci_transform == "log"){
      ests$lower <- ests$est * exp(-q * ests$se / abs(ests$est))
      ests$upper <- ests$est * exp(q * ests$se / abs(ests$est))
    }else if (ci_transform == "log-log"){
      ests$upper <- ests$est^exp(-q * ests$se /
                                   abs(ests$est * log(abs(ests$est))))
      ests$lower <- ests$est^exp(q * ests$se /
                                   abs(ests$est * log(abs(ests$est))))
    }

    startindx <- which(cumhaz == 0)
    if (target == "survival"){
      ests$lower[startindx] <- ests$lower[startindx + 1]
      ests$upper[startindx] <- pmax(1, ests$upper[startindx + 1])
    }else{
      ests$lower[startindx] <- pmin(0, ests$lower[startindx + 1])
      ests$upper[startindx] <- ests$upper[startindx + 1]
    }
  }

  plt <- with(ests,
              ggplot2::ggplot(ests,
                              ggplot2::aes(x = time,
                                           color = x_label,
                                           group = x_number)))
  if (ci == TRUE){
    plt <- with(ests,
                plt +
                  ggplot2::geom_step(ggplot2::aes(y = lower),
                                     linetype = "dashed",
                                     direction = "hv") +
                  ggplot2::geom_step(ggplot2::aes(y = upper),
                                     linetype = "dashed",
                                     direction = "hv"))
  }
  plt <- with(ests,
              plt +
                ggplot2::geom_step(ggplot2::aes(y = est), direction = "hv",
                                   key_glyph = "rect") +
                ggplot2::labs(x = "Time", y = ylab, color = "") +
                ggplot2::scale_color_viridis_d(begin = 0, end = 0.8) +
                ggplot2::coord_cartesian(xlim = c(min(ests$time) / 1.01,
                                                  max(ests$time) * 1.01),
                                         ylim = c(0, max(ests$est) * 1.01),
                                         expand = FALSE) +
                ggplot2::theme_classic() +
                ggplot2::theme(legend.position = "bottom",
                               legend.key.size = ggplot2::unit(0.9, "lines")))

  print(plt)
  invisible(list(ests = ests, plot = plt))
}

#' @rdname coxph_indtrunc
#' @export
vcov.coxph_indtrunc <- function(object, ...){
  if (is.null(object[["var_matrix"]]))
    stop("Standard errors not available for fitted model. Refit with se=TRUE.")
  object[["var_matrix"]]
}

#' @rdname coxph_indtrunc
#' @export
coef.coxph_indtrunc <- function(object, ...){
  coefs <- object[["coef"]][, "coef", drop = TRUE]
  names(coefs) <- rownames(object[["coef"]])
  coefs
}

#' @rdname coxph_indtrunc
#' @export
confint.coxph_indtrunc <- function(object, parm, level = 0.95, ...){
  if (missing(parm)) parm <- rownames(object[["coef"]])
  stopifnot(is.vector(parm), is.numeric(parm) | is.character(parm),
            is.vector(level), is.numeric(level), length(level) == 1,
            all(!is.na(level)), all(level < 1), all(level > 0))
  if (is.numeric(parm) && any((parm < 1) | (parm > nrow(object[["coef"]]))))
    stop("Numerical parm contains indices that are out of bounds")
  if (is.character(parm) && !all(parm %in% rownames(object[["coef"]])))
    stop("Character vector parm contains names that do not match coefficients")
  parm <- unique(parm)
  p <- 1 - ((1 - level) / 2)
  q <- stats::qnorm(p)
  ci <- cbind(object[["coef"]][parm, "coef", drop = TRUE] -
                (q * object[["coef"]][parm, "se", drop = TRUE]),
              object[["coef"]][parm, "coef", drop = TRUE] +
                (q * object[["coef"]][parm, "se", drop = TRUE]))
  rownames(ci) <- rownames(object[["coef"]][parm, , drop = FALSE])
  colnames(ci) <- paste0(c((1 - p), p) * 100, "%")
  ci
}

#' @rdname coxph_indtrunc
#' @export
residuals.coxph_indtrunc <- function(object, type = c("martingale", "deviance"),
                                     per_subject = type == "deviance", ...){
  type <- match.arg(type)
  per_subject <- per_subject
  stopifnot(is.logical(per_subject), is.vector(per_subject),
            length(per_subject) == 1, all(!is.na(per_subject)))
  if (type == "martingale"){
    if (per_subject == TRUE)
      resids <- c(rowsum(object[["martingale_resid"]], object[["id_indx"]]))
    else
      resids <- object[["martingale_resid"]]

  }else{
    if (per_subject == TRUE)
      resids <- object[["deviance_resid"]]
    else{
      y <- object[["cox_fit"]][["y"]]
      event_indx <- y[order(object[["sort_indx"]]) , ncol(y), drop = TRUE]
      obs_mresid <- object[["martingale_resid"]]
      obs_wt <- object[["cox_fit"]][["time_wt"]] * event_indx
      if (any(obs_wt == 0)) obs_wt[obs_wt == 0] <- 1
      resids <- sign(obs_mresid) *
        sqrt(-2 * (obs_mresid +
                     ifelse(event_indx > 0,
                            obs_wt * log(1 - obs_mresid / obs_wt),
                            0)))
    }
  }
  resids
}

#' @rdname coxph_indtrunc
#' @export
print.coxph_indtrunc <- function(x, ...){
  timewt_name <- ifelse(x[["wt_name"]] == "W-1", "none", x[["wt_name"]])
  cat(paste0("\n",
             "IPW Cox model under quasi-independent truncation.",
             "\n",
             "\n",
             "  Call: ", paste0(format(x[["call_"]]),
                                collapse = paste0("\n",
                                                  paste0(rep(" ", 8),
                                                         collapse = ""))),
             "\n",
             "\n",
             "  Model was fit on ", x[["n"]], " complete samples (",
             x[["n_miss"]], " dropped due to missingness).",
             "\n",
             "  Min event time: ", x[["eventrange"]][1],
             "\n",
             "  Max event time: ", x[["eventrange"]][2],
             "\n",
             "  ", x[["n_event"]], " observed events in the score function, ",
                  "time range [", x[["timerange"]][1], ", ",
                                  x[["timerange"]][2], "].",
             "\n",
             "  Inverse probability of non-truncation weights estimated ",
                "nonparametrically.",
             "\n",
             "  Time-varying weight component: ", timewt_name,
             "\n",
             "\n",
             "Model estimates:",
             "\n"))
  print(x[["coef"]])
  cat("\n")

  if (!is.null(x[["strata_summary"]])){
    sort_indx <- order(x[["strata_summary"]][, "strata_size", drop = TRUE],
                       decreasing = TRUE)
    strata_summary <- x[["strata_summary"]][sort_indx,
                                            c("strata",
                                              "strata_size",
                                              "unique_event_times"),
                                            drop = FALSE]
    colnames(strata_summary) <- c("stratum", "size", "unique event times")

    cat(paste0("Summary of the ", nrow(x[["strata_summary"]]), " strata:",
               "\n"))

    if (nrow(x[["strata_summary"]]) > 10){
      print(strata_summary[1:10, , drop = FALSE],
            row.names = FALSE)
      cat("  ...")
      cat("\n\n")
    }else{
      print(strata_summary, row.names = FALSE)
      cat("\n")
    }
  }

  invisible(x)
}

#' Sensitivity analysis for positivity violations in a Cox model with
#' quasi-independent double truncation.
#'
#' @description
#' Computes inverse probability weighted sensitivity analysis estimators for the
#' Cox model using nonparametric weights (see [npmle()]).
#' Given a grid of values for the potential amount of baseline probability mass
#' lost due to zero sampling probability, shows how a positivity violation could
#' influence the results of a fitted Cox model.
#'
#' @param formula A model formula with a `Survdt` response.
#'   Subjects with missing values are discarded.
#'   Time-varying covariates are not supported.
#'   See [Survdt()].
#' @param data An optional data frame used to evaluate `formula` and
#'   `strata_formula`.
#' @param strata_formula An optional one-sided formula with stratification
#'   variables on the right hand side, for stratified Cox models.
#' @param trunc_mass A numerical vector containing the grid of truncated
#'   baseline mass values to use in the sensitivity analysis.
#' @param se Whether to calculate (robust) standard errors (default `TRUE`).
#' @param level Confidence level for constructing sensitivity intervals
#'   (default 0.95).
#' @param robust_infl A square matrix of influence values for the event weight,
#'   with one column per individual and one row per follow up time, both in the
#'   order that they appear in the data.
#'   Only used if `ipw_type` is a function and `se`=`TRUE`.
#' @param max_iter Maximum number of iterations before non-convergence
#'   is declared at a given truncated mass (default 100).
#' @param error Tolerance for determining numerical convergence of coefficient
#'   estimates (default `1e-7`).
#' @param beta_init Optional numerical vector with initial values for each
#'   coefficient estimate.
#' @param center_covariates Whether to center the covariates before the
#'   sensitivity analysis (default `FALSE`).
#'   Covariates with values of 0/1 are not centered.
#' @param x A `coxph_pos_indtrunc` object.
#' @param hazratio Whether to plot the hazard ratios exp(beta) and their
#'   sensitivity intervals instead of the regression coefficients beta (default
#'   `FALSE`).
#' @param dropNA Whether to remove truncation masses where the estimator did not
#'   converge, resulting in missing values, from the plot (default `TRUE`).
#'
#' @inheritParams coxph_indtrunc
#'
#' @details
#' Consider a sample of doubly truncated data with event times \eqn{T_i},
#' truncation times \eqn{L_i} and \eqn{R_i}, and covariates \eqn{X_i}.
#'
#' A positivity violation occurs when the sampling probability
#' \eqn{P(L<T \leq R|T=t)} is zero for \eqn{t\leq a} and \eqn{t > b}, where the
#' observation bounds \eqn{a}
#' and \eqn{b} lie within the event time support.
#' The truncated baseline probability mass is defined as the amount of probability
#' mass right-truncated by the positivity violation at the baseline covariate
#' level, conditional on not being left-truncated by the positivity violation,
#' i.e. \eqn{P(T> b|X=0,T > a)}.
#' Note that the interpretation of this truncated mass depends on how the baseline
#' covariate value is defined, e.g. whether the covariates are centered.
#'
#' At each specified truncated mass, [positivity_sens_indtrunc()] fits a Cox model
#' that accounts for the positivity violation through an adjustment to the
#' at-risk process (see Vazquez and Xie, 2025).
#' This modified at-risk process is inserted into the weighted score function from
#' [coxph_indtrunc()].
#'
#' The end result of the sensitivity analysis is a series of adjusted coefficient
#' estimators at the grid of truncated masses, which show how sensitive the Cox
#' model estimates are to potential positivity violations, as well as sensitivity
#' intervals giving a range of plausible values for the coefficients at the
#' user-supplied confidence `level` for each truncated mass.
#'
#' In general, the sensitivity analysis estimators may not converge at
#' some truncated mass values.
#' Usually this is due to large weights, near-zero risk scores
#' \eqn{\exp(X_i^T\beta)}, and/or truncated masses too close to one.
#' Some options to consider are:
#' * Changing the baseline covariate value by shifting the covariates or reordering
#'   factor levels.
#' * Using time-varying weights (see `ipw_type`).
#' * Restricting the time window in the score function through `restr_times`
#'   (see below).
#' * Altering the range and step size of `trunc_mass`.
#'
#' @inheritSection coxph_indtrunc Stratified Cox model
#'
#' @inheritSection coxph_indtrunc Restricted time window in score function
#'
#' @returns
#' [positivity_sens_indtrunc()] returns an object of class `coxph_pos_indtrunc`,
#' i.e. a list with the following components:
#' \item{n}{The complete case sample size.}
#' \item{n_miss}{The number of individuals removed due to missing data.}
#' \item{n_event}{The number of observed events used in the score function.}
#' \item{level}{The user-specified confidence level for constructing sensitivity
#'   intervals.}
#' \item{trunc_mass}{The truncated baseline masses supplied in `trunc_mass`
#'   (sorted).
#'   The baseline covariate values are provided in `col_centerings`.}
#' \item{coef}{A matrix of coefficient estimates with `j`th column corresponding
#'   to `trunc_mass[j]`.
#'   Missing values indicate non-convergence or collinearity.}
#' \item{se}{A matrix of standard errors with `j`th column corresponding to
#'   `trunc_mass[j]`.}
#' \item{lower,upper}{Matrices containing the lower and upper sensitivity
#'   interval endpoints, with `j`th column corresponding to `trunc_mass[j]`.}
#' \item{timerange}{A vector with the minimum and maximum times used in the
#'   score function.}
#' \item{eventrange}{A vector with the unrestricted range of observed event times.}
#' \item{wt_name}{A string giving the name of the weights used in `ipw_type`.}
#' \item{col_centerings}{A numerical vector with the values used to center the
#'   columns of the design matrix if `center_covariates`=`TRUE`.
#'   Otherwise, a vector of zeros.}
#' \item{strata_summary}{Some summary statistics for each stratum, computed prior
#'   to applying `restr_times`.}
#' \item{call_}{The call to [positivity_sens_indtrunc()].}
#'
#' The `print` method for `coxph_pos_indtrunc` objects displays a summary of the
#' sensitivity analysis results, while the `plot` method returns a visualization
#' of the estimators and their sensitivity intervals across different truncated
#' mass values, generated through [ggplot2::ggplot()].
#'
#' @export
#'
#' @seealso
#' * [coxph_indtrunc()]: fitting a Cox model under quasi-independent truncation.
#' * `vignette("cox-regression-quasiindep", package = "survdt")`: an illustration
#'   of this sensitivity analysis applied to the [aids] data, if vignettes were
#'   built when installing [survdt].
#'
#' @examples
#' sensfit <- positivity_sens_indtrunc(Survdt(incu, ltrunc, rtrunc) ~ age_gp,
#'                                     data = aids,
#'                                     trunc_mass = seq(0, 0.11, by = 0.02),
#'                                     ipw_type = "W-asurv")
#' sensfit
#' plot(sensfit)
#'
#' @references
#' Vazquez, O., & Xie, S. X. (2025). Robust inverse probability weighted
#' estimators for doubly truncated Cox regression with closed-form standard
#' errors. *Lifetime Data Analysis*, **31**(2), 364–-393.
positivity_sens_indtrunc <- function(formula, data = NULL, strata_formula = NULL,
                                     trunc_mass = seq(0, 0.9, by = 0.1),
                                     ipw_type = c("W-asurv", "W-surv", "W-1", "W-a"),
                                     se = TRUE, level = 0.95, restr_times = NULL,
                                     npmle_fit = NULL,
                                     robust_infl = NULL, max_iter = 100,
                                     error = 1e-7, beta_init = NULL,
                                     center_covariates = FALSE,
                                     simple_labels = FALSE){
  fn_call <- match.call()

  # check input
  if (missing(formula))
    stop('\n  argument "formula" is missing, with no default')
  stopifnot(inherits(formula, "formula"))
  if (length(formula) < 3) stop("The formula is missing a LHS or a RHS")
  if (!is.null(data) && is.matrix(data)) data <- as.data.frame(data)
  stopifnot(is.null(data) || is.data.frame(data))
  if (!is.null(strata_formula)){
    stopifnot(inherits(strata_formula, "formula"))
    if (length(strata_formula) != 2)
      stop("The strata formula should have the form: ~var1 + var2 + ...")
    strata_vars <- stats::model.frame(stats::update(strata_formula, ~.-1),
                                      data = data,
                                      na.action = stats::na.pass)
    if (attr(attr(strata_vars, "terms"), "response") != 0)
      stop("The strata formula should not have a response variable")
    if (length(strata_vars) == 0)
      stop("strata formula is missing stratification variables on RHS")
    complete_cases_s <- rowSums(is.na(strata_vars)) == 0
    factor_names <- c(apply(strata_vars[1, , drop = FALSE],
                            MARGIN = 1, FUN = names))
    stopifnot(is.logical(simple_labels), is.vector(simple_labels),
              length(simple_labels) == 1, all(!is.na(simple_labels)))
    if (simple_labels == TRUE) factor_names <- rep("", length(factor_names))
    s <- apply(strata_vars, MARGIN = 1, FUN = function(x) {
      paste0(factor_names, x, sep = "", collapse = ", ")
    })
  }else{
    complete_cases_s <- TRUE
    s <- NULL
  }

  mf <- stats::model.frame(stats::update(formula, ~.+1), data = data,
                           na.action = stats::na.pass)
  if (attr(attr(mf, "terms"), "response") == 0)
    stop("formula is missing a response")
  if (sum(sapply(mf, FUN = function(x) any(is.infinite(x)))) > 0)
    stop("Infinite values occurred when evaluating 'formula'")
  y <- mf[[attr(attr(mf, "terms"), "response")]]
  ofst <- stats::model.offset(mf)
  if (is.Survdt2(y))
    stop("Response should be a Survdt object with one observation",
         " per subject\n (not Survdt2, time-varying covariates are",
         " not supported).\n See ?survdt::survdt2to1")
  stopifnot(is.Survdt(y))
  if (!is.null(s) && (length(s) != length(y)))
    stop("Strata formula terms and main formula terms have different lengths")


  if (is.null(restr_times)) restr_times <- c(NA_real_, NA_real_)
  stopifnot(is.vector(trunc_mass), is.numeric(trunc_mass),
            all(!is.na(trunc_mass)), all(trunc_mass < 1), all(trunc_mass>=0),
            length(trunc_mass) != 0,
            (is.vector(ipw_type) && is.character(ipw_type)) ||
              is.function(ipw_type),
            is.logical(se), is.vector(se), length(se) == 1, all(!is.na(se)),
            is.numeric(level), is.vector(level), length(level) == 1,
            all(!is.na(level)), is.numeric(restr_times), is.vector(restr_times),
            is.null(npmle_fit) || is.npmle(npmle_fit),
            is.null(robust_infl) || is.matrix(robust_infl),
            is.numeric(max_iter), is.vector(max_iter), length(max_iter) == 1,
            all(!is.na(max_iter)), max_iter > 1,
            is.numeric(error), is.vector(error), length(error) == 1,
            all(!is.na(error)), error > 0,
            is.null(beta_init) ||
              (is.numeric(beta_init) && is.vector(beta_init)),
            is.logical(center_covariates), is.vector(center_covariates),
            length(center_covariates) == 1, all(!is.na(center_covariates)))
  stopifnot(length(restr_times) == 2, 0 < level, level < 1)
  trunc_mass <- sort(trunc_mass)
  max_iter <- floor(max_iter)

  # missing values
  n_init <- length(y)
  complete_cases <- (rowSums(is.na(mf)) == 0) & complete_cases_s
  n_miss <- sum(!complete_cases)
  n <- sum(complete_cases)
  if (n == 0) stop("All individuals have missing values")
  X <- stats::model.matrix(mf, data = data)
  X <- X[, -which(attr(X, "assign") == 0), drop = FALSE]
  if (ncol(X) == 0) stop("No covariates provided in RHS of formula")

  if (!is.null(s)) {
    if (length(unique(s)) != length(unique(s[complete_cases])))
      warning("Dropping strata with no complete cases")
    s <- s[complete_cases]
    s <- factor(s, levels = unique(s[order(c(strata_vars[complete_cases, 1,
                                                         drop = TRUE])[1:n])]))
    sort_indx <- order(s, y[complete_cases, "time", drop = TRUE])
    s <- s[sort_indx]
    strata_vals <- unique(s)
    n_strata <- length(strata_vals)
    if (n_strata < 2)
      stop(paste0("Only one stratum in complete cases.\n Either update the",
                  " strata formula or drop it altogether"))

    n_unique_strata <- stats::aggregate(y[, "time", drop = TRUE] ~ s,
                                        FUN = function(x) length(unique(x)))
    small_strata <- n_unique_strata[n_unique_strata[, 2, drop = TRUE] < 2, 1,
                                    drop = TRUE]
    if (length(n_unique_strata) > 5)
      n_unique_strata <- c(as.character(n_unique_strata[1:5]), "...")
    if (length(small_strata) > 0)
      warning(paste0("Less than two unique event times found in stratum: ",
                     paste0(small_strata, collapse = "; ")))
    strata_summary <- data.frame(strata = n_unique_strata[, 1, drop = TRUE],
                                 strata_size =
                                   sapply(n_unique_strata[, 1, drop = TRUE],
                                          FUN = function(x) sum(s == x)),
                                 unique_event_times =
                                   n_unique_strata[, 2, drop = TRUE])
  }else{
    sort_indx <- order(y[complete_cases, "time", drop = TRUE])
    strata_vals <- NULL
    n_strata <- 1
    strata_summary <- NULL
  }
  sorted_cc <- which(complete_cases)[sort_indx]
  y <- y[sorted_cc, ]
  X <- X[sorted_cc, , drop = FALSE]
  if (is.null(ofst)){
    ofst <- rep(0, n)
  }else{
    ofst <- ofst[sorted_cc]
  }

  # time window
  eventrange <- range(y[, "time", drop = TRUE])
  time_lims <- eventrange
  if (!is.na(restr_times[1]) && (restr_times[1] >= time_lims[2]))
    stop("restr_times[1] must be less than the maximum observed event time")
  if (!is.na(restr_times[2]) && (restr_times[2] <= time_lims[1]))
    stop("restr_times[2] must be greater than the minimum observed event time")
  if (all(!is.na(restr_times)) && (restr_times[1] >= restr_times[2]))
    stop("Must have restr_times[1]<restr_times[2]")
  test_window <- ifelse(is.na(restr_times), time_lims,
                        c(max(restr_times[1], time_lims[1]),
                          min(restr_times[2], time_lims[2])))
  ltrunc_indx <- y[, "time", drop = TRUE] < test_window[1]
  rcens_indx <- test_window[2] < y[, "time", drop = TRUE]
  n_events <- sum(!ltrunc_indx & !rcens_indx)
  if (n_events == 0) stop("No observed event times in restricted window")

  # initial coefficient value
  if (is.null(beta_init)) beta_init <- rep(0, ncol(X))
  stopifnot(all(!is.na(beta_init)), all(!is.infinite(beta_init)))
  if (length(beta_init) != ncol(X))
    stop(paste0("Length of beta_init does not match the number of predictors (",
                ncol(X), ")"))

  # check custom weights
  if (is.character(ipw_type)){
    ipw_type <- match.arg(ipw_type)
    wt_name <- ipw_type
  }else{
    time_wt <- ipw_type(y[, "time", drop = TRUE])

    if (length(time_wt) != length(y))
      stop("Length of output from function ipw_type does not match input")
    if (any(is.na(time_wt) | is.infinite(time_wt)))
      stop("Missing or infinite values found in output of function ipw_type")
    if (any(time_wt <= 0))
      stop("Function ipw_type did not return strictly positive values")

    if (se == TRUE){
      if (is.null(robust_infl))
        stop("Must supply robust_infl for custom event weights when se=TRUE")
      if ((nrow(robust_infl) != n_init) || (ncol(robust_infl) != n_init)){
        stop("Matrix robust_infl should have one row & column per individual")
      }
      robust_infl <- robust_infl[sorted_cc, sorted_cc, drop = FALSE]
      if (any(is.na(robust_infl) | is.infinite(robust_infl)))
        stop("Missing or infinite values found in robust_infl")
    }
    wt_name <- paste0("user-supplied function ", format(substitute(ipw_type)))
  }

  # npmle
  if (!is.null(npmle_fit)){
    stopifnot(is.npmle(npmle_fit))
    if (nrow(npmle_fit[["masses"]]) != n)
      stop(paste0("npmle_fit sample size (", nrow(npmle_fit[["masses"]]), ") ",
                  "does not match the number of complete cases ",
                  "in the model frame (", n, ")"))
    if (any(abs(y[, "time", drop = TRUE] -
                npmle_fit[["masses"]]$time[sort_indx]) >
            1e-9))
      stop("Follow up times in formula response and npmle_fit do not match")

    if (se == FALSE){
      if (is.null(npmle_fit[["ainv"]])){
        npmle_fit <- npmle(with(npmle_fit[["masses"]],
                                Survdt(time, ltrunc, rtrunc)),
                           se = FALSE, ainv = TRUE, infl = FALSE, max_iter = 0,
                           init = npmle_fit[["masses"]]$f,
                           convergence_warning = FALSE,
                           check_identifiability = FALSE)
      }
      ipw <- npmle_fit[["ainv"]][sort_indx, "ainv", drop = TRUE]
    }else{
      if (is.null(npmle_fit[["infl_ainv"]])){
        infl_vals <- npmle_infl(npmle_fit[["masses"]], ainv = TRUE,
                                L_F = npmle_fit[["infl_cdf"]],
                                already_sorted = FALSE)
        npmle_fit$infl_cdf <- infl_vals$L_F[, sort_indx, drop = FALSE]
        npmle_fit$infl_ainv <- infl_vals$L_ainv[, sort_indx, drop = FALSE]
        npmle_fit$infl_ainv <- npmle_fit[["infl_ainv"]][sort_indx, ,
                                                        drop = FALSE]
        ipw <- 1 / infl_vals$a[sort_indx]
        infl_vals <- NULL
      }else{
        npmle_fit$infl_cdf <- npmle_fit[["infl_cdf"]][, sort_indx,
                                                      drop = FALSE]
        npmle_fit$infl_ainv <- npmle_fit[["infl_ainv"]][, sort_indx,
                                                        drop = FALSE]
        npmle_fit$infl_ainv <- npmle_fit[["infl_ainv"]][sort_indx, ,
                                                        drop = FALSE]
        ipw <- npmle_fit[["ainv"]][sort_indx, "ainv", drop = TRUE]
      }
    }
  }else{
    npmle_fit <- npmle(y, ainv = TRUE, se = se, infl = se)
    if (!npmle_fit[["converged"]]) stop("Fitted npmle did not converge")
    ipw <- npmle_fit[["ainv"]][, "ainv", drop = TRUE]
  }

  # time-varying weights
  if (is.character(ipw_type)){
    robust_infl <- NULL
    calculate_infl <- se == TRUE
    if (ipw_type == "W-1"){
      time_wt <- rep(1, n)
    }else if (ipw_type == "W-a"){
      time_wt <- 1 / ipw
      if (calculate_infl == TRUE){
        robust_infl <- apply(npmle_fit[["infl_ainv"]], MARGIN = 2,
                             FUN = function(x) {-x / (ipw^2)})
      }
    }else{
      cdf_time_match <- pmin(findInterval(y[, "time", drop = TRUE],
                                          npmle_fit[["cdf"]][,
                                                             "time",
                                                             drop = TRUE],
                                          left.open = TRUE) + 1,
                             nrow(npmle_fit[["cdf"]]))
      lcontsurv <- c(1,
                     1 - npmle_fit[["cdf"]][,
                                            "cdf",
                                            drop = TRUE])[cdf_time_match]

      if (ipw_type == "W-surv"){
        time_wt <- lcontsurv
        if (calculate_infl == TRUE){
          robust_infl <- -rbind(0,
                                npmle_fit[["infl_cdf"]])[cdf_time_match,
                                                         ,
                                                         drop = FALSE]
        }
      }else if (ipw_type == "W-asurv"){
        time_wt <- lcontsurv / ipw
        if (calculate_infl == TRUE){
          robust_infl <- sapply(seq_len(n), FUN = function(i){
            (-npmle_fit[["infl_ainv"]][, i] * lcontsurv / (ipw^2)) -
              (c(0,
                 npmle_fit[["infl_cdf"]][,
                                         i,
                                         drop = TRUE])[cdf_time_match] / ipw)
          })
        }
      }
    }
  }
  time_wt <- time_wt + 1e-7

  # fit sensitivity analysis estimators
  p <- ncol(X)
  col_centerings <- rep(0, p)
  if (center_covariates == TRUE){
    col_centerings <- c(apply(X, MARGIN = 2, FUN = function(x) {
      ifelse(all(x %in% c(0, 1)), 0, mean(x))
    }))
    X <- scale(X, center = col_centerings, scale = FALSE)
  }
  var_names <- colnames(X)
  if (is.null(var_names)) var_names <- paste0("X", 1:p)
  names(col_centerings) <- var_names

  coef_ests <- matrix(NA, ncol = length(trunc_mass), nrow = p)
  stderrs <- matrix(NA, ncol = length(trunc_mass), nrow = p)
  for (j in seq_along(trunc_mass)){
    ests <- tryCatch(coxph_rtrunc(y = y, X = X, s = s, ofst = ofst,
                                  ipw = ipw, time_wt = time_wt,
                                  trunc_mass = trunc_mass[j],
                                  ltrunc_indx = ltrunc_indx,
                                  rcens_indx = rcens_indx,
                                  max_iter = max_iter, error = error,
                                  beta_init = beta_init,
                                  ipw_infl = npmle_fit[["infl_ainv"]],
                                  time_wt_infl = robust_infl),
                     error = function(e) {
                       paste0("\n  Sensitivity analysis failed at ",
                              "trunc_mass ", trunc_mass[j],
                              " with error message ",
                              "\n  '",
                              e$message,
                              "'\n  ",
                              ifelse(is.null(e$call), "",
                                     paste0("in ", e$call, ".\n")),
                              "\n",
                              "See ?positivity_sens_indtrunc for ",
                              "suggestions to improve convergence.")
                     })

    if ((length(ests) == 1) && is.character(ests)){
      warning(ests)
      ests <- NULL
    }

    if (!is.null(ests) && any(!is.na(ests$coef))){
      coef_ests[, j] <- ests$coef
      stderrs[, j] <- ests$se
      beta_init <- ifelse(is.na(ests$coef), 0, ests$coef) # improve convergence
    }else{
      break
    }
  }

  # sensitivity intervals
  q <- stats::qnorm(1 - (1 - level) / 2)
  if (length(trunc_mass) == 1){
    lower <- coef_ests - q * stderrs
    upper <- coef_ests + q * stderrs
  }else{
    lower <- t(apply(coef_ests - q * stderrs, MARGIN = 1, FUN = cummin))
    upper <- t(apply(coef_ests + q * stderrs, MARGIN = 1, FUN = cummax))
  }

  rownames(coef_ests) <- var_names
  rownames(stderrs) <- var_names
  rownames(lower) <- var_names
  rownames(upper) <- var_names
  res <- list(n = n, n_miss = n_miss, n_event = n_events, level = level,
              trunc_mass = trunc_mass, coef = coef_ests, se = stderrs,
              lower = lower, upper = upper, timerange = test_window,
              eventrange = eventrange, wt_name = wt_name,
              col_centerings = col_centerings, strata_summary = strata_summary,
              call_ = fn_call)
  class(res) <- "coxph_pos_indtrunc"
  res
}

# internal function to extract data from formulas
prep_data_for_fit <- function(formula, data = NULL, strata_formula = NULL,
                              simple_labels = FALSE, drop_NA = TRUE,
                              check_post_event = TRUE, min_time = "sample",
                              strata_levels = NULL, ...){
  stopifnot(inherits(formula, "formula"))
  if (length(formula) < 3) stop("The formula is missing a LHS or a RHS")
  if (!is.null(data) && is.matrix(data)) data <- as.data.frame(data)
  stopifnot(is.null(data) || is.data.frame(data))
  if (!is.null(strata_formula)){
    stopifnot(is.logical(simple_labels), is.vector(simple_labels),
              length(simple_labels) == 1, all(!is.na(simple_labels)))
    stopifnot(inherits(strata_formula, "formula"))
    if (length(strata_formula) != 2)
      stop("The strata formula should have the form: ~var1 + var2 + ...")
    strata_vars <- stats::model.frame(stats::update(strata_formula, ~.-1),
                                      data = data, na.action = stats::na.pass)
    if (attr(attr(strata_vars, "terms"), "response") != 0)
      stop("The strata formula should not have a response variable")
    if (length(strata_vars) == 0)
      stop("The strata formula is missing stratification variables on RHS")
    complete_cases_s <- rowSums(is.na(strata_vars)) == 0
    if (!drop_NA && any(!complete_cases_s))
      stop("Missing values found in strata")
    factor_names <- c(apply(strata_vars[1, , drop = FALSE],
                            MARGIN = 1, FUN = names))
    if (simple_labels == TRUE) factor_names <- rep("", length(factor_names))
    s <- apply(strata_vars, MARGIN = 1, FUN = function(x) {
      paste0(factor_names, x, sep = "", collapse = ", ")
    })
  }else{
    complete_cases_s <- TRUE
    s <- NULL
  }

  mf <- stats::model.frame(stats::update(formula, ~.+1), data = data,
                           na.action = stats::na.pass)
  if (attr(attr(mf, "terms"), "response") == 0)
    stop("formula is missing a response")
  if (sum(sapply(mf, FUN = function(x) any(is.infinite(x)))) > 0)
    stop("Infinite values occurred when evaluating 'formula'")
  y <- mf[[attr(attr(mf, "terms"), "response")]]
  ofst <- stats::model.offset(mf)
  stopifnot(is.Survdt(y) || is.Survdt2(y))
  if (!is.null(s) && (length(s) != length(y)))
    stop("Strata formula terms and main formula terms have different lengths")
  complete_cases <- (rowSums(is.na(mf)) == 0) & complete_cases_s
  time_col <- "time"
  if (is.Survdt2(y)){
    time_col <- "stop"
    complete_cases <- !(y[, "id", drop = TRUE] %in%
                          y[!complete_cases, "id", drop = TRUE])
    id <- y[, "id", drop = TRUE]
    event_indx <- y[, "event", drop = TRUE] == 1

    # check observation windows
    event_wndws <- y[complete_cases & (y[, "event", drop = TRUE] == 1), ]
    if (check_post_event &&
        any(y[complete_cases, "start", drop = TRUE] >=
            event_wndws[match(id[complete_cases],
                              event_wndws[, "id", drop = TRUE]),
                        "stop", drop = TRUE]))
      stop(paste0("Observation windows for a given id should not extend past",
                  " the event time"))

    min_required_time <- min_time
    if (min_required_time == "sample") {
      min_required_time <- min(event_wndws[, "stop", drop = TRUE])
    }
    starts_before_min_time <- y[complete_cases, "start", drop = TRUE] <
      min_required_time
    if (any(stats::aggregate(starts_before_min_time ~ id[complete_cases],
                             FUN = sum)[, 2] == 0))
      stop(paste0("Observation windows for any given id should extend back past",
                  " the ",
                  ifelse(min_time == "sample",
                         "minimum observed event time",
                         paste0("minimum time ", min_time))))
    complete_id_nums <-
      as.numeric(unique(factor(id, levels = unique(id))[complete_cases]))
    id <- as.numeric(factor(id, levels = unique(id)))
  }else{
    id <- seq_len(length(y))
    complete_id_nums <- id[complete_cases]
    event_indx <- rep(TRUE, length(y))
  }
  match_id_indx <- match(id, complete_id_nums)

  # missing values
  n_init <- length(unique(id))
  if (!drop_NA && any(!complete_cases))
    stop("Missing values found evaluating model formula")
  n <- length(unique(id[complete_cases]))
  if (n == 0) stop("All individuals have missing values")
  X <- stats::model.matrix(mf, data = data)
  X <- X[, -which(attr(X, "assign") == 0), drop = FALSE]
  if (ncol(X) == 0) stop("No covariates provided in RHS of formula")

  if (!is.null(s)) {
    if (length(unique(s)) != length(unique(s[complete_cases])))
      warning("Dropping strata with no complete cases")
    s <- s[complete_cases]

    if (is.null(strata_levels)) {
      strata_levels <- unique(s[order(c(strata_vars[complete_cases, 1,
                                                    drop = TRUE])[1:n])])
    }

    s <- factor(s, levels = strata_levels)
    sort_indx <- order(s, y[complete_cases, time_col, drop = TRUE])
  }else{
    sort_indx <- order(y[complete_cases, time_col, drop = TRUE])
  }
  sorted_cc <- which(complete_cases)[sort_indx]
  if (is.null(ofst)) ofst <- rep(0, length(sorted_cc))

  # s may be NULL; ofst is numerical vector
  list(y = y, X = X, ofst = ofst, s = s, n = n, n_init = n_init,
       complete_id_nums = complete_id_nums, id = id,
       match_id_indx = match_id_indx, sorted_cc = sorted_cc,
       sort_indx = sort_indx, time_col = time_col,
       complete_cases = complete_cases, complete_cases_s = complete_cases_s,
       event_indx = event_indx)
}

#' @rdname positivity_sens_indtrunc
#' @export
print.coxph_pos_indtrunc <- function(x, ...){
  timewt_name <- ifelse(x[["wt_name"]] == "W-1", "none", x[["wt_name"]])
  cat(paste0("\n",
             "Sensitivity analysis for positivity violations in a Cox model",
             " under quasi-\n independent truncation.",
             "\n",
             "\n",
             "  Call: ", paste0(format(x[["call_"]]),
                                collapse = paste0("\n",
                                                  paste0(rep(" ", 8),
                                                         collapse = ""))),
             "\n",
             "\n",
             "  Model was fit on ", x[["n"]], " complete samples (",
             x[["n_miss"]], " dropped due to missingness).",
             "\n",
             "  Min event time: ", x[["eventrange"]][1],
             "\n",
             "  Max event time: ", x[["eventrange"]][2],
             "\n",
             "  ", x[["n_event"]], " observed events in the score function, ",
             "time range [", x[["timerange"]][1], ", ",
             x[["timerange"]][2], "].",
             "\n",
             "  Inverse probability of non-truncation weights estimated ",
             "nonparametrically.",
             "\n",
             "  Time-varying weight component: ", timewt_name,
             "\n",
             "\n",
             "Sensitivity analysis estimates:",
             "\n"))

  coefmat <- x[["coef"]]
  colnames(coefmat) <- paste0("trunc_mass=", x[["trunc_mass"]])
  print(coefmat)
  cat("\n")

  if (!is.null(x[["strata_summary"]])){
    sort_indx <- order(x[["strata_summary"]][, "strata_size", drop = TRUE],
                       decreasing = TRUE)
    strata_summary <- x[["strata_summary"]][sort_indx,
                                            c("strata",
                                              "strata_size",
                                              "unique_event_times"),
                                            drop = FALSE]
    colnames(strata_summary) <- c("stratum", "size", "unique event times")

    cat(paste0("Summary of the ", nrow(x[["strata_summary"]]), " strata:",
               "\n"))

    if (nrow(x[["strata_summary"]]) > 10){
      print(strata_summary[1:10, , drop = FALSE],
            row.names = FALSE)
      cat("  ...")
      cat("\n\n")
    }else{
      print(strata_summary, row.names = FALSE)
      cat("\n")
    }
  }
  invisible(x)
}



#' @rdname positivity_sens_indtrunc
#' @export
plot.coxph_pos_indtrunc <- function(x, parm, level, hazratio = FALSE,
                                    dropNA = TRUE, ...){
  if (missing(parm)) parm <- rownames(x[["coef"]])
  if (missing(level)) level <- x[["level"]]
  stopifnot(is.vector(parm), is.numeric(parm) | is.character(parm),
            is.numeric(level), is.vector(level), all(!is.na(level)),
            length(level) == 1, all(level > 0), all(level < 1),
            is.vector(hazratio), is.logical(hazratio), length(hazratio) == 1,
            all(!is.na(hazratio)), is.vector(dropNA), is.logical(dropNA),
            length(dropNA) == 1, all(!is.na(dropNA)))

  if (is.numeric(parm) && any((parm < 1) | (parm > nrow(x[["coef"]]))))
    stop("Numerical parm contains indices that are out of bounds")
  if (is.character(parm) && !all(parm %in% rownames(x[["coef"]])))
    stop("Character vector parm contains names that do not match coefficients")
  parm <- unique(parm)

  q <- stats::qnorm(1-(1-level)/2)
  n_mass <- ncol(x[["coef"]])
  if (level != x[["level"]]){
    if (n_mass == 1){
      x$lower <- x[["coef"]] - q * x[["se"]]
      x$upper <- x[["coef"]] + q * x[["se"]]
    }else{
      x$lower <- t(apply(x[["coef"]] - q * x[["se"]], MARGIN = 1, FUN = cummin))
      x$upper <- t(apply(x[["coef"]] + q * x[["se"]], MARGIN = 1, FUN = cummax))
    }
  }

  ests <- data.frame(trunc_mass = rep(x[["trunc_mass"]], each = length(parm)),
                     coef = c(x[["coef"]][parm, ]),
                     se = c(x[["se"]][parm, ]),
                     lower = c(x[["lower"]][parm, ]),
                     upper = c(x[["upper"]][parm, ]),
                     parm = rep(rownames(x[["coef"]][parm, 1, drop = FALSE]),
                                n_mass))
  ylab <- "Coefficient"
  if (hazratio == TRUE){
    ylab <- "Hazard ratio"
    ests$coef <- exp(ests$coef)
    ests$lower <- exp(ests$lower)
    ests$upper <- exp(ests$upper)
  }
  if (dropNA == TRUE) ests <- ests[!is.na(ests$coef), , drop = FALSE]
  plt <- ggplot2::ggplot(ests)
  if (any(!is.na(ests$upper))){
    if (nrow(ests) == 1){
      plt <- with(ests,
                  plt +
                    ggplot2::geom_errorbar(ggplot2::aes(x = trunc_mass,
                                                        ymin = lower,
                                                        ymax = upper)))
    }else{
      plt <- with(ests,
                  plt +
                    ggplot2::geom_line(ggplot2::aes(x = trunc_mass, y = lower),
                                       linetype = "dashed") +
                    ggplot2::geom_line(ggplot2::aes(x = trunc_mass, y = upper),
                                       linetype = "dashed"))
    }
  }
  if (nrow(ests) > 1){
    plt <- with(ests,
                plt +
                  ggplot2::geom_line(ggplot2::aes(x = trunc_mass, y = coef),
                                     linetype = "dotted"))
  }
  with(ests, plt +
         ggplot2::geom_point(ggplot2::aes(x = trunc_mass, y = coef)) +
         ggplot2::facet_wrap(~parm, scales = "free_y") +
         ggplot2::labs(x = "Truncated mass", y = ylab) +
         ggplot2::theme_classic())
}


# internal function to compute the sensitivity analysis estimator
coxph_rtrunc <- function(y, X, s = NULL, ofst = NULL, ipw, time_wt, trunc_mass,
                         ltrunc_indx = NULL, rcens_indx = NULL,
                         max_iter, error, beta_init = NULL, ipw_infl = NULL,
                         time_wt_infl = NULL){
  n <- nrow(X)
  p <- ncol(X)
  rownames(y) <- NULL
  rownames(X) <- NULL
  use_strata <- !is.null(s)
  if (is.null(ofst)) ofst <- rep(0, n)
  if (is.null(ltrunc_indx)) ltrunc_indx <- rep(FALSE, n)
  if (is.null(rcens_indx)) rcens_indx <- rep(FALSE, n)
  if (is.null(beta_init)) beta_init <- rep(0, p)

  # data sorted by event time and strata
  fit_indx <- !ltrunc_indx
  n_fit <- sum(fit_indx)
  outer_wt <- ipw * time_wt
  ofst_wt <- ofst - log(time_wt)
  y_coxph_aug <- survival::Surv(c(y[fit_indx, "time", drop = TRUE],
                                  rep(max(y[fit_indx, "time", drop = TRUE]), n)),
                                c(as.numeric(!rcens_indx[fit_indx]), rep(0, n)))
  X_aug <- rbind(X[fit_indx, , drop = FALSE], X[, , drop = FALSE])
  s_aug <- NULL
  if (use_strata) {
    s_num <- as.numeric(s)
    s_aug <- c(s_num[fit_indx], s_num)
  }
  ofst_aug <- c(ofst_wt[fit_indx], ofst_wt)

  offset <- stats::offset
  strata <- survival::strata
  beta_est <- beta_init
  coef_na <- is.na(beta_est)
  converged <- FALSE
  for (i in seq_len(max_iter)){
    exps <- exp(c(X %*% ifelse(coef_na, rep(0, p), beta_est)) + ofst)
    p_trunc <- trunc_mass^exps
    odds_trunc <- (p_trunc) / (1 - p_trunc)
    wt_aug <- c(outer_wt[fit_indx], outer_wt * odds_trunc)
    poswt <- c(rep(TRUE, n_fit), odds_trunc != 0)
    if (use_strata){
      augfit <- survival::coxph(y_coxph_aug ~ X_aug + offset(ofst_aug) +
                                  strata(s_aug), subset = poswt,
                                na.action = stats::na.fail, weights = wt_aug,
                                init = ifelse(coef_na, rep(0, p), beta_est),
                                robust = FALSE, ties = "breslow")
    }else{
      augfit <- survival::coxph(y_coxph_aug ~ X_aug + offset(ofst_aug),
                                subset = poswt, na.action = stats::na.fail,
                                weights = wt_aug,
                                init = ifelse(coef_na, rep(0, p), beta_est),
                                robust = FALSE, ties = "breslow")
    }
    beta_next <- stats::coef(augfit)
    coef_na_next <- is.na(beta_next)
    updated_coef_indx <- !coef_na_next & !coef_na
    if (all(coef_na_next))
      stop("Failed to update coefficients through augmented Cox model",
           call. = FALSE)

    if (any(updated_coef_indx)){
      # check convergence and exit without updating if achieved
      if (all(abs(beta_next[updated_coef_indx] -
                  beta_est[updated_coef_indx]) < error)){
        converged <- TRUE
        break
      }
    }

    beta_est <- beta_next
    coef_na <- coef_na_next
    next
  }

  if (!converged){
    stop("Algorithm reached max_iter without convergence",
         call. = FALSE)
  }

  coef_na <- is.na(beta_est)
  effective_p <- sum(!coef_na)
  stderrs <- rep(NA, p)
  if (!is.null(ipw_infl)){
    exps <- exp(c(X %*% ifelse(coef_na, rep(0, p), beta_est)) + ofst)
    p_trunc <- trunc_mass^exps
    odds_trunc <- (p_trunc) / (1 - p_trunc)
    wt_aug <- c(outer_wt[fit_indx], outer_wt * odds_trunc)
    posoddswt <- odds_trunc != 0
    poswt <- c(rep(TRUE, n_fit), posoddswt)
    if (use_strata){
      lastfit <- survival::coxph(y_coxph_aug ~ X_aug + offset(ofst_aug) +
                                   strata(s_aug), subset = poswt,
                                 na.action = stats::na.fail, weights = wt_aug,
                                 init = ifelse(coef_na, rep(0, p), beta_est),
                                 robust = FALSE, ties = "breslow",
                                 control = survival::coxph.control(iter.max = 0))
    }else{
      lastfit <- survival::coxph(y_coxph_aug ~ X_aug + offset(ofst_aug),
                                 subset = poswt, na.action = stats::na.fail,
                                 weights = wt_aug,
                                 init = ifelse(coef_na, rep(0, p), beta_est),
                                 robust = FALSE, ties = "breslow",
                                 control = survival::coxph.control(iter.max = 0))
    }

    # derivative
    adj_wts <- odds_trunc * (1 + odds_trunc) * exps^2 * ipw *
      ifelse(trunc_mass > 0, log(trunc_mass), 0)
    if (use_strata == FALSE) s <- rep(1, n)
    derivative <- matrix(0, nrow = effective_p, ncol = effective_p)
    for (stratum_val in unique(s)){
      strata_indx <- s == stratum_val
      adj1 <- c(t(adj_wts[strata_indx]) %*% X[strata_indx, !coef_na,
                                              drop = FALSE])
      adj2 <- t(X[strata_indx, !coef_na, drop = FALSE]) %*%
        diag(adj_wts[strata_indx], sum(strata_indx)) %*%
        X[strata_indx, !coef_na, drop = FALSE]

      unique_etimes <- unique(y[strata_indx & !ltrunc_indx & !rcens_indx, "time",
                                drop = TRUE])
      riskset_sums <- sapply(unique_etimes, FUN = function(ti){
        riskindx <- y[strata_indx, "time", drop = TRUE] >= ti
        time_match <- y[strata_indx, "time", drop = TRUE] == ti
        # sum over all strata members
        riskwts <- (riskindx + odds_trunc[strata_indx]) * exps[strata_indx] *
          ipw[strata_indx]
        outer_wt_derivative <- sum((outer_wt[strata_indx])[time_match])
        S0 <- sum(riskwts)
        S1 <- c(t(riskwts) %*% X[strata_indx, !coef_na, drop = FALSE])
        S2 <- t(X[strata_indx, !coef_na, drop = FALSE]) %*%
          diag(riskwts, length(riskwts)) %*%
          X[strata_indx, !coef_na, drop = FALSE]

        # contribution to derivative
        c(outer_wt_derivative *
            (((S2 + adj2) / S0) - (S1 %*% t(S1 + adj1) / S0^2)))
      })
      if (effective_p == 1)
        dentries <- sum(riskset_sums)
      else
        dentries <- rowSums(riskset_sums)

      derivative <- derivative + matrix(dentries,
                                        nrow = effective_p,
                                        ncol = effective_p)
    }

    # score variance components: {number of coefs without NA} x n
    score_terms <- matrix(0, nrow = effective_p, ncol = n)
    score_resids <- stats::residuals(lastfit, weighted = TRUE, type = "score",
                                     collapse = c(which(fit_indx),
                                                  which(posoddswt)))
    if (p == 1){
      score_indx <- as.numeric(names(score_resids))
      score_terms[1, score_indx] <- c(score_resids)
    }
    else{
      score_indx <- as.numeric(rownames(score_resids))
      score_terms[, score_indx] <- t(score_resids[, !coef_na, drop = FALSE])
    }
    score_terms <- score_terms +
      (score_terms %*% diag(1 / ipw, length(ipw)) %*% ipw_infl / n)

    if (!is.null(time_wt_infl)){
      sch_resids <- stats::residuals(lastfit,
                                     weighted = TRUE,
                                     type = "schoenfeld")
      sch_indx <- !ltrunc_indx & !rcens_indx
      score_terms <- score_terms +
        (t(sch_resids)[!coef_na, , drop = FALSE] %*%
           diag(1 / time_wt[sch_indx], sum(sch_indx)) %*%
           time_wt_infl[sch_indx, , drop = FALSE] / n)
    }

    meat <- apply(score_terms, MARGIN = 2, FUN = function(x) x %o% x)
    if (effective_p == 1){
      meat <- sum(meat)
    }else{
      meat <- rowSums(meat)
    }
    inv_derivative <- solve(derivative)
    var_est <- inv_derivative %*%
      matrix(meat, nrow = effective_p) %*%
      t(inv_derivative)
    stderrs[!coef_na] <- sqrt(diag(var_est))
  }
  list(coef = beta_est, se = stderrs)
}
