# Likelihood functions and gradient functions for xreg2

# ---------------------------------------------------------------------------
# Internal: negative log-likelihood dispatcher
# ---------------------------------------------------------------------------

negll_list2 <- function(par, dtalist, ...) {
  sum(sapply(dtalist, function(dtl) {
    do.call(dtl$ll_fun, args = c(list(dtalist = dtl, par = par), list(...)))
  }))
}

# Internal: gradient dispatcher
gradll_list2 <- function(par, dtalist, ...) {
  rowSums(as.matrix(sapply(dtalist, function(dtl) {
    do.call(dtl$grad_fun, args = c(list(dtalist = dtl, par = par), list(...)))
  })))
}

# ---------------------------------------------------------------------------
# Internal: forward gradient for continuous normal
# ---------------------------------------------------------------------------

grad_cont_normal2 <- function(par, dtalist, ret_mat = FALSE, eps = 1e-7, centered = TRUE) {

  if (centered) par <- par - eps / 2

  d_df <- dtalist$unique
  f_df <- dtalist$full
  nobs <- NROW(d_df)

  names(par) <- dtalist$coeff_names
  plen  <- length(par)
  plen1 <- plen + 1

  parm <- as.data.frame(as.list(par))
  parm <- parm[rep(1, plen + 1), ]

  for (pn in 1:plen) {
    parm[pn + 1, pn] <- parm[pn + 1, pn] + eps
  }
  par <- parm

  formulas <- dtalist$formulas

  for (formulanum in seq_along(formulas)) {
    formula  <- formulas[[formulanum]]
    new_var  <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    tmp <- sapply(1:plen1, function(x) {
      theseargs <- c(
        as.list(par),
        as.list(dtalist$fixed),
        as.list(d_df[, colnames(d_df)[colnames(d_df) %in% all.vars(formula)], drop = FALSE])
      )
      theseargs <- lapply(theseargs, function(thisarg) {
        if (!is.null(dim(thisarg)) && dim(thisarg)[2] == plen1) return(thisarg[, x])
        if (length(thisarg) == plen1) return(thisarg[x])
        thisarg
      })
      with(theseargs, eval(parse(text = formula)))
    })
    if (is.null(dim(tmp)) && length(tmp) == plen1)
      tmp <- matrix(tmp, nrow = nobs, ncol = plen1, byrow = TRUE)
    d_df[, new_var] <- tmp
  }

  pvals <- matrix(NA, nrow = nobs, ncol = plen1)

  sel  <- d_df$internal_type == 1
  pvals[sel, ]  <- dnorm(x    = d_df$Xb[sel, ],
                         mean = d_df$internal_lb[sel],
                         sd   = d_df$sigma_est[sel, ],
                         log  = TRUE)
  sel <- !sel
  pvals[sel, ] <- log(
    pnorm(q = d_df$Xb[sel, ], mean = d_df$internal_lb[sel],
          sd = d_df$sigma_est[sel, ], lower.tail = TRUE,  log.p = FALSE) -
    pnorm(q = d_df$Xb[sel, ], mean = d_df$internal_ub[sel],
          sd = d_df$sigma_est[sel, ], lower.tail = TRUE,  log.p = FALSE)
  )

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  pvals <- pvals[f_df$umatch, ] * f_df$internal_count
  if (ret_mat) return(pvals)
  pvals <- colSums(pvals)
  pvals <- (pvals[1] - pvals[2:plen1]) / eps
  return(pvals)
}

# ---------------------------------------------------------------------------
# Internal: forward gradient for dichotomous logistic
# ---------------------------------------------------------------------------

grad_dich_logistic2 <- function(par, dtalist, ret_mat = FALSE, eps = 1e-7, centered = TRUE) {

  if (centered) par <- par - eps / 2

  d_df <- dtalist$unique
  f_df <- dtalist$full
  nobs <- NROW(d_df)

  names(par) <- dtalist$coeff_names
  plen  <- length(par)
  plen1 <- plen + 1

  parm <- as.data.frame(as.list(par))
  parm <- parm[rep(1, plen + 1), ]

  for (pn in 1:plen) {
    parm[pn + 1, pn] <- parm[pn + 1, pn] + eps
  }
  par <- parm

  formulas <- dtalist$formulas

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    tmp <- sapply(1:plen1, function(x) {
      theseargs <- c(as.list(par[x, ]), as.list(dtalist$fixed))
      with(theseargs, with(d_df, eval(parse(text = formula))))
    })
    if (is.null(dim(tmp)) && length(tmp) == plen1)
      tmp <- matrix(tmp, nrow = nobs, ncol = plen1, byrow = TRUE)
    d_df[, new_var] <- tmp
  }

  Xb <- d_df$Xb
  if ("theta_est" %in% colnames(d_df)) Xb <- Xb * d_df$theta_est

  logistic_tmp <- .5 + .5 * tanh(Xb / 2)
  pvals <- log(d_df$internal_lb *      logistic_tmp +
               (1 - d_df$internal_lb) * (1 - logistic_tmp))
  pvals <- matrix(pvals, nrow = nobs, ncol = plen1)

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  pvals <- pvals[f_df$umatch, ] * f_df$internal_count
  if (ret_mat) return(pvals)
  pvals <- colSums(pvals)
  pvals <- (pvals[1] - pvals[2:plen1]) / eps
  return(pvals)
}

