# printing
#' @param x An object of class `DTRreg`.
#' @rdname DTRreg
#' @importFrom utils tail
#' @export
print.DTRreg <- function(x,...) {

  K <- x$K

  cat(ifelse(x$setup$type == "DTR", 
             paste("DTR estimation over ", K, " stages:\n\n"), 
             ""),
      "Blip parameter estimates",
      ifelse(x$setup$var.estim == "none", 
             " (standard errors not requested)\n", 
             "\n"))
  
  for (k in 1L:x$K) {
    if (x$setup$var.estim == "none" || {is.list(x$covmat) && any(x$covmat[[k]] == "aborted")}) {
      cat("Stage ", k, " (n = ", x$analysis$n[[k]], ")\n", sep = "")
      result <- data.frame("Estimate" = sprintf(x$psi[[k]], fmt = '%#.4f'))
      rownames(result) <- names(x$psi[[k]])
      print(result)
      cat("\n")
    } else {
      cat("Stage ", k, " (n = ", x$analysis$n[[k]], ")\n", sep = "")
      
      se <- utils::tail(sqrt(diag(x$covmat[[k]])), length(x$psi[[k]]))
      ci <- cbind(x$psi[[k]] - 1.96 * se, x$psi[[k]] + 1.96 * se)
      result <- data.frame("Estimate" = sprintf(x$psi[[k]], fmt = '%#.4f'),
                           "Std_Error" = sprintf(se, fmt = '%#.4f'),
                           "95%_CI_Lower" = sprintf(pmin(ci[, 1L], ci[, 2L]), fmt = '%#.4f'), 
                           "95%_CI_Upper" = sprintf(pmax(ci[, 1L], ci[, 2L]), fmt = '%#.4f'))
      rownames(result) <- names(x$psi[[k]])
      print(result)
    }
  }

  # print warnings if non-regularity concerns
  if (x$setup$type == "DTR") {
    if (any(!is.na(x$nonreg))) {
      
      tst <- unlist(x$analysis$cts) == "bin" & x$nonreg > 0.05
      if (any(tst)) {
        use_plural <- sum(tst) > 1L
        cat("\nWarning: possible non-regularity at ",
            ifelse(use_plural, "stages ", "stage "),
            paste(which(tst), collapse = ", "),
            ifelse(use_plural, " (probs = ", "(prob = "),
            paste(format(x$nonreg[tst], digits = 3), collapse = ", "), 
            ifelse(use_plural, ", respectively)\n", ")\n"), sep = "")
      }
    }
    
    # print DTRregs explicitly if bin or cts.l
    if (x$analysis$cts[[1L]] == "bin") cat("Recommended dynamic treatment regimen:\n")
    
    for (j in 1L:K) {
      
      # format parameter estimate to be 3 digits w/ scientific notation
      est <- sprintf(abs(x$psi[[j]]), fmt = '%#.4f')
      
      # identify the appropriate +/- connector
      sg <- c("", rep("+", length(est) - 1L))
      sg[x$psi[[j]] < 0.0] <- "-"
      
      # names of parameters
      nms <- names(x$psi[[j]])
      nms[1L] <- ""
      
      if (is.null(x$optimization)) x$optimization <- "max"
      
      # combine these to form the rule  
      rule <- paste0(paste(sg, est, nms, collapse = " "), " > 0")
      
      if (x$analysis$cts[[j]] == "bin") {
        cat("Stage ", j, ": treat if", rule, "\n", sep = "")
      } else if (x$analysis$cts[[j]] == "cts.l") {
        cat("Stage ", j, ": maximum treatment if",
            rule, ", otherwise minimum treatment.\n", sep = "")
      }
    }
  }
}


#' @param object An object of class `DTRreg`.
#' @param ... Ignored.
#' @rdname DTRreg
#' @export
summary.DTRreg <- function(object,...) {
  print.DTRreg(object)
}

