#' Insurance models with varying benefits
#'
#' Functions for discrete and continuous contingent payment models with
#' varying benefits.
#'
#' @name insurance_varying_benefits
#' @keywords internal
NULL

# ------------------------------------------------------------
# Internal helpers
# ------------------------------------------------------------

.discrete_prob_k <- function(x, k, tbl = NULL, model = NULL, ...) {
  if (!is.null(tbl)) {
    p0 <- npx(tbl, x = x, n = k)
    p1 <- npx(tbl, x = x, n = k + 1)

    if (is.na(p0)) p0 <- 0
    if (is.na(p1)) p1 <- 0

    return(p0 - p1)
  }

  if (is.null(model)) {
    stop("Supply either tbl or model.", call. = FALSE)
  }

  tpx(k, x = x, model = model, ...) -
    tpx(k + 1, x = x, model = model, ...)
}

.max_k_varying <- function(x, tbl = NULL, model = NULL, ...,
                           tol = 1e-12, k_max = 5000) {
  for (k in 0:k_max) {
    if (!is.null(tbl)) {
      surv <- npx(tbl, x = x, n = k)
    } else {
      if (is.null(model)) {
        stop("Supply either tbl or model.", call. = FALSE)
      }
      surv <- tpx(k, x = x, model = model, ...)
    }

    if (is.na(surv)) {
      return(max(k - 1L, 0L))
    }

    if (!is.finite(surv) || surv <= tol) {
      return(k)
    }
  }

  k_max
}

.get_upper_varying_cont <- function(x, n = NULL, model, ...) {
  if (tolower(model) == "uniform") {
    dots <- list(...)

    if (is.null(dots$omega)) {
      stop("For model = 'uniform', omega must be supplied.", call. = FALSE)
    }

    upper <- dots$omega - x

    if (!is.null(n)) {
      upper <- pmin(upper, n)
    }

    return(pmax(upper, 0))
  }

  if (is.null(n)) {
    return(Inf)
  }

  n
}

.integrate_varying_cont <- function(f, lower, upper) {
  stats::integrate(
    f,
    lower = lower,
    upper = upper,
    rel.tol = 1e-9,
    subdivisions = 1000L
  )$value
}

# ------------------------------------------------------------
# Discrete varying-benefit models
# ------------------------------------------------------------

#' Increasing whole life insurance
#'
#' Computes
#' \deqn{(IA)_x = \sum_{k=0}^{\infty} (k+1) v^{k+1} \Pr(K_x = k).}
#'
#' @param x Age.
#' @param i Effective annual interest rate.
#' @param tbl Optional life table object.
#' @param model Optional parametric survival model.
#' @param ... Additional model parameters.
#' @param tol Numerical tolerance for truncation.
#' @param k_max Maximum number of terms.
#'
#' @return Numeric vector.
#' @export
IAx <- function(x, i, tbl = NULL, model = NULL, ...,
                tol = 1e-12, k_max = 5000) {
  x <- as.numeric(x)
  i <- .check_i_discrete(i)

  if (length(x) == 0L || any(!is.finite(x))) {
    stop("x must contain finite numeric values.", call. = FALSE)
  }

  if (any(x < 0)) {
    stop("x must be nonnegative.", call. = FALSE)
  }

  v <- 1 / (1 + i)

  sapply(x, function(xx) {
    kk_max <- .max_k_varying(
      xx,
      tbl = tbl,
      model = model,
      ...,
      tol = tol,
      k_max = k_max
    )

    k <- 0:kk_max

    probs <- vapply(
      k,
      function(kk) .discrete_prob_k(xx, kk, tbl = tbl, model = model, ...),
      numeric(1)
    )

    sum((k + 1) * v^(k + 1) * probs, na.rm = TRUE)
  })
}