# ---------------------------------------------------------------------------
# Exported: continuous normal log-likelihood (fast list-based implementation)
# ---------------------------------------------------------------------------

#' Continuous normal log-likelihood for xreg2
#'
#' @title Continuous Normal Negative Log-Likelihood
#' @description Computes the negative log-likelihood for a censored/interval
#'   normal regression model. Designed as a \code{p_fun} argument for
#'   \code{\link{xreg2_control}}.
#'
#' @param par Named numeric vector of current parameter values.
#' @param dtalist Internal data list produced by \code{xreg2_fit}.
#' @param log.p Logical; use log-scale probabilities. Default \code{TRUE}.
#' @param aggregate.p Logical; return scalar negloglik. Default \code{TRUE}.
#' @param print_pars Logical; print parameters and value. Default \code{FALSE}.
#'
#' @return Scalar negative log-likelihood (when \code{aggregate.p = TRUE}).
#'
#' @examples
#' # Typically used as p_fun in xreg2_control():
#' # ctrl <- xreg2_control(formulas = list(y ~ INTERCEPT + x * SLOPE),
#' #                        p_fun = xreg2_cont_normal)
#'
#' @export
xreg2_cont_normal <- function(par, dtalist, log.p = TRUE, aggregate.p = TRUE,
                               print_pars = FALSE) {
  names(par) <- dtalist$coeff_names
  if (print_pars)
    message(do.call(paste, lapply(names(par),
                                  function(x) paste0(x, ": ", round(par[[x]], 3)))))

  d_l  <- as.list(dtalist$unique)
  nobs <- length(d_l[[1]])

  lapply(dtalist$formulas, function(formula) {
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(names(d_l), dtalist$valuevar)) new_var <- "Xb"
    d_l[[new_var]] <<- rep_len(
      eval(formula[[3]], envir = c(as.list(c(par, dtalist$fixed)), d_l)),
      nobs)
    NULL
  })

  logf  <- if (log.p) function(x) log(x) else function(x) x
  pvals <- matrix(NA, nrow = nobs, ncol = 1)

  sel <- which(d_l$internal_type == 1)
  pvals[sel, ] <- dnorm(x    = d_l$internal_lb[sel],
                        mean = d_l$Xb[sel],
                        sd   = d_l$sigma_est[sel],
                        log  = log.p)

  sel <- which(d_l$internal_type != 1)
  pvals[sel, ] <- logf(
    pnorm(q = d_l$internal_lb[sel], mean = d_l$Xb[sel],
          sd = d_l$sigma_est[sel], lower.tail = FALSE, log.p = FALSE) -
    pnorm(q = d_l$internal_ub[sel], mean = d_l$Xb[sel],
          sd = d_l$sigma_est[sel], lower.tail = FALSE, log.p = FALSE)
  )

  if (!aggregate.p) return(pvals)

  if (log.p) {
    if ("boundNA" %in% names(dtalist))
      if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
    pvals[pvals == -Inf] <- log(.Machine$double.xmin)
    pvals <- -sum(pvals[dtalist$full$umatch, ] * dtalist$full$internal_count)
  } else {
    if ("boundNA" %in% names(dtalist))
      if (dtalist$boundNA) pvals[is.na(pvals)] <- .Machine$double.xmin
    pvals[pvals == 0] <- .Machine$double.xmin
    pvals <- -prod(pvals[dtalist$full$umatch, ] * dtalist$full$internal_count)
  }

  if (is.na(pvals)) {
    if (print_pars) message("Encountered NA")
    pvals <- 10e10
  }
  if (print_pars) message(paste0("Value: ", pvals))
  return(pvals)
}

attr(x = xreg2_cont_normal, which = "required") <- list(
  sigma_est = list(
    formula      = sigma_est ~ exp(LN_SIGMA),
    start_values = c(LN_SIGMA = 0)
  )
)