#' Optimal Outcome Prediction for DTRs
#' 
#' Predicted outcome assuming optimal treatment (according to analysis via 
#'   G-estimation or dWOLS) was followed.  Assumes blip and treatment-free 
#'   models correctly specified.
#'  
#' This function may be used in a similar fashion to more traditional modeling 
#'   commands (such as lm).  Users are referred to the primary `DTRreg()` 
#'   and `DTRSurv()` help command
#'   (and associated literature) for information concerning model specification.
#'   In particular, we note that the predict function assumes that the 
#'   treatment-free model has been correctly specified, as the treatment-free 
#'   parameters are used in the prediction process.
#'   
#' @param object A model object generated by the function `DTRreg()` or `DWSurv()`.
#' @param newdata A dataset (usually the data analyzed by DTRreg for which 
#'   predicted outcomes are desired.  If a new dataset is provided, variable 
#'   names should correspond to those presented to `DTRreg()` or `DWSurv()`.
#' @param treat.range If treatment is continuous (rather than binary), a 
#'   vectors of the form c(min,max) which specify the minimum and maximum value 
#'   the treatment may take at stage 1.  
#'   If unspecified, this will be inferred from the 
#'   treat.range provided with use of the original DTRreg command.  As such, if 
#'   no treatment range was specified there either, treat.range will be the 
#'   minimum and maximum observed first stage treatment.
#' @param ... Space for additional arguments (not currently used)
#' 
#' @return An \eqn{n}{n} x 1 matrix of predicted outcome values.
#' 
#' @references
#' Chakraborty, B., Moodie, E. E. M. (2013) \emph{Statistical Methods for 
#'   Dynamic Treatment Regimes}. New York: Springer.
#'   
#' Robins, J. M. (2004) \emph{Optimal structural nested models for optimal 
#'   sequential decisions}. In Proceedings of the Second Seattle Symposium on 
#'   Biostatistics, D. Y. Lin and P. J. Heagerty (eds), 189-326. 
#'   New York: Springer.
#'   
#' Wallace, M. P., Moodie, E. M. (2015) Doubly-Robust Dynamic Treatment Regimen 
#'   Estimation Via Weighted Least Squares. \emph{Biometrics} \bold{71}(3), 
#'   636-644 (doi:10.1111/biom.12306.)
#'   
#' @author Michael Wallace
#' 
#' @examples
#' # example single run of a 2-stage g-estimation analysis
#' 
#' set.seed(1)
#' 
#' # expit function
#' expit <- function(x) { 1.0 / (1.0 + exp(-x)) }
#' 
#' # sample size
#' n <- 10000
#' 
#' # variables (X = patient information, A = treatment)
#' X1 <- rnorm(n)
#' A1 <- rbinom(n, 1, expit(X1))
#' X2 <- rnorm(n)
#' A2 <- rbinom(n, 1, expit(X2))
#' 
#' # blip functions
#' gamma1 <- A1 * (1 + X1)
#' gamma2 <- A2 * (1 + X2)
#' 
#' # observed outcome: treatment-free outcome plus blip functions
#' Y <- exp(X1) + exp(X2) + gamma1 + gamma2 + rnorm(n)
#'  
#' # models to be passed to DTRreg
#' # blip model
#' blip.mod <- list(~ X1, ~ X2)
#' # treatment model (correctly specified)
#' treat.mod <- list(A1 ~ X1, A2 ~ 1)
#' # treatment-free model (incorrectly specified)
#' tf.mod <- list(~ X1, ~ X2)
#' 
#' # perform G-estimation
#' mod1 <- DTRreg(twoStageCont$Y, blip.mod, treat.mod, tf.mod, 
#'               data = twoStageCont, method = "gest")
#'               
#' # predicted Y for optimal treatment
#' dat <- data.frame(X1, X2, A1, A2)
#' predict(mod1, newdata = dat)
#' @concept dynamic treatment regimens
#' @concept adaptive treatment strategies
#' @concept personalized medicine
#' @concept g-estimation
#' @concept dynamic weighted ordinary least squares
#' @export
predict.DTRreg <- function(object, newdata, treat.range = NULL, ...) {
  K <- object$K
  
  stopifnot("`treat.range` must be a list of length `stage`" = is.null(treat.range) ||
              {is.list(treat.range) && all(lengths(treat.range) == 2L)})
  
  tx.range <- if (!is.list(treat.range)) object$setup$tx.range[[1L]] else treat.range[[1L]]
  cts_obj <- switch(object$analysis$cts[[1L]],
                    "bin"      = Binary$new(tf.model   = object$setup$models[[1L]]$tf,
                                            blip.model = object$setup$models[[1L]]$blip,
                                            tx.var     = object$setup$tx.vars[1L]),
                    "multinom" = MultiNom$new(tf.model   = object$setup$models[[1L]]$tf,
                                              blip.model = object$setup$models[[1L]]$blip,
                                              tx.var     = object$setup$tx.vars[1L],
                                              tx.levels  = levels(object$training_data$A[[1L]])),
                    "cts.q"    = ContQuadraticBlip$new(tf.model    = object$setup$models[[1L]]$tf,
                                                       blip.model  = object$setup$models[[1L]]$blip,
                                                       tx.var      = object$setup$tx.vars[1L],
                                                       treat.range = tx.range)
  )
  
  if (!missing(newdata)) {
    if (!all(complete.cases(newdata))) 
      stop("if provided, `newdata` must be complete", call. = FALSE)
    newdata[[object$setup$tx.vars[1L]]] <-
      cts_obj$opt(outcome.fit = object$analysis$outcome.fit[[1L]], data = newdata)
  } else {
    newdata <- object$training_data$data
    newdata[[object$setup$tx.vars[1L]]] <- object$analysis$opt.treat[[1L]]
  }
  
  newdata <- cts_obj$prep(newdata, newdata[[object$setup$tx.vars[1L]]])
  
  pred <- predict(object$analysis$outcome.fit[[1L]], newdata = newdata)
  if ("DWSurv" %in% class(object)) pred <- exp(pred)
  pred
}

