
#' Likelihood for composition vectors
#'
#' Returns the log-likelihood for composition data, e.g., length, age, or stock composition,
#' with various statistical distributions supported.
#'
#' @param obs A vector of observed values. Internally converted to proportions.
#' @param pred A vector of predicted values. Same length as `obs`. Internally converted to proportions.
#' @param type Character for the desired distribution
#' @param N Numeric, the sample size corresponding to `obs` for multinomial or Dirichlet multinomial distributions.
#' @param theta Numeric, the linear (`type = "dirmult1"`) or saturating (`type = "dirmult2"`) Dirichlet-multinomial parameter, respectively. See Thorson et al. (2017)
#' @param stdev Numeric or vectorized for `obs`, the likelihood standard deviation for lognormal distribution.
#' @return Numeric representing the log-likelihood.
#'
#' @details
#' Observed and predicted vectors are internally converted to proportions.
#'
#' For `type = "lognormal"` or `"logitnormal"`, zero observations are removed from the likelihood calculation.
#' @references
#' Thorson et al. 2017. Model-based estimates of effective sample size in stock assessment models using the
#' Dirichlet-multinomial distribution. Fish. Res. 192:84-93. \doi{10.1016/j.fishres.2016.06.005}
#' @examples
#' M <- 0.1
#' age <- seq(1:10)
#' pred <- exp(-M * age)
#' obs <- pred * rlnorm(10, sd = 0.05)
#' like_comp(obs, pred, N = 10, type = "multinomial")
#' like_comp(obs, pred, N = 100, type = "multinomial")
#' like_comp(obs, pred, N = 10, type = "dirmult1", theta = 1)
#' like_comp(obs, pred, N = 10, type = "dirmult1", theta = 20)
#' @importFrom stats rmultinom rnorm
#' @export
like_comp <- function(obs, pred, type = c("multinomial", "dirmult1", "dirmult2", "lognormal", "logitnormal"),
                      N = sum(obs), theta, stdev) {

  if (!inherits(obs, "simref") && (all(is.na(obs)) || !sum(obs))) { # estimation or report mode
    v <- if (inherits(pred, "advector")) advector(0) else 0
    return(v)
  }

  stopifnot(length(obs) == length(pred))
  type <- match.arg(type)

  #p_pred <- (pred + 1e-8)/sum(pred + 1e-8) # This doesn't work
  pred <- CondExpGt(pred, 1e-8, pred, 1e-8)
  p_pred <- pred/sum(pred)
  p_obs <- obs/sum(obs)

  if (type == "multinomial") {
    if (inherits(obs, "simref")) {
      if (sum(pred) && !is.na(N)) {
        obs[] <- stats::rmultinom(1, size = N, prob = p_pred)
      } else {
        obs[] <- NA
      }
    } else {
      # Do not use stats::dmultinom which rounds observations to whole numbers!
      v <- dmultinom_(N * p_obs, prob = p_pred, log = TRUE)
    }

  } else if (grepl("dirmult", type)) {

    if (type == "dirmult1") {
      alpha <- theta * N * p_pred
    } else if (type == "dirmult2") {
      alpha <- theta * p_pred
    }

    if (inherits(obs, "simref")) {
      if (!requireNamespace("RTMBdist", quietly = TRUE)) {
        stop("Need the RTMBdist package to simulate from the Dirichlet-multinomial distribution")
      }
      if (sum(pred) && !is.na(N)) {
        obs[] <- RTMBdist::rdirmult(1, size = N, alpha = alpha)
      } else {
        obs[] <- NA
      }
    } else {
      v <- ddirmult_(obs, size = N, alpha = alpha, log = TRUE)
    }

  } else if (type == "lognormal") {

    if (missing(stdev)) stdev <- 1/sqrt(p_pred)

    if (inherits(obs, "simref")) {
      if (sum(pred)) {
        obs[] <- exp(stats::rnorm(length(pred), log(p_pred), stdev))
      } else {
        obs[] <- NA
      }
    } else {
      resid <- p_obs/p_pred
      v <- dnorm(log(resid[obs > 0]), 0, stdev[obs > 0], log = TRUE) |> sum()
    }

  } else if (type == "logitnormal") {

    if (missing(stdev)) stdev <- 1/sqrt(p_pred)

    i_fit <- obs > 0
    i_ref <- rep(FALSE, length(obs))
    i_ref[which(i_fit)[1]] <- TRUE

    if (inherits(obs, "simref")) {
      if (sum(pred)) {
        xpred <- log(p_pred[!i_ref]/p_pred[i_ref])
        xsamp <- stats::rnorm(length(xpred), xpred, stdev[!i_ref])
        y <- exp(xsamp)/(1 + exp(xsamp))

        obs[!i_ref] <- y
        obs[i_ref] <- 1 - sum(y)
      } else {
        obs[] <- NA
      }
    } else {
      xobs <- log(p_obs[i_fit & !i_ref]/p_obs[i_ref])
      xpred <- log(p_pred[i_fit & !i_ref]/p_pred[i_ref])

      v <- dnorm(xobs, xpred, stdev[i_fit & !i_ref], log = TRUE) |> sum()
    }

  }

  return(v)
}

#' Likelihood for CKMR
#'
#' Returns the log-likelihood for a set of pairwise comparisons. For a parent-offspring pair, a comparison
#' is defined by the capture year of parent, capture age of parent, and birth year of offspring.
#' For a half-sibling pair, a comparison is defined by the cohort year of each sibling.
#' Binomial and Poisson distributions are supported (Conn et al. 2020).
#'
#' @param n The number of pairwise comparisons
#' @param m The number of kinship matches
#' @param p The probability of kinship match
#' @param type The statistical distribution for the likelihood calculation
#' @return Numeric representing the log-likelihood.
#' @seealso [calc_POP()] [calc_HSP()]
#' @references
#' Conn, P.B. et al. 2020. Robustness of close-kin mark-recapture estimators to dispersal
#' limitation and spatially varying sampling probabilities. Ecol. Evol. 10: 5558-5569. \doi{10.1002/ece3.6296}
#' @export
like_CKMR <- function(n, m, p, type = c("binomial", "poisson")) {
  type <- match(type)
  if (is.null(n) || all(!n)) {
    v <- 0
  } else if (type == "binomial") {
    v <- dbinom(m, n, p, log = TRUE)
  } else if (type == "poisson") {
    v <- dpois(m, n * p, log = TRUE)
  }
  return(v)
}

ddirmult_ <- function(x, size, alpha, log = FALSE) {
  x <- size * x/sum(x)
  alpha0 <- sum(alpha)
  val <- lgamma(alpha0) + lgamma(size + 1) - lgamma(alpha0 + size)
  val2 <- lgamma(x + alpha) - lgamma(alpha) - lgamma(x + 1)

  log_res <- val + sum(val2)

  if (log) log_res else exp(log_res)
}

dmultinom_ <- function(x, size = NULL, prob, log = FALSE) {
  K <- length(prob)
  if (length(x) != K)
    stop("x[] and prob[] must be equal length vectors.")
  if (any(x < 0))
    stop("'x' must be non-negative")
  N <- sum(x)
  if (is.null(size))
    size <- N
  else if (size != N)
    stop("size != sum(x), i.e. one is wrong")

  r <- lgamma(size + 1) + sum(x * log(prob) - lgamma(x + 1))
  if (log)
    r
  else exp(r)
}