# ---------------------------------------------------------------------------
# Exported: general continuous log-likelihood (multi-distribution)
# ---------------------------------------------------------------------------

#' General continuous log-likelihood for xreg2
#'
#' @title General Continuous Negative Log-Likelihood
#' @description Computes the negative log-likelihood for a censored/interval
#'   regression model supporting multiple distributions (normal, lognormal,
#'   exponential, beta, gamma, weibull). Designed as a \code{p_fun} argument
#'   for \code{\link{xreg2_control}}.
#'
#' @param par Named numeric vector of current parameter values.
#' @param dtalist Internal data list produced by \code{xreg2_fit}.
#' @param log.p Logical; use log-scale probabilities. Default \code{TRUE}.
#' @param aggregate.p Logical; return scalar negloglik. Default \code{TRUE}.
#' @param print_pars Logical; print parameters and value. Default \code{FALSE}.
#' @param dist Character; distribution family. One of \code{"normal"},
#'   \code{"lognormal"}, \code{"exp"}/\code{"exponential"}, \code{"beta"},
#'   \code{"gamma"}, \code{"weibull"}. Default \code{"normal"}.
#'
#' @return Scalar negative log-likelihood (when \code{aggregate.p = TRUE}).
#'
#' @examples
#' # ctrl <- xreg2_control(formulas = list(y ~ INTERCEPT + x * SLOPE),
#' #                        p_fun = xreg2_cont,
#' #                        start_values = c(INTERCEPT = 0, SLOPE = 0, LN_SIGMA = 0))
#'
#' @export
xreg2_cont <- function(par, dtalist, log.p = TRUE, aggregate.p = TRUE,
                        print_pars = FALSE, dist = "normal") {
  p_type <- dtalist$type
  names(par) <- dtalist$coeff_names

  formulas <- dtalist$formulas
  d_df <- dtalist$unique
  f_df <- dtalist$full

  nobs <- NROW(d_df)

  args <- c(as.list(par), as.list(dtalist$fixed))
  d_df <- cbind(d_df, args)

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    d_df[, new_var] <- with(d_df, eval(parse(text = formula)))
  }

  pvals <- matrix(NA, nrow = nobs, ncol = 1)
  sel   <- d_df$internal_type == 1

  if (dist == "normal") {
    pvals[sel, ] <- dnorm(x    = d_df$internal_lb[sel],
                          mean = d_df$Xb[sel],
                          sd   = d_df$sigma_est[sel],
                          log  = log.p)
    sel <- !sel
    if (log.p) {
      pvals[sel, ] <- log(
        pnorm(q = d_df$internal_lb[sel], mean = d_df$Xb[sel],
              sd = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE) -
        pnorm(q = d_df$internal_ub[sel], mean = d_df$Xb[sel],
              sd = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE))
    } else {
      pvals[sel, ] <-
        pnorm(q = d_df$internal_lb[sel], mean = d_df$Xb[sel],
              sd = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE) -
        pnorm(q = d_df$internal_ub[sel], mean = d_df$Xb[sel],
              sd = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE)
    }

  } else if (dist == "lognormal") {
    pvals[sel, ] <- dlnorm(x       = d_df$internal_lb[sel],
                           meanlog = d_df$Xb[sel],
                           sdlog   = d_df$sigma_est[sel],
                           log     = log.p)
    sel <- !sel
    if (log.p) {
      pvals[sel, ] <- log(
        plnorm(q = d_df$internal_lb[sel], meanlog = d_df$Xb[sel],
               sdlog = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE) -
        plnorm(q = d_df$internal_ub[sel], meanlog = d_df$Xb[sel],
               sdlog = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE))
    } else {
      pvals[sel, ] <-
        plnorm(q = d_df$internal_lb[sel], meanlog = d_df$Xb[sel],
               sdlog = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE) -
        plnorm(q = d_df$internal_ub[sel], meanlog = d_df$Xb[sel],
               sdlog = d_df$sigma_est[sel], lower.tail = FALSE, log.p = FALSE)
    }

  } else if (dist %in% c("exp", "exponential")) {
    d_df$Xb[sel & d_df$Xb < 0] <- 0
    pvals[sel, ] <- dexp(x    = d_df$internal_lb[sel],
                         rate = d_df$Xb[sel],
                         log  = log.p)
    sel <- !sel
    if (log.p) {
      pvals[sel, ] <- log(
        pexp(q = d_df$internal_lb[sel], rate = d_df$Xb[sel],
             lower.tail = FALSE, log.p = FALSE) -
        pexp(q = d_df$internal_ub[sel], rate = d_df$Xb[sel],
             lower.tail = FALSE, log.p = FALSE))
    } else {
      pvals[sel, ] <-
        pexp(q = d_df$internal_lb[sel], rate = d_df$Xb[sel],
             lower.tail = FALSE, log.p = FALSE) -
        pexp(q = d_df$internal_ub[sel], rate = d_df$Xb[sel],
             lower.tail = FALSE, log.p = FALSE)
    }

  } else if (dist == "beta") {
    d_df$Xb[sel & d_df$Xb < 0] <- 0
    pvals[sel, ] <- dbeta(x      = d_df$internal_lb[sel] * d_df$SCALE[sel],
                          shape1 = d_df$shape1[sel],
                          shape2 = d_df$shape2[sel],
                          log    = log.p)
    sel <- !sel
    if (log.p) {
      pvals[sel, ] <- log(
        pbeta(q = d_df$internal_lb[sel] * d_df$SCALE[sel],
              shape1 = d_df$shape1[sel], shape2 = d_df$shape2[sel],
              lower.tail = FALSE, log.p = FALSE) -
        pbeta(q = d_df$internal_ub[sel] * d_df$SCALE[sel],
              shape1 = d_df$shape1[sel], shape2 = d_df$shape2[sel],
              lower.tail = FALSE, log.p = FALSE))
    } else {
      pvals[sel, ] <-
        pbeta(q = d_df$internal_lb[sel] * d_df$SCALE[sel],
              shape1 = d_df$shape1[sel], shape2 = d_df$shape2[sel],
              lower.tail = FALSE, log.p = FALSE) -
        pbeta(q = d_df$internal_ub[sel] * d_df$SCALE[sel],
              shape1 = d_df$shape1[sel], shape2 = d_df$shape2[sel],
              lower.tail = FALSE, log.p = FALSE)
    }

  } else if (dist == "gamma") {
    d_df$Xb[sel & d_df$Xb < 0] <- 0
    pvals[sel, ] <- dgamma(x     = d_df$internal_lb[sel],
                           shape = d_df$shape[sel],
                           rate  = d_df$rate[sel],
                           log   = log.p)
    sel <- !sel
    if (log.p) {
      pvals[sel, ] <- log(
        pgamma(q = d_df$internal_lb[sel], shape = d_df$shape[sel],
               rate = d_df$rate[sel], lower.tail = FALSE, log.p = FALSE) -
        pgamma(q = d_df$internal_ub[sel], shape = d_df$shape[sel],
               rate = d_df$rate[sel], lower.tail = FALSE, log.p = FALSE))
    } else {
      pvals[sel, ] <-
        pgamma(q = d_df$internal_lb[sel], shape = d_df$shape[sel],
               rate = d_df$rate[sel], lower.tail = FALSE, log.p = FALSE) -
        pgamma(q = d_df$internal_ub[sel], shape = d_df$shape[sel],
               rate = d_df$rate[sel], lower.tail = FALSE, log.p = FALSE)
    }

  } else if (dist == "weibull") {
    d_df$Xb[sel & d_df$Xb < 0] <- 0
    pvals[sel, ] <- dweibull(x     = d_df$internal_lb[sel],
                             shape = d_df$shape[sel],
                             scale = d_df$scale[sel],
                             log   = log.p)
    sel <- !sel
    if (log.p) {
      pvals[sel, ] <- log(
        pweibull(q = d_df$internal_lb[sel], shape = d_df$shape[sel],
                 scale = d_df$scale[sel], lower.tail = FALSE, log.p = FALSE) -
        pweibull(q = d_df$internal_ub[sel], shape = d_df$shape[sel],
                 scale = d_df$scale[sel], lower.tail = FALSE, log.p = FALSE))
    } else {
      pvals[sel, ] <-
        pweibull(q = d_df$internal_lb[sel], shape = d_df$shape[sel],
                 scale = d_df$scale[sel], lower.tail = FALSE, log.p = FALSE) -
        pweibull(q = d_df$internal_ub[sel], shape = d_df$shape[sel],
                 scale = d_df$scale[sel], lower.tail = FALSE, log.p = FALSE)
    }
  }

  if (!aggregate.p) return(pvals)

  if (log.p) {
    if ("boundNA" %in% names(dtalist))
      if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
    pvals[pvals == -Inf] <- log(.Machine$double.xmin)
    pvals <- pvals[f_df$umatch, ]
    pvalso <- -sum(pvals * f_df$internal_count)
  } else {
    if ("boundNA" %in% names(dtalist))
      if (dtalist$boundNA) pvals[is.na(pvals)] <- .Machine$double.xmin
    pvals[pvals == 0] <- .Machine$double.xmin
    pvals <- pvals[f_df$umatch, ]
    pvalso <- -prod(pvals * f_df$internal_count)
  }
  return(pvalso)
}