#' @rdname DTRreg
#' @export
coef.DTRreg <- function(object,...) {
  K <- object$K
  psi <- object$psi
  names(psi) <- paste0("stage_", 1L:K)
  psi
}

#' @noRd
#' @keywords internal
.ciStep.se <- function(psij, level, covmatj) {
  psi <- matrix(NA, nrow = length(psij), ncol = 2L)
  rownames(psi) <- names(psij)
  colnames(psi) <- c(paste(100.0 * {{1.0 - level} / 2.0},"%"), 
                          paste(100.0 * {{1.0 + level} / 2.0},"%"))
  
  psi[, 1L] <- psij + 
    stats::qnorm({1.0 - level} / 2.0) * sqrt(diag(covmatj))
    
  psi[, 2L] <- psij + 
    stats::qnorm({1.0 + level} / 2.0) * sqrt(diag(covmatj))

  psi    
}

#' @noRd
#' @keywords internal
.ciStep.percentile <- function(psij, level, var.estim, psi.bootj) {
  if (var.estim != "bootstrap") {
    stop("Percentile confidence intervals only available with bootstrap.",
         call. = FALSE)
  }
  
  psi <- matrix(NA, nrow = length(psij), ncol = 2L)
  rownames(psi) <- names(psij)
  colnames(psi) <- c(paste(100.0 * {{1.0 - level} / 2.0},"%"), 
                     paste(100.0 * {{1.0 + level} / 2.0},"%"))
  
  psi[, 1L] <- apply(psi.bootj, 2L, stats::quantile, probs = {1.0 - level} / 2.0)
  psi[, 2L] <- apply(psi.bootj, 2L, stats::quantile, probs = {1.0 + level} / 2.0)
  psi
}

