#' Make predictions from a \code{"l0ara"} object
#' @description Generate linear predictors, response-scale predictions, class
#'   labels, or fitted coefficients from a \code{"l0ara"} fit.
#' @param object Fitted "l0ara" object.
#' @param newx Matrix of new predictor values. If omitted, predictions are made
#'   for the training design matrix stored in \code{object$x}.
#' @param type Type of prediction required. \code{"link"} returns the linear
#'   predictor. \code{"response"} returns the response-scale mean implied by the
#'   model family. \code{"coefficients"} returns the fitted coefficient vector.
#'   \code{"class"} is available only for \code{family = "logit"} and returns
#'   0/1 class labels using a cutoff of 0 on the linear predictor.
#' @param ... Not used argument.
#' @details For Gaussian models, \code{type = "link"} and
#'   \code{type = "response"} return the same values.
#' @return The return value depends on \code{type}: a numeric vector of linear
#'   predictors or response-scale predictions, a named coefficient vector, or a
#'   numeric vector of 0/1 class labels.
#' @author
#' Wenchuan Guo <wguo007@ucr.edu>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{coef.l0ara}}, \code{\link{l0ara}}.
#' @export
predict.l0ara <- function(object, newx, type=c("link", "response", "coefficients", "class"), ...){
  type <- match.arg(type)
  beta <- coef.l0ara(object)
  if (missing(newx)) {
    newx <- object$x
  } else {
    newx <- .l0ara_as_matrix(newx)
  }
  if (type == "coefficients") {
    return(beta)
  }
  eta <- .l0ara_linear_predictor(beta, newx)
  if (type == "link" || object$family == "gaussian") {
    return(eta)
  }
  response <- drop(.l0ara_family_response(eta, object$family))
  if (type == "response") {
    return(response)
  }
  if (type=="class") {
    if (object$family=="logit") {
      return(drop(1 * (eta > 0)))
    } else {
      stop("type='class' can only be used with family='logit'", call. = FALSE)
    }
  }
}

#' Extract coefficients from a \code{"l0ara"} object
#' @description Return the fitted coefficient vector.
#' @param object Fitted "l0ara" object.
#' @param ... Not used argument.
#' @return A named numeric vector of fitted coefficients.
#' @author
#' Wenchuan Guo <wguo1017@gmail.com>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{predict.l0ara}}, \code{\link{l0ara}}.
#' @export
coef.l0ara <- function(object, ...){
  coefs <- as.vector(object$beta)
  names(coefs) <- .l0ara_coef_names(object)
  coefs
}

#' Extract coefficients from a \code{"cv.l0ara"} object
#' @description Return coefficients from the model refit at the selected
#'   \code{lam.min}.
#' @param object Fitted "cv.l0ara" object.
#' @param ... Not used argument.
#' @details If \code{object$fit.min} is missing, the model is refit on the full
#'   data using \code{lam.min}.
#' @return A named numeric vector of fitted coefficients.
#' @author
#' Wenchuan Guo <wguo1017@gmail.com>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{predict.l0ara}}, \code{\link{l0ara}},
#'   \code{\link{cv.l0ara}}.
#' @export
coef.cv.l0ara <- function(object, ...){
  fit <- object$fit.min
  if (is.null(fit)) {
    fit <- l0ara(object$x, object$y, object$family, object$lam.min)
  }
  coefs <- as.vector(fit$beta)
  names(coefs) <- .l0ara_coef_names(fit)
  coefs
}


#' Print a \code{"l0ara"} object
#' @description Print a short summary of a fitted model.
#' @param x Fitted "l0ara" object.
#' @param ... Not used argument.
#' @return Called for its side effect of printing model information.
#' @author
#' Wenchuan Guo <wguo1017@gmail.com>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{predict.l0ara}}, \code{\link{coef.l0ara}},
#'   \code{\link{l0ara}}.
#' @export
print.l0ara <- function(x, ...){
  family <- .l0ara_family_label(x$family)
  cat("Lambda used : ", .l0ara_lambda_value(x), "\n")
  cat("Model : ", family, "\n")
  cat("Iterations : ", x$iter, "\n")
  cat("Degree of freedom : ", x$df,"\n")
}