# ---------------------------------------------------------------------------
# Exported: dichotomous logistic log-likelihood
# ---------------------------------------------------------------------------

#' Dichotomous logistic log-likelihood for xreg2
#'
#' @title Dichotomous Logistic Negative Log-Likelihood
#' @description Computes the negative log-likelihood for a binary logistic
#'   regression model. Designed as a \code{p_fun} argument for
#'   \code{\link{xreg2_control}}.
#'
#' @param par Named numeric vector of current parameter values.
#' @param dtalist Internal data list produced by \code{xreg2_fit}.
#'
#' @return Scalar negative log-likelihood.
#'
#' @examples
#' # ctrl <- xreg2_control(formulas = list(y ~ INTERCEPT + x * SLOPE),
#' #                        p_fun = xreg2_dich_logistic,
#' #                        start_values = c(INTERCEPT = 0, SLOPE = 0))
#'
#' @export
xreg2_dich_logistic <- function(par, dtalist) {
  p_type <- dtalist$type
  names(par) <- dtalist$coeff_names

  formulas <- dtalist$formulas
  d_df <- dtalist$unique
  f_df <- dtalist$full

  nobs <- NROW(d_df)

  args <- c(as.list(par), as.list(dtalist$fixed))

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    d_df[, new_var] <- with(args, with(d_df, eval(parse(text = formula))))
  }

  Xb <- d_df$Xb
  if ("theta_est" %in% colnames(d_df)) Xb <- Xb * d_df$theta_est

  logistic_tmp <- .5 + .5 * tanh(Xb / 2)
  pvals <- log(d_df$internal_lb *      logistic_tmp +
               (1 - d_df$internal_lb) * (1 - logistic_tmp))
  pvals <- matrix(pvals, nrow = nobs, ncol = 1)

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  pvals <- pvals[f_df$umatch, ]
  tout  <- -sum(pvals * f_df$internal_count)
  return(tout)
}