#' Confidence Interval Calculations for DTRs
#' 
#' Confidence intervals for parameters, with the option of 
#'   constructing the confidence intervals using the percentile method 
#'   when bootstrap is used.
#'   
#' @param object A model object generated by the function DTRreg.
#' @param parm Not available for DTRreg objects.
#' @param level The confidence level required.
#' @param type Typical Wald-type confidence interval "se" (default) or 
#'   confidence intervals derived with the percentile method "percentile" 
#'   (bootstrap variance estimates only).
#' @param ... Space for additional arguments (not currently used).
#' 
#' @return A list with columns giving lower and upper confidence limits for each
#'   parameter. These will be labelled as (1-level)/2 and 1 - (1-level)/2 in 
#'   percentage (by default 2.5\% and 97.5\%).
#'   
#' @importFrom stats confint qnorm quantile
#' @export
confint.DTRreg <- function(object, parm = NULL, level = 0.95, 
                           type = c("se", "percentile"), ...) {
  
  type <- match.arg(type)
  
  if (object$setup$var.estim == "none") {
    message("confint() only available when DTRreg called with a variance",
            "estimation option")
    return(invisible(NULL))
  }
  
  if (any(object$covmat[[1]] == "aborted")) {
    message("variance estimation aborted; unable to obtain confidence intervals")
    return(invisible(NULL))
  }
    
  K <- object$K
  psi <- object$psi
  names(psi) <- paste0("stage_", 1L:K)
  
  for (j in 1L:K) {
    psi[[j]] <- if (type == "se") {
      .ciStep.se(object$psi[[j]], level, object$covmat[[j]])
    } else {
      .ciStep.percentile(object$psi[[j]], level, object$setup$var.estim, object$psi.boot[[j]])
    }
  }
  psi
}

#' Diagnostic Plots for DTR Estimation
#' 
#' Diagnostic plots for assessment of treatment, treatment-free, and blip models 
#'   following DTR estimation using DTRreg or DWSurv.
#'   
#' DTR estimation using G-estimation and dWOLS requires the specification of 
#'   three models: the treatment, treatment-free, and blip.  The treatment model 
#'   may be assessed via standard diagnostics, whereas the treatment-free and 
#'   blip models may be simultaneously assessed using diagnostic plots 
#'   introduced by Rich et al.  The plot() function first presents diagnostic 
#'   plots that assess the latter, plotting fitted values against residuals and 
#'   covariates following DTR estimation.  If there is any evidence of a 
#'   relationship between the variables in these plots, this is evidence that at 
#'   least one of the blip or treatment-free models is mis-specified.
#'   
#' Following these plots, the plot() function will present standard diagnostic 
#'   plots for the treatment model.  These are produced directly by the standard 
#'   plot() command applied to the models that were fit.  For example, if 
#'   treatment is binary, the resulting plots are the same as those that are 
#'   generated by the plot() command applied to a glm object for logistic 
#'   regression.
#' 
#' @param x A model object generated by the functions DTRreg and DWSurv.
#' @param ... Space for additional arguments (not currently used)
#' @returns Invisibly, a list of gg objects.
#' 
#' @examples
#' # example single run of a 2-stage g-estimation analysis
#' 
#' set.seed(1)
#' 
#' # expit function
#' expit <- function(x) { 1.0 / (1.0 + exp(-x)) }
#' 
#' # sample size
#' n <- 10000
#' 
#' # variables (X = patient information, A = treatment)
#' X1 <- rnorm(n)
#' A1 <- rbinom(n, 1, expit(X1))
#' X2 <- rnorm(n)
#' A2 <- rbinom(n, 1, expit(X2))
#' 
#' # blip functions
#' gamma1 <- A1 * (1 + X1)
#' gamma2 <- A2 * (1 + X2)
#' 
#' # observed outcome: treatment-free outcome plus blip functions
#' Y <- exp(X1) + exp(X2) + gamma1 + gamma2 + rnorm(n)
#'  
#' # models to be passed to DTRreg
#' # blip model
#' blip.mod <- list(~ X1, ~ X2)
#' # treatment model (correctly specified)
#' treat.mod <- list(A1 ~ X1, A2 ~ 1)
#' # treatment-free model (incorrectly specified)
#' tf.mod <- list(~ X1, ~ X2)
#' 
#' # perform G-estimation
#' mod1 <- DTRreg(twoStageCont$Y, blip.mod, treat.mod, tf.mod, 
#'               data = twoStageCont, method = "gest")
#'               
#' # model diagnostics: note treatment-free model is mis-specified
#' plot(mod1)
#' 
#' @export
plot.DTRreg <- function(x, ...) {
  # call the autoplot method to get the list of ggplots
  plots <- ggplot2::autoplot(x, ...)
  
  # handle the interactive display
  if (interactive() && length(plots) > 1) {
    old_ask <- grDevices::devAskNewPage(TRUE)
    on.exit(grDevices::devAskNewPage(old_ask), add = TRUE)
  }
  
  # print each plot to the active graphics device
  for (p in plots) {
    # if it's a ggplot, print it; if it's a base plot caught in the list, plot it.
    if (inherits(p, "ggplot") || inherits(p, "grob")) {
      print(p)
    } else {
      plot(p)
    }
  }
  
  # Return the list invisibly
  invisible(plots)
}


