#' Normal Skew Model fit Function
#'
#' @description
#' Estimates the parameters of the Sample Selection Model with Skew-Normal
#' Distribution
#'
#' @return
#'
#' Returns a list with the following components.
#'
#' Coefficients: Returns a numerical vector with the best estimated values
#' of the model parameters;
#'
#' Value: The value of function to be minimized (or maximized) corresponding
#' to par.
#'
#' loglik: Maximized value of the log-likelihood function calculated from the estimated coefficients.
#'
#' gradient: Analytical score vector evaluated at the estimated coefficients.
#'
#' counts: Component of the Optim function. A two-element integer vector
#' giving the number of calls to fn and gr respectively. This excludes
#' those calls needed to compute the Hessian, if requested, and any calls
#' to fn to compute a finite-difference approximation to the gradient.
#'
#' hessian: Component of the Optim function, with pre-defined option
#' hessian=TRUE. A symmetric matrix giving an estimate of the Hessian
#' at the solution found. Note that this is the Hessian of the unconstrained
#' problem even if the box constraints are active.
#'
#' fisher_infoSK: Fisher information matrix
#'
#' prop_sigmaSK: Square root of the Fisher information matrix diagonal
#'
#' level: Selection variable levels
#'
#' nObs: Numeric value representing the size of the database
#'
#' nParam: Numerical value representing the number of model parameters
#'
#' N0: Numerical value representing the number of unobserved entries
#'
#' N1: Numerical value representing the number of complete entries
#'
#' NXS: Numerical value representing the number of parameters of the
#' selection model
#'
#' NXO: Numerical value representing the number of parameters of the
#' regression model
#'
#' df: Numerical value that represents the difference between the size
#' of the response vector of the selection equation and the number of
#' model parameters
#'
#' aic: Numerical value representing Akaike's information criterion.
#'
#' bic: Numerical value representing Schwarz's Bayesian Criterion
#'
#' initial.value: Numerical vector that represents the input values
#' (Initial Values) used in the parameter estimation.
#'
#' @details
#' The HeckmanSK() function fits the Sample Selection Model
#' based on the Skew-normal distribution. For more information see Ogundimu
#' and Hutton (2016) \doi{10.1111/sjos.12171}.
#'
#'
#' @param selection Selection equation.
#' @param outcome Primary Regression Equation.
#' @param lambda Numeric scalar or vector with initial values for the skewness parameter.
#' If more than one value is supplied and \code{start = NULL}, the function uses
#' a multi-start strategy and keeps the fit with the largest log-likelihood.
#' @param start initial values.
#' @param data Database.
#' @examples
#' data("Mroz87")
#' attach(Mroz87)
#' selectEq <- lfp ~ huswage + kids5 + mtr + fatheduc + educ + city
#' outcomeEq <- log(wage) ~ educ+city
#' HeckmanSK(selectEq, outcomeEq, data = Mroz87, lambda = -1.5)
#' @importFrom Rdpack reprompt
#' @references Ogundimu, E. O. and Hutton, J. L. (2016). A sample
#' selection model with skew-normal distribution. \emph{Scandinavian
#' Journal of Statistics}. \doi{10.1111/sjos.12171}.
#' @export HeckmanSK
#' @export
HeckmanSK <- function(selection, outcome, data = sys.frame(sys.parent()),
                      lambda = c(-5, -3, -2, -1, -0.5, 0, 0.5, 1, 2, 3, 5),
                      start = NULL) {

  # Extract design matrices and outcomes from the selection and outcome models
  components <- extract_model_components(selection = selection, outcome = outcome, data = data)
  XS  <- components$XS       # Design matrix for selection equation
  YS  <- components$YS       # Response variable for selection equation
  NXS <- components$NXS      # Number of selection parameters
  XO  <- components$XO       # Design matrix for outcome equation
  YO  <- components$YO       # Response variable for outcome equation
  NXO <- components$NXO      # Number of outcome parameters
  YSLevels <- components$YSLevels  # Levels of the selection variable

  # Replace missing values in the outcome with zeros
  YO[is.na(YO)] <- 0

  #############################
  # Log-likelihood function
  #############################
  loglik_SK <- function(start) {
    # Define parameter indices
    NXS <- dim(model.matrix(~XS))[2] - 1
    NXO <- dim(model.matrix(~XO))[2] - 1
    istartS <- 1:NXS
    istartO <- seq(tail(istartS, 1) + 1, length = NXO)
    isigma  <- tail(istartO, 1) + 1
    irho    <- tail(isigma, 1) + 1
    ilamb1  <- tail(irho, 1) + 1

    # Extract parameters
    g      <- start[istartS]
    b      <- start[istartO]
    sigma  <- start[isigma]
    if (!is.finite(sigma) || sigma <= 0) return(NA_real_)
    rho    <- start[irho]
    if (!is.finite(rho) || abs(rho) >= 1) return(NA_real_)
    lamb1  <- start[ilamb1]

    # Linear predictors
    XS.g <- XS %*% g
    XO.b <- XO %*% b
    u2   <- YO - XO.b
    z    <- u2 / sigma
    r    <- sqrt(1 - rho^2)
    u    <- 1 + lamb1^2 - (lamb1 * rho)^2
    lstar <- -lamb1 * rho / sqrt(u)

    # Compute selection log-likelihood
    h <- pmax(sn::psn(-XS.g, 0, 1, -lstar), .Machine$double.xmin)

    # Log-likelihood vector
    ll <- ifelse(YS == 0,
                 log(h),
                 log(2) - log(sigma) + stats::dnorm(z, log = TRUE) +
                   stats::pnorm(lamb1 * z, log.p = TRUE) +
                   stats::pnorm((XS.g + rho * z) / r, log.p = TRUE))

    return(sum(ll))
  }

  #############################
  # Gradient of the log-likelihood
  #############################
  gradlik_SK <- function(start) {
    # Dimensions and parameter positions
    NXS <- ncol(XS)
    NXO <- ncol(XO)
    nObs <- length(YS)
    nParam <- NXS + NXO + 3

    istartS <- seq_len(NXS)
    istartO <- seq(from = NXS + 1L, length.out = NXO)
    isigma  <- NXS + NXO + 1L
    irho    <- NXS + NXO + 2L
    ilamb1  <- NXS + NXO + 3L

    # Extract parameters
    g     <- start[istartS]
    b     <- start[istartO]
    sigma <- start[isigma]
    rho   <- start[irho]
    lamb1 <- start[ilamb1]

    if (!is.finite(sigma) || sigma <= 0) {
      return(rep(NA_real_, nParam))
    }
    if (!is.finite(rho) || abs(rho) >= 1) {
      return(rep(NA_real_, nParam))
    }
    if (!is.finite(lamb1)) {
      return(rep(NA_real_, nParam))
    }

    # Subsets
    sel0 <- YS == 0
    sel1 <- YS == 1

    XS0 <- XS[sel0, , drop = FALSE]
    XS1 <- XS[sel1, , drop = FALSE]
    XO1 <- XO[sel1, , drop = FALSE]
    YO1 <- YO[sel1]

    # Linear predictors and auxiliary quantities
    XS0.g <- as.numeric(XS0 %*% g)
    XS1.g <- as.numeric(XS1 %*% g)
    XO1.b <- as.numeric(XO1 %*% b)

    z <- (YO1 - XO1.b) / sigma
    r <- 1 - rho^2
    u <- 1 + lamb1^2 - (lamb1 * rho)^2

    # The censored contribution is log F_SN(-eta; q), with
    # q = lambda * rho / sqrt(u). In the notation used in the paper,
    # lambda* = -q is used for P(S = 1), whereas P(S = 0) uses q.
    q <- lamb1 * rho / sqrt(u)

    # Selected observations: K1 = phi(omega)/Phi(omega),
    # K3 = phi(lambda z)/Phi(lambda z).
    omega <- (XS1.g + rho * z) / sqrt(r)
    eta   <- lamb1 * z

    K1 <- exp(stats::dnorm(omega, log = TRUE) -
                stats::pnorm(omega, log.p = TRUE))
    K3 <- exp(stats::dnorm(eta, log = TRUE) -
                stats::pnorm(eta, log.p = TRUE))

    # Censored observations: F0 = F_SN(-eta_s; q),
    # K2 = phi(eta_s) Phi(-q eta_s) / F_SN(-eta_s; q),
    # K4 = phi(eta_s sqrt(1 + lambda^2) / sqrt(u)) / F_SN(-eta_s; q).
    F0 <- sn::psn(-XS0.g, xi = 0, omega = 1, alpha = q)
    F0 <- pmax(F0, .Machine$double.xmin)

    K2 <- exp(stats::dnorm(XS0.g, log = TRUE) +
                stats::pnorm(-q * XS0.g, log.p = TRUE) -
                log(F0))

    K4 <- exp(stats::dnorm(XS0.g * sqrt(1 + lamb1^2) / sqrt(u), log = TRUE) -
                log(F0))

    gradient <- matrix(0, nrow = nObs, ncol = nParam)

    # d ell / d gamma
    # Correct selected term: K1 / sqrt(r), not K1 / r.
    if (nrow(XS1) > 0L) {
      gradient[sel1, istartS] <- XS1 * as.numeric(K1 / sqrt(r))
    }
    if (nrow(XS0) > 0L) {
      gradient[sel0, istartS] <- XS0 * as.numeric(-2 * K2)
    }

    # d ell / d beta
    if (nrow(XO1) > 0L) {
      gradient[sel1, istartO] <- XO1 * as.numeric(
        z / sigma - (lamb1 * K3) / sigma - (rho * K1) / (sigma * sqrt(r))
      )
    }

    # d ell / d sigma
    if (sum(sel1) > 0L) {
      gradient[sel1, isigma] <-
        -1 / sigma +
        z^2 / sigma -
        lamb1 * z * K3 / sigma -
        rho * z * K1 / (sigma * sqrt(r))
    }

    # d ell / d rho
    if (sum(sel1) > 0L) {
      gradient[sel1, irho] <- K1 * (z + rho * XS1.g) / (r^(3 / 2))
    }
    if (sum(sel0) > 0L) {
      gradient[sel0, irho] <- -2 * lamb1 * K4 / sqrt(2 * pi * u)
    }

    # d ell / d lambda
    # Correct censored term: (1 + lambda^2) is in the denominator.
    if (sum(sel1) > 0L) {
      gradient[sel1, ilamb1] <- K3 * z
    }
    if (sum(sel0) > 0L) {
      gradient[sel0, ilamb1] <-
        -2 * rho * K4 / ((1 + lamb1^2) * sqrt(2 * pi * u))
    }

    colSums(gradient)
  }

  #############################
  # Initial values and multi-start strategy
  #############################

  # Internal optimizer used for a single starting vector.
  fit_one_SK <- function(start_i) {
    tryCatch(
      stats::optim(
        par     = start_i,
        fn      = loglik_SK,
        gr      = gradlik_SK,
        method  = "BFGS",
        hessian = TRUE,
        control = list(fnscale = -1)
      ),
      error = function(e) {
        structure(
          list(error = conditionMessage(e), par = start_i, value = NA_real_),
          class = "try-error"
        )
      }
    )
  }

  if (is.null(start)) {
    message("Start not provided; using multi-start initial values for lambda.")

    lambda_grid <- unique(as.numeric(lambda))
    lambda_grid <- lambda_grid[is.finite(lambda_grid)]

    if (length(lambda_grid) == 0L) {
      stop("Argument 'lambda' must contain at least one finite numeric value.")
    }

    start_base <- step2(YS, XS, YO, XO)
    start_list <- lapply(lambda_grid, function(lam) c(start_base, lam))

    fit_list <- lapply(start_list, fit_one_SK)

    fit_values <- vapply(fit_list, function(fit) {
      if (inherits(fit, "try-error") || is.null(fit$value) || !is.finite(fit$value)) {
        return(NA_real_)
      }
      fit$value
    }, numeric(1))

    if (all(is.na(fit_values))) {
      stop("All HeckmanSK multi-start optimizations failed. Try supplying 'start' manually.")
    }

    best_id <- which.max(fit_values)
    theta_SK <- fit_list[[best_id]]
    start <- start_list[[best_id]]

    lambda_search <- data.frame(
      lambda_start = lambda_grid,
      logLik       = fit_values,
      convergence  = vapply(fit_list, function(fit) {
        if (inherits(fit, "try-error") || is.null(fit$convergence)) return(NA_integer_)
        fit$convergence
      }, integer(1)),
      message      = vapply(fit_list, function(fit) {
        if (inherits(fit, "try-error")) return(fit$error)
        if (is.null(fit$message)) return(NA_character_)
        as.character(fit$message)
      }, character(1)),
      sigma_hat    = vapply(fit_list, function(fit) {
        if (inherits(fit, "try-error") || is.null(fit$par) || length(fit$par) < NXS + NXO + 3) return(NA_real_)
        fit$par[NXS + NXO + 1]
      }, numeric(1)),
      rho_hat      = vapply(fit_list, function(fit) {
        if (inherits(fit, "try-error") || is.null(fit$par) || length(fit$par) < NXS + NXO + 3) return(NA_real_)
        fit$par[NXS + NXO + 2]
      }, numeric(1)),
      lambda_hat   = vapply(fit_list, function(fit) {
        if (inherits(fit, "try-error") || is.null(fit$par) || length(fit$par) < NXS + NXO + 3) return(NA_real_)
        fit$par[NXS + NXO + 3]
      }, numeric(1))
    )

  } else {
    start <- as.numeric(start)

    expected_length <- NXS + NXO + 3
    if (length(start) != expected_length) {
      stop("Argument 'start' must have length ", expected_length,
           " = NXS + NXO + 3.")
    }

    theta_SK <- fit_one_SK(start)

    if (inherits(theta_SK, "try-error") || !is.finite(theta_SK$value)) {
      stop("HeckmanSK optimization failed from the supplied 'start' values.")
    }

    lambda_search <- data.frame(
      lambda_start = start[expected_length],
      logLik       = theta_SK$value,
      convergence  = theta_SK$convergence,
      message      = if (is.null(theta_SK$message)) NA_character_ else as.character(theta_SK$message),
      sigma_hat    = theta_SK$par[NXS + NXO + 1],
      rho_hat      = theta_SK$par[NXS + NXO + 2],
      lambda_hat   = theta_SK$par[NXS + NXO + 3]
    )
  }

  #############################
  # Organize and return results
  #############################
  names(theta_SK$par) <- c(colnames(XS), colnames(XO), "sigma", "rho", "lambda")

  result <- list(
    coefficients   = theta_SK$par,
    value          = theta_SK$value,
    loglik         = theta_SK$value,
    gradient       = gradlik_SK(theta_SK$par),
    counts         = theta_SK$counts[2],
    hessian        = theta_SK$hessian,
    fisher_infoSK  = .ssmodels_safe_vcov(theta_SK$hessian),
    prop_sigmaSK   = .ssmodels_safe_se(.ssmodels_safe_vcov(theta_SK$hessian)),
    level          = YSLevels,
    nObs           = length(YS),
    nParam         = length(start),
    N0             = sum(YS == 0),
    N1             = sum(YS == 1),
    NXS            = ncol(XS),
    NXO            = ncol(XO),
    df             = length(YS) - length(start),
    aic            = -2 * theta_SK$value + 2 * length(start),
    bic            = -2 * theta_SK$value + length(start) * log(length(YS)),
    initial.value  = start,
    lambda_search  = lambda_search,
    lambda_start   = start[length(start)],
    multi_start    = nrow(lambda_search) > 1L
  )

  class(result) <- c("HeckmanSK", class(theta_SK))
  return(result)
}