# ---------------------------------------------------------------------------
# Exported: continuous normal with random intercept (Gauss-Hermite quadrature)
# ---------------------------------------------------------------------------

#' Continuous normal log-likelihood with random intercept
#'
#' @title Continuous Normal Negative Log-Likelihood with Random Intercept
#' @description Computes the negative log-likelihood for a mixed-effects normal
#'   regression model with a random intercept, integrated via Gauss-Hermite
#'   quadrature. Designed as a \code{p_fun} argument for
#'   \code{\link{xreg2_control}}.
#'
#' @param par Named numeric vector of current parameter values. Must include
#'   either \code{LN_OMEGA} or \code{OMEGA} (between-subject SD).
#' @param dtalist Internal data list produced by \code{xreg2_fit}. Must
#'   include an \code{internal_id} column.
#' @param print_pars Logical; print parameters and value. Default \code{FALSE}.
#'
#' @return Scalar negative log-likelihood.
#'
#' @examples
#' # ctrl <- xreg2_control(formulas = list(y ~ INTERCEPT + x * SLOPE),
#' #                        p_fun = xreg2_cont_r_normal)
#'
#' @export
xreg2_cont_r_normal <- function(par, dtalist, print_pars = FALSE) {
  names(par) <- dtalist$coeff_names

  if (print_pars)
    message(do.call(paste, lapply(names(par),
                                  function(x) paste0(x, ": ", round(par[[x]], 3)))))

  formulas <- dtalist$formulas
  d_df <- dtalist$unique
  f_df <- dtalist$full

  args <- c(as.list(par), as.list(dtalist$fixed))

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    d_df[, new_var] <- with(args, with(d_df, eval(parse(text = formula))))
  }

  nobs <- NROW(d_df)

  if (!sum(c("LN_OMEGA", "OMEGA") %in% names(par)) == 1)
    stop("Ambiguous or non-unique OMEGA/LN_OMEGA parameter defined.")
  if ("LN_OMEGA" %in% names(par)) betwSD <- exp(par["LN_OMEGA"])
  if ("OMEGA"    %in% names(par)) betwSD <- par["OMEGA"]

  uid  <- sort(unique(f_df$internal_id))
  nuid <- NROW(uid)

  GH_order <- 8L
  if (!is.environment((xregenv <- getOption("xreg")))) {
    options(xreg = (xregenv <- new.env(parent = emptyenv())))
  }
  if (exists("Gauss_Hermite_Order", envir = xregenv)) {
    GH_order <- get("Gauss_Hermite_Order", envir = xregenv)
  } else {
    assign(x = "Gauss_Hermite_Order", value = GH_order, envir = xregenv)
  }
  ws   <- as.matrix(xreg_GHS[[GH_order]]$w)
  xs   <- xreg_GHS[[GH_order]]$x
  mus  <- sqrt(2) * betwSD * xs
  tmpm <- matrix(NA, nrow = nobs, ncol = GH_order)
  Xbm  <- Outer(y = d_df$Xb, x = mus, "+")

  sel  <- d_df$internal_type == 1
  sel2 <- !sel
  tmpm[sel, ]  <- dnorm(x    = d_df$internal_lb[sel],
                        mean = Xbm[sel, ],
                        sd   = d_df$sigma_est[sel],
                        log  = TRUE)
  tmpm[sel2, ] <- Log(
    pnorm(q = d_df$internal_lb[sel2], mean = Xbm[sel2, ],
          sd = d_df$sigma_est[sel2], lower.tail = FALSE) -
    pnorm(q = d_df$internal_ub[sel2], mean = Xbm[sel2, ],
          sd = d_df$sigma_est[sel2], lower.tail = FALSE)
  )

  tmpm <- tmpm[f_df$umatch, ] * f_df$internal_count

  pvals <- Log(1 / sqrt(pi) *
                 exp(apply(X      = tmpm,
                           MARGIN = 2,
                           FUN    = group.sum,
                           ina    = f_df$internal_id) +
                       rep_row(Log(ws), nuid)) %*% rep(1, GH_order))

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  tout <- -sum(pvals)
  if (is.na(tout)) {
    if (print_pars) message("Encountered NA")
    tout <- 10e10
  }
  if (print_pars) message(paste0("Value: ", tout))
  return(tout)
}