#' Diagnostic Plots for DTR Estimation via autoplot()
#' 
#' Diagnostic plots for assessment of treatment, treatment-free, and blip models 
#'   following DTR estimation using DTRreg or DWSurv.
#'   
#' @param object A model object generated by the functions DTRreg and DWSurv.
#' @param ... Space for additional arguments (not currently used)
#' @returns A named list of gg objects.
#' 
#' @examples
#' # example single run of a 2-stage g-estimation analysis
#' 
#' set.seed(1)
#' 
#' # expit function
#' expit <- function(x) { 1.0 / (1.0 + exp(-x)) }
#' 
#' # sample size
#' n <- 10000
#' 
#' # variables (X = patient information, A = treatment)
#' X1 <- rnorm(n)
#' A1 <- rbinom(n, 1, expit(X1))
#' X2 <- rnorm(n)
#' A2 <- rbinom(n, 1, expit(X2))
#' 
#' # blip functions
#' gamma1 <- A1 * (1 + X1)
#' gamma2 <- A2 * (1 + X2)
#' 
#' # observed outcome: treatment-free outcome plus blip functions
#' Y <- exp(X1) + exp(X2) + gamma1 + gamma2 + rnorm(n)
#'  
#' # models to be passed to DTRreg
#' # blip model
#' blip.mod <- list(~ X1, ~ X2)
#' # treatment model (correctly specified)
#' treat.mod <- list(A1 ~ X1, A2 ~ 1)
#' # treatment-free model (incorrectly specified)
#' tf.mod <- list(~ X1, ~ X2)
#' 
#' # perform G-estimation
#' mod1 <- DTRreg(twoStageCont$Y, blip.mod, treat.mod, tf.mod, 
#'               data = twoStageCont, method = "gest")
#'               
#' # model diagnostics: note treatment-free model is mis-specified
#' gg_list <- ggplot2::autoplot(mod1)
#' 
#' @import ggplot2
#' @export
autoplot.DTRreg <- function(object, ...) {
  plot_list <- list()
  
  # residual Analysis
  for (j in seq_len(object$K)) {
    cases <- object$analysis$last.stage >= j
    obsK   <- object$analysis$fitted.values[[j]]
    fitK <- object$analysis$outcome.fit[[j]]$fitted.values
    resids <- object$analysis$residuals[[j]]
    
    # helper to generate the standardized diagnostic plot
    make_resid_plot <- function(x_val, y_val, x_label, title) {
      df <- data.frame(x = x_val, y = y_val)
      ggplot(df, aes(.data[["x"]], .data[["y"]])) +
        geom_point(alpha = 0.5, shape = 1) +
        geom_hline(yintercept = 0, linetype = "dashed", color = "blue") +
        geom_smooth(method = "loess", color = "red", se = FALSE, linewidth = 0.8) +
        labs(title = title, x = x_label, y = "Residuals") +
        theme_minimal()
    }
    
    # plot: Residuals vs Fitted
    name_fitted <- paste0("Stage", j, "_Fitted")
    plot_list[[name_fitted]] <- make_resid_plot(
      fitK, resids, "Fitted values", paste("Stage", j, "Residuals vs Fitted")
    )
    
    # plots: Blip Covariates
    blip <- object$setup$models[[j]]$blip
    blip_mm <- model.matrix(blip, object$analysis$blip.data[[j]])
    if (attr(terms(blip), "intercept") == 1L) blip_mm <- blip_mm[, -1L, drop = FALSE]
    
    if (ncol(blip_mm) > 0L) {
      for (i in seq_len(ncol(blip_mm))) {
        col_name <- colnames(blip_mm)[i]
        name_cov <- paste0("Stage", j, "_", col_name)
        plot_list[[name_cov]] <- make_resid_plot(
          blip_mm[, i], resids, col_name, paste("Stage", j, "Residuals vs", col_name)
        )
      }
    }
  }
  
  # Sub-Model Diagnostics (Treatment, Censoring, etc.)
  # use ggplotify or ggfortify if possible to convert them.
  add_submodel_plots <- function(plot_list, mod, prefix) {
    if (is.logical(mod)) return(plot_list)
    
    tryCatch({
      if (inherits(mod, c("glm", "lm"))) {
        # handle directly - don't rely on ggfortify being loaded
        df  <- data.frame(fitted   = fitted(mod),
                          residuals = stats::residuals(mod, type = "deviance"))
        p   <- ggplot2::ggplot(df, ggplot2::aes(.data[["fitted"]], .data[["residuals"]])) +
          ggplot2::geom_point(alpha = 0.5, shape = 1) +
          ggplot2::geom_hline(yintercept = 0, linetype = "dashed",
                              color = "blue") +
          ggplot2::geom_smooth(method = "loess", color = "red",
                               se = FALSE, linewidth = 0.8) +
          ggplot2::labs(title    = paste(prefix, "Residuals vs Fitted"),
                        x = "Fitted values", y = "Deviance Residuals") +
          ggplot2::theme_minimal()
        plot_list[[prefix]] <- p
      } else {
        p <- ggplot2::autoplot(mod)
        plot_list[[prefix]] <- p
      }
    }, error = function(e1) {
      if (requireNamespace("ggplotify", quietly = TRUE)) {
        tryCatch({
          plot_list[[prefix]] <<- ggplotify::as.ggplot(function() {
            old_ask <- graphics::par(ask = FALSE)
            on.exit(graphics::par(old_ask))
            plot(mod)
          })
        }, error = function(e2) {
          # silently skip
        })
      }
    })
    
    plot_list
  }  
  for (j in seq_len(object$K)) {
    plot_list <- add_submodel_plots(plot_list,
                                    mod = object$analysis$tx.mod.fitted[[j]], 
                                    paste0("Stage", j, "_Treatment"))
    
    if ("DWSurv" %in% class(object)) {
      plot_list <- add_submodel_plots(plot_list,
                                      object$analysis$cens.mod.fitted[[j]], 
                                      paste0("Stage", j, "_Censoring"))
    } else {
      plot_list <- add_submodel_plots(plot_list,
                                      object$analysis$cc.mod.fitted[[j]], 
                                      paste0("Stage", j, "_CompleteCases"))
    }
  }
  
  plot_list
}