#' Print a \code{"cv.l0ara"} object
#' @description Print a short summary of a cross-validated fit.
#' @param x Fitted "cv.l0ara" object.
#' @param ... Not used argument.
#' @return Called for its side effect of printing cross-validation information.
#' @author
#' Wenchuan Guo <wguo1017@gmail.com>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{coef.cv.l0ara}}, \code{\link{plot.cv.l0ara}},
#'   \code{\link{cv.l0ara}}, \code{\link{l0ara}}.
#' @export
print.cv.l0ara <- function(x, ...){
  measure <- .l0ara_measure_label(x$measure)
  family <- .l0ara_family_label(x$family)
  cat("Number of Lambda used : ", length(x$lambda), "\n")
  cat("Optimal Lambda : ", x$lam.min, "\n")
  cat("Model : ", family, "\n")
  cat("Measure : ", measure, "\n")
  cat("Minimumn error : ",min(x$cv.error), "\n")
}

#' Plot a \code{"l0ara"} object
#' @description Plot fitted values against the linear predictor and, for
#'   logistic models, optionally add an ROC curve based on in-sample fitted
#'   probabilities.
#' @param x Fitted "l0ara" object.
#' @param auc Logical; if \code{TRUE} and \code{x$family == "logit"}, also draw
#'   an ROC curve.
#' @param split Logical; if \code{TRUE}, do not change the graphics layout with
#'   \code{par(mfrow = ...)}.
#' @param col Color used for observed points and the ROC line.
#' @param ... Not used argument.
#' @return Called for its side effect of producing one or two plots.
#' @author
#' Wenchuan Guo <wguo1017@gmail.com>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{predict.l0ara}}, \code{\link{coef.l0ara}},
#'   \code{\link{l0ara}}.
#' @export
plot.l0ara <- function(x, auc = FALSE, split = FALSE, col = 4, ...){
  nplots <- if (auc) 2 else 1
  if (!split) {
    par(mfrow = c(1, nplots))
  }
  lp <- predict(x, type = "link")
  resp <- drop(.l0ara_family_response(lp, x$family))
  plot(lp, resp, xlab = "Linear predictor", ylab = "Fitted value", pch = 20, main = "Linear predictor v.s. Fitted", ...)
  points(lp, x$y, col = col, pch = 20)
  legend("bottomright", legend = c("Fitted", "Truth"), col = c(1, col), pch = rep(20, 2))

  if (x$family == "logit" && auc) {
    n.thres <- 50
    thres <- seq(0, 1, length.out = n.thres)
    pred <- outer(resp, thres, `>`)
    fp <- colMeans(pred[x$y == 0, , drop = FALSE] == 1)
    tp <- colMeans(pred[x$y == 1, , drop = FALSE] == 1)
    area <- (tp - fp + 1) / 2
    plot(
      NA,
      xlab = "False positive rate",
      ylab = "True positive rate",
      xlim = c(0, 1),
      ylim = c(0, 1),
      main = paste("ROC curve for lambda =", .l0ara_lambda_value(x)),
      sub = paste("max auc = ", round(max(area), 2), ";", "best cut-off = ", round(thres[which.max(area)], 2)),
      ...
    )
    lines(fp, tp, pch = 19, col = col, type = "b", lty = 2)
  }
}

#' Plot a \code{"cv.l0ara"} object
#' @description Plot the cross-validation scores against the supplied penalty
#'   values and mark the selected \code{lam.min}.
#' @param x Fitted "cv.l0ara" object.
#' @param col Color used for the plotted points.
#' @param ... Not used argument.
#' @return Called for its side effect of producing a plot.
#' @author
#' Wenchuan Guo <wguo1017@gmail.com>, Shujie Ma <shujie.ma@ucr.edu>, Zhenqiu Liu <Zhenqiu.Liu@cshs.org>
#' @seealso \code{\link{coef.cv.l0ara}}, \code{\link{cv.l0ara}},
#'   \code{\link{l0ara}}.
#' @export

plot.cv.l0ara <- function(x, col = 3, ...) {
  cv.u <- x$cv.error + x$cv.std
  cv.l <- x$cv.error - x$cv.std
  plot(x = x$lambda, y = x$cv.error, main = .l0ara_family_label(x$family), ylim = range(cv.l, cv.u), xlab = "Lambda", ylab = x$name, col = col, pch = 19, sub = paste("Optimal lambda = ", round(x$lam.min, 2)))
  abline(v = x$lam.min, lty = 3)
}