attr(x = xreg2_cont_r_normal, which = "required") <- list(
  sigma_est = list(
    formula      = sigma_est ~ exp(LN_SIGMA),
    start_values = c(LN_SIGMA = 0)
  ),
  omega_est = list(
    formula      = omega_est ~ exp(LN_OMEGA),
    start_values = c(OMEGA = 0)
  )
)

# ---------------------------------------------------------------------------
# Exported: continuous log-normal with random intercept
# ---------------------------------------------------------------------------

#' Continuous log-normal log-likelihood with random intercept
#'
#' @title Continuous Log-Normal Negative Log-Likelihood with Random Intercept
#' @description Computes the negative log-likelihood for a mixed-effects
#'   log-normal regression model with a random intercept, integrated via
#'   Gauss-Hermite quadrature. Designed as a \code{p_fun} argument for
#'   \code{\link{xreg2_control}}.
#'
#' @param par Named numeric vector of current parameter values. Must include
#'   either \code{LN_OMEGA} or \code{OMEGA}.
#' @param dtalist Internal data list produced by \code{xreg2_fit}. Must
#'   include an \code{internal_id} column.
#'
#' @return Scalar negative log-likelihood.
#'
#' @examples
#' # ctrl <- xreg2_control(formulas = list(y ~ INTERCEPT + x * SLOPE),
#' #                        p_fun = xreg2_cont_r_lnorm)
#'
#' @export
xreg2_cont_r_lnorm <- function(par, dtalist) {
  names(par) <- dtalist$coeff_names

  formulas <- dtalist$formulas
  d_df <- dtalist$unique
  f_df <- dtalist$full

  args <- c(as.list(par), as.list(dtalist$fixed))

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    d_df[, new_var] <- with(args, with(d_df, eval(parse(text = formula))))
  }

  nobs <- NROW(d_df)

  if (!sum(c("LN_OMEGA", "OMEGA") %in% names(par)) == 1)
    stop("Ambiguous or non-unique OMEGA/LN_OMEGA parameter defined.")
  if ("LN_OMEGA" %in% names(par)) betwSD <- exp(par["LN_OMEGA"])
  if ("OMEGA"    %in% names(par)) betwSD <- par["OMEGA"]

  uid  <- sort(unique(f_df$internal_id))
  nuid <- NROW(uid)

  GH_order <- 8L
  if (!is.environment((xregenv <- getOption("xreg")))) {
    options(xreg = (xregenv <- new.env(parent = emptyenv())))
  }
  if (exists("Gauss_Hermite_Order", envir = xregenv)) {
    GH_order <- get("Gauss_Hermite_Order", envir = xregenv)
  } else {
    assign(x = "Gauss_Hermite_Order", value = GH_order, envir = xregenv)
  }
  ws   <- as.matrix(xreg_GHS[[GH_order]]$w)
  xs   <- xreg_GHS[[GH_order]]$x
  mus  <- sqrt(2) * betwSD * xs
  tmpm <- matrix(NA, nrow = nobs, ncol = GH_order)
  Xbm  <- Outer(y = d_df$Xb, x = mus, "+")

  sel  <- d_df$internal_type == 1
  sel2 <- !sel
  tmpm[sel, ]  <- dlnorm(x       = d_df$internal_lb[sel],
                         meanlog = Xbm[sel, ],
                         sdlog   = d_df$sigma_est[sel],
                         log     = TRUE)
  tmpm[sel2, ] <- Log(
    plnorm(q = d_df$internal_ub[sel2], meanlog = Xbm[sel2, ],
           sdlog = d_df$sigma_est[sel2]) -
    plnorm(q = d_df$internal_lb[sel2], meanlog = Xbm[sel2, ],
           sdlog = d_df$sigma_est[sel2])
  )

  tmpm <- tmpm[f_df$umatch, ] * f_df$internal_count

  pvals <- Log(1 / sqrt(pi) *
                 exp(apply(X      = tmpm,
                           MARGIN = 2,
                           FUN    = group.sum,
                           ina    = f_df$internal_id) +
                       rep_row(Log(ws), nuid)) %*% rep(1, GH_order))

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  tout <- -sum(pvals)
  return(tout)
}

