# ==========================================================================
# Closed-form k-out-of-n of heterogeneous Weibull components
# ==========================================================================
#
# Same subset-enumeration formula as exp_kofn, with Weibull per-component
# survivals S_j(t) = exp(-(t / scale_j)^shape_j).
# ==========================================================================


#' k-out-of-n system of independent Weibull components (closed form)
#'
#' Constructs a `dist_structure` for a k-out-of-m system whose components
#' are independent (possibly heterogeneous) Weibulls. Closed-form `surv`,
#' `cdf`, `sampler`, `density`, and `hazard` via subset enumeration,
#' the critical-state density formula, and component order statistics.
#'
#' @param k Minimum functioning components for system operation.
#' @param shapes Positive numeric vector of length `m`.
#' @param scales Positive numeric vector of length `m`.
#' @return
#' `wei_kofn()` returns an object of class
#'   `c("wei_kofn", "kofn_dist", "coherent_dist", "dist_structure",
#'   "univariate_dist", "continuous_dist", "dist")`.
#'
#' The associated S3 methods return:
#' - `surv()`, `density()`, `hazard()`: a closure `function(t, ...)`.
#' - `cdf()` is derived via the `dist_structure` default and returns
#'   a closure `function(t, ...)` equal to `1 - surv(x)(t)`.
#' - `sampler()`: a closure `function(n, ...)` returning `n` random
#'   variates from the system lifetime distribution.
#' @examples
#' sys <- wei_kofn(k = 2, shapes = c(1, 2, 3), scales = c(1, 2, 3))
#' algebraic.dist::surv(sys)(1)
#' @export
wei_kofn <- function(k, shapes, scales) {
  stopifnot(length(shapes) == length(scales),
            all(shapes > 0), all(scales > 0))
  m <- length(shapes)
  stopifnot(k >= 1L, k <= m)
  components <- lapply(seq_len(m), function(j) {
    algebraic.dist::weibull_dist(shape = shapes[j], scale = scales[j])
  })
  obj <- kofn_dist(k, components)
  obj$shapes <- as.numeric(shapes)
  obj$scales <- as.numeric(scales)
  class(obj) <- c("wei_kofn", class(obj))
  obj
}


#' @rdname wei_kofn
#' @param x A `wei_kofn` object.
#' @param ... Ignored.
#' @export
surv.wei_kofn <- function(x, ...) {
  shapes <- x$shapes
  scales <- x$scales
  k <- x$k
  function(t, ...) {
    vapply(t, function(ti) {
      kofn_surv_probability(exp(-(ti / scales)^shapes), k)
    }, numeric(1L))
  }
}


#' @rdname wei_kofn
#' @export
sampler.wei_kofn <- function(x, ...) {
  order_idx <- length(x$shapes) - x$k + 1L
  samplers <- make_component_samplers(stats::rweibull,
                                      shape = x$shapes, scale = x$scales)
  function(n, ...) {
    apply(sample_component_matrix(samplers, n), 1L,
          function(row) sort(row)[order_idx])
  }
}


#' @rdname wei_kofn
#' @method hazard wei_kofn
#' @importFrom algebraic.dist hazard
#' @export
hazard.wei_kofn <- function(x, ...) {
  # h_sys(t) = f_sys(t) / S_sys(t); both factors have closed forms here.
  f_fn <- density.wei_kofn(x)
  S_fn <- surv.wei_kofn(x)
  function(t, ...) f_fn(t) / S_fn(t)
}


#' @rdname wei_kofn
#' @importFrom stats density dweibull
#' @export
density.wei_kofn <- function(x, ...) {
  shapes <- x$shapes
  scales <- x$scales
  k <- x$k
  m <- length(shapes)
  function(t, log = FALSE, ...) {
    vals <- vapply(t, function(ti) {
      surv <- exp(-(ti / scales)^shapes)
      dens <- vapply(seq_len(m), function(j) {
        stats::dweibull(ti, shape = shapes[j], scale = scales[j])
      }, numeric(1L))
      kofn_density_value(dens, surv, k)
    }, numeric(1L))
    if (isTRUE(log)) log(vals) else vals
  }
}