#' Increasing n-year term insurance
#'
#' Computes
#' \deqn{(IA)_{x:\overline{n}|}^{1}
#' = \sum_{k=0}^{n-1} (k+1) v^{k+1} \Pr(K_x = k).}
#'
#' @param x Age.
#' @param n Term.
#' @param i Effective annual interest rate.
#' @param tbl Optional life table object.
#' @param model Optional parametric survival model.
#' @param ... Additional model parameters.
#'
#' @return Numeric vector.
#' @export
IAxn1 <- function(x, n, i, tbl = NULL, model = NULL, ...) {
  xn <- .recycle_x_n(x, n)
  x <- xn$x
  n <- xn$n
  i <- .check_i_discrete(i)

  v <- 1 / (1 + i)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    nn <- n[j]

    if (nn == 0) {
      return(0)
    }

    k <- 0:(nn - 1)

    probs <- vapply(
      k,
      function(kk) .discrete_prob_k(xx, kk, tbl = tbl, model = model, ...),
      numeric(1)
    )

    sum((k + 1) * v^(k + 1) * probs, na.rm = TRUE)
  })
}

#' Decreasing n-year term insurance
#'
#' Computes
#' \deqn{(DA)_{x:\overline{n}|}^{1}
#' = \sum_{k=0}^{n-1} (n-k) v^{k+1} \Pr(K_x = k).}
#'
#' @inheritParams IAxn1
#'
#' @return Numeric vector.
#' @export
DAxn1 <- function(x, n, i, tbl = NULL, model = NULL, ...) {
  xn <- .recycle_x_n(x, n)
  x <- xn$x
  n <- xn$n
  i <- .check_i_discrete(i)

  v <- 1 / (1 + i)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    nn <- n[j]

    if (nn == 0) {
      return(0)
    }

    k <- 0:(nn - 1)

    probs <- vapply(
      k,
      function(kk) .discrete_prob_k(xx, kk, tbl = tbl, model = model, ...),
      numeric(1)
    )

    sum((nn - k) * v^(k + 1) * probs, na.rm = TRUE)
  })
}

# ------------------------------------------------------------
# Continuous varying-benefit models
# ------------------------------------------------------------

#' Fully continuous increasing whole life insurance
#'
#' Computes
#' \deqn{(\bar{I}\bar{A})_x =
#' \int_0^\infty t v^t {}_tp_x \mu_{x+t}\,dt.}
#'
#' @param x Age.
#' @param i Effective annual interest rate.
#' @param model Parametric survival model.
#' @param ... Additional model parameters.
#'
#' @return Numeric vector.
#' @export
IbarAbarx <- function(x, i, model, ...) {
  args <- .recycle_cont_vectors(x, .check_i_cont(i))
  x <- args[[1L]]
  i <- args[[2L]]

  .check_x_cont(x)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    delta <- log(1 + i[j])

    upper <- .get_upper_varying_cont(xx, model = model, ...)

    if (upper <= 0) {
      return(0)
    }

    f <- function(t) {
      val <- t * exp(-delta * t) *
        tpx(t, x = xx, model = model, ...) *
        hazard0(xx + t, model = model, ...)
      val[!is.finite(val)] <- 0
      val
    }

    .integrate_varying_cont(f, 0, upper)
  })
}

#' Piecewise-continuous increasing whole life insurance
#'
#' Computes
#' \deqn{(I\bar{A})_x =
#' \int_0^\infty \lfloor t+1 \rfloor v^t {}_tp_x \mu_{x+t}\,dt.}
#'
#' @inheritParams IbarAbarx
#'
#' @return Numeric vector.
#' @export
IAbarx <- function(x, i, model, ...) {
  args <- .recycle_cont_vectors(x, .check_i_cont(i))
  x <- args[[1L]]
  i <- args[[2L]]

  .check_x_cont(x)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    delta <- log(1 + i[j])

    upper <- .get_upper_varying_cont(xx, model = model, ...)

    if (upper <= 0) {
      return(0)
    }

    f <- function(t) {
      val <- floor(t + 1) * exp(-delta * t) *
        tpx(t, x = xx, model = model, ...) *
        hazard0(xx + t, model = model, ...)
      val[!is.finite(val)] <- 0
      val
    }

    .integrate_varying_cont(f, 0, upper)
  })
}