# ---------------------------------------------------------------------------
# Internal: gradient for cont_r_normal (3D array version)
# ---------------------------------------------------------------------------

grad_cont_r_normal1 <- function(par, dtalist, ret_mat = FALSE, eps = 1e-7, centered = TRUE) {

  if (centered) par <- par - eps / 2

  d_df <- dtalist$unique
  f_df <- dtalist$full
  nobs <- NROW(d_df)

  names(par) <- dtalist$coeff_names
  formulas   <- dtalist$formulas
  plen  <- length(par)
  plen1 <- plen + 1

  parm <- as.data.frame(as.list(par))
  parm <- parm[rep(1, plen + 1), ]
  for (pn in 1:plen) parm[pn + 1, pn] <- parm[pn + 1, pn] + eps

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    tmp <- sapply(1:plen1, function(x) {
      theseargs <- c(as.list(parm[x, ]), as.list(dtalist$fixed))
      with(theseargs, with(d_df, eval(parse(text = formula))))
    })
    if (is.null(dim(tmp)) && length(tmp) == plen1)
      tmp <- matrix(tmp, nrow = nobs, ncol = plen1, byrow = TRUE)
    d_df[, new_var] <- tmp
  }

  betwSD <- exp(parm[["LN_OMEGA"]])
  uid    <- sort(unique(f_df$internal_id))
  nuid   <- NROW(uid)

  GH_order <- 8L
  if (!is.environment((xregenv <- getOption("xreg")))) {
    options(xreg = (xregenv <- new.env(parent = emptyenv())))
  }
  if (exists("Gauss_Hermite_Order", envir = xregenv)) {
    GH_order <- get("Gauss_Hermite_Order", envir = xregenv)
  } else {
    assign(x = "Gauss_Hermite_Order", value = GH_order, envir = xregenv)
  }
  ws   <- as.matrix(xreg_GHS[[GH_order]]$w)
  xs   <- xreg_GHS[[GH_order]]$x
  mus  <- sqrt(2) * Outer(betwSD, xs)

  Xba  <- array(Outer(rep(0, GH_order), d_df$Xb, "+"), c(dim(d_df$Xb), GH_order))
  musa <- aperm(array(Outer(rep(0, nobs), mus, "+"), c(dim(mus), nobs)))
  pa   <- Xba <- Xba + musa

  sel  <- d_df$internal_type == 1
  sel2 <- !sel
  pa[sel, , ]  <- dnorm(x    = d_df$internal_lb[sel],
                        mean = Xba[sel, , ],
                        sd   = d_df$sigma_est[sel, ],
                        log  = TRUE)
  pa[sel2, , ] <- log(
    pnorm(q = d_df$internal_lb[sel2], mean = Xba[sel2, , ],
          sd = d_df$sigma_est[sel2, ], lower.tail = FALSE) -
    pnorm(q = d_df$internal_ub[sel2], mean = Xba[sel2, , ],
          sd = d_df$sigma_est[sel2, ], lower.tail = FALSE)
  )

  pa <- pa[f_df$umatch, , ]

  pvals <- apply(X = pa, MARGIN = 2, FUN = function(slice) {
    Log(1 / sqrt(pi) *
          exp(apply(X      = slice,
                    MARGIN = 2,
                    FUN    = group.sum,
                    ina    = f_df$internal_id) +
                rep_row(Log(ws), nuid)) %*% rep(1, GH_order))
  })

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  if (ret_mat) return(pvals)
  pvals <- colSums(pvals)
  pvals <- (pvals[1] - pvals[2:plen1]) / eps
  return(pvals)
}

# ---------------------------------------------------------------------------
# Internal: gradient for cont_r_normal (column-wise version, faster)
# ---------------------------------------------------------------------------

grad_cont_r_normal2 <- function(par, dtalist, ret_mat = FALSE, eps = 1e-7, centered = TRUE) {

  if (centered) par <- par - eps / 2

  d_df <- dtalist$unique
  f_df <- dtalist$full
  nobs <- NROW(d_df)

  names(par) <- dtalist$coeff_names
  formulas   <- dtalist$formulas
  plen  <- length(par)
  plen1 <- plen + 1

  parm <- as.data.frame(as.list(par))
  parm <- parm[rep(1, plen + 1), ]
  for (pn in 1:plen) parm[pn + 1, pn] <- parm[pn + 1, pn] + eps

  for (formulanum in seq_along(formulas)) {
    formula <- formulas[[formulanum]]
    new_var <- as.character(formula[[2]])
    if (new_var %in% c(colnames(d_df), dtalist$valuevar)) new_var <- "Xb"
    tmp <- sapply(1:plen1, function(x) {
      theseargs <- c(as.list(parm[x, ]), as.list(dtalist$fixed))
      with(theseargs, with(d_df, eval(parse(text = formula))))
    })
    if (is.null(dim(tmp)) && length(tmp) == plen1)
      tmp <- matrix(tmp, nrow = nobs, ncol = plen1, byrow = TRUE)
    d_df[, new_var] <- tmp
  }

  uid  <- sort(unique(f_df$internal_id))
  nuid <- NROW(uid)

  GH_order <- 8L
  if (!is.environment((xregenv <- getOption("xreg")))) {
    options(xreg = (xregenv <- new.env(parent = emptyenv())))
  }
  if (exists("Gauss_Hermite_Order", envir = xregenv)) {
    GH_order <- get("Gauss_Hermite_Order", envir = xregenv)
  } else {
    assign(x = "Gauss_Hermite_Order", value = GH_order, envir = xregenv)
  }
  ws <- as.matrix(xreg_GHS[[GH_order]]$w)
  xs <- xreg_GHS[[GH_order]]$x

  pvals <- sapply(1:plen1, function(x) {
    Xb     <- d_df$Xb[, x]
    betwSD <- exp(parm["LN_OMEGA"][x, ])
    mus    <- sqrt(2) * betwSD * xs
    tmpm   <- matrix(NA, nrow = nobs, ncol = GH_order)
    Xbm    <- Outer(y = Xb, x = mus, "+")

    sel  <- d_df$internal_type == 1
    sel2 <- !sel
    tmpm[sel, ]  <- dnorm(x    = d_df$internal_lb[sel],
                          mean = Xbm[sel, ],
                          sd   = d_df$sigma_est[sel, x],
                          log  = TRUE)
    tmpm[sel2, ] <- Log(
      pnorm(q = d_df$internal_lb[sel2], mean = Xbm[sel2, ],
            sd = d_df$sigma_est[sel2, x], lower.tail = FALSE) -
      pnorm(q = d_df$internal_ub[sel2], mean = Xbm[sel2, ],
            sd = d_df$sigma_est[sel2, x], lower.tail = FALSE)
    )

    tmpm <- tmpm[f_df$umatch, ] * f_df$internal_count

    Log(1 / sqrt(pi) *
          exp(apply(X      = tmpm,
                    MARGIN = 2,
                    FUN    = group.sum,
                    ina    = f_df$internal_id) +
                rep_row(Log(ws), nuid)) %*% rep(1, GH_order))
  })

  if ("boundNA" %in% names(dtalist))
    if (dtalist$boundNA) pvals[is.na(pvals)] <- log(.Machine$double.xmin)
  pvals[pvals == -Inf] <- log(.Machine$double.xmin)

  if (ret_mat) return(pvals)
  pvals <- colSums(pvals)
  pvals <- (pvals[1] - pvals[2:plen1]) / eps
  return(pvals)
}

# Internal: default gradient using numDeriv
grad_default <- function(func, x, method = "simple", ...) {
  numDeriv::grad(func = func, x = x, method = method, ...)
}