#' Fully continuous increasing n-year term insurance
#'
#' Computes
#' \deqn{(\bar{I}\bar{A})_{x:\overline{n}|}^{1}
#' = \int_0^n t v^t {}_tp_x \mu_{x+t}\,dt.}
#'
#' @param x Age.
#' @param n Term.
#' @param i Effective annual interest rate.
#' @param model Parametric survival model.
#' @param ... Additional model parameters.
#'
#' @return Numeric vector.
#' @export
IbarAbarxn1 <- function(x, n, i, model, ...) {
  args <- .recycle_cont_vectors(x, n, .check_i_cont(i))
  x <- args[[1L]]
  n <- args[[2L]]
  i <- args[[3L]]

  .check_x_cont(x)
  .check_n_cont(n)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    nn <- n[j]
    delta <- log(1 + i[j])

    upper <- .get_upper_varying_cont(xx, n = nn, model = model, ...)

    if (upper <= 0) {
      return(0)
    }

    f <- function(t) {
      val <- t * exp(-delta * t) *
        tpx(t, x = xx, model = model, ...) *
        hazard0(xx + t, model = model, ...)
      val[!is.finite(val)] <- 0
      val
    }

    .integrate_varying_cont(f, 0, upper)
  })
}

#' Fully continuous decreasing n-year term insurance
#'
#' Computes
#' \deqn{(\bar{D}\bar{A})_{x:\overline{n}|}^{1}
#' = \int_0^n (n-t) v^t {}_tp_x \mu_{x+t}\,dt.}
#'
#' @inheritParams IbarAbarxn1
#'
#' @return Numeric vector.
#' @export
DbarAbarxn1 <- function(x, n, i, model, ...) {
  args <- .recycle_cont_vectors(x, n, .check_i_cont(i))
  x <- args[[1L]]
  n <- args[[2L]]
  i <- args[[3L]]

  .check_x_cont(x)
  .check_n_cont(n)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    nn <- n[j]
    delta <- log(1 + i[j])

    upper <- .get_upper_varying_cont(xx, n = nn, model = model, ...)

    if (upper <= 0) {
      return(0)
    }

    f <- function(t) {
      val <- (nn - t) * exp(-delta * t) *
        tpx(t, x = xx, model = model, ...) *
        hazard0(xx + t, model = model, ...)
      val[!is.finite(val)] <- 0
      val
    }

    .integrate_varying_cont(f, 0, upper)
  })
}

#' Piecewise-continuous decreasing n-year term insurance
#'
#' Computes
#' \deqn{(D\bar{A})_{x:\overline{n}|}^{1}
#' = \int_0^n \lfloor n+1-t \rfloor v^t {}_tp_x \mu_{x+t}\,dt.}
#'
#' @inheritParams IbarAbarxn1
#'
#' @return Numeric vector.
#' @export
DAbarxn1 <- function(x, n, i, model, ...) {
  args <- .recycle_cont_vectors(x, n, .check_i_cont(i))
  x <- args[[1L]]
  n <- args[[2L]]
  i <- args[[3L]]

  .check_x_cont(x)
  .check_n_cont(n)

  sapply(seq_along(x), function(j) {
    xx <- x[j]
    nn <- n[j]
    delta <- log(1 + i[j])

    upper <- .get_upper_varying_cont(xx, n = nn, model = model, ...)

    if (upper <= 0) {
      return(0)
    }

    f <- function(t) {
      val <- floor(nn + 1 - t) * exp(-delta * t) *
        tpx(t, x = xx, model = model, ...) *
        hazard0(xx + t, model = model, ...)
      val[!is.finite(val)] <- 0
      val
    }

    .integrate_varying_cont(f, 0, upper)
  })
}
