#' Effective infectiousness duration of the SEAIR chain
#'
#' Computes the scalar
#' \deqn{c \;=\; \frac{\alpha}{\kappa + \gamma_A}
#'       \;+\; \frac{\delta\,\kappa}{(\kappa + \gamma_A)\,\gamma_I},}
#' the integrated contribution of a unit of newly infected mass to
#' the force of infection, collapsing the \eqn{E \to A \to I} stage
#' chain into a single number.  The latency rate \eqn{\sigma} does
#' not appear in \eqn{c}.
#'
#' @param params A \code{"seair_params"} object.
#'
#' @return A positive scalar.
#'
#' @examples
#' pars <- seair_params(beta = 1.5, sigma = 0.3, kappa = 0.2,
#'                      gamma_A = 0.1, gamma_I = 0.13,
#'                      alpha = 0.5, delta = 0.3)
#' c_effective(pars)
#'
#' @export
c_effective <- function(params) {
  stopifnot(inherits(params, "seair_params"))
  with(params,
       alpha / (kappa + gamma_A)
       + delta * kappa / ((kappa + gamma_A) * gamma_I))
}

#' Basic reproduction number of the mobility-based SEAIR model
#'
#' For a mobility distribution \eqn{f} on \eqn{(0,1)}, the basic
#' reproduction number at the disease-free equilibrium is
#' \deqn{\mathcal{R}_0 \;=\; \beta\,c\,\langle m^2 f \rangle,}
#' where \eqn{c} is as in \code{\link{c_effective}} and
#' \eqn{\langle m^2 f \rangle = \int_0^1 m^2 f(m)\,dm}.
#'
#' @param params A \code{"seair_params"} object.
#' @param f A function returning the density \eqn{f(m)} on \eqn{[0,1]},
#'   or a numeric vector of \eqn{f} values evaluated on \code{m_grid}.
#' @param m_grid Mobility grid (used only if \code{f} is numeric).
#'
#' @return The basic reproduction number \eqn{\mathcal{R}_0}.
#'
#' @examples
#' pars <- seair_params(beta = 1.5, sigma = 0.3, kappa = 0.2,
#'                      gamma_A = 0.1, gamma_I = 0.13,
#'                      alpha = 0.5, delta = 0.3)
#'
#' ## f given as a function (normalised automatically by the caller).
#' f <- function(m) 6 * m * (1 - m)    # Beta(2, 2) density on [0, 1]
#' R0_seair(pars, f)
#'
#' ## f given as values on a grid.
#' m <- seq(0, 1, length.out = 101)
#' R0_seair(pars, f(m), m_grid = m)
#'
#' @export
R0_seair <- function(params, f, m_grid = NULL) {
  stopifnot(inherits(params, "seair_params"))
  ce <- c_effective(params)
  if (is.function(f)) {
    Q2 <- stats::integrate(function(m) m^2 * f(m), 0, 1)$value
    total <- stats::integrate(f, 0, 1)$value
    if (total <= 0) stop("f integrates to zero.")
    Q2 <- Q2 / total
  } else if (is.numeric(f)) {
    if (is.null(m_grid) || length(m_grid) != length(f))
      stop("If 'f' is numeric, 'm_grid' of matching length is required.")
    if (length(m_grid) < 3 || any(!is.finite(m_grid)) ||
        any(m_grid < 0) || any(m_grid > 1) || any(diff(m_grid) <= 0) ||
        any(!is.finite(f)) || any(f < 0)) {
      stop("'m_grid' must be an increasing grid in [0, 1] and 'f' must be finite and nonnegative.")
    }
    w  <- .trap_weights(m_grid)
    total <- sum(w * f)
    if (total <= 0) stop("f integrates to zero on the grid.")
    Q2 <- sum(w * m_grid^2 * f) / total
  } else stop("'f' must be a function or a numeric vector.")
  params$beta * ce * Q2
}

#' Final epidemic size under the vanishing-seed regime
#'
#' Solves the scalar fixed-point equation
#' \deqn{M_R^\infty \;=\; \int_0^1 m f(m)\,(1 - e^{-\beta c M_R^\infty m})\,dm}
#' for the first mobility moment of the recovered compartment at the
#' end of the epidemic (with \eqn{I_0 \equiv 0}, \eqn{R_0 \equiv 0},
#' \eqn{S_0 = f}), and returns the total final epidemic size
#' \deqn{R_\infty \;=\; 1 - \int_0^1 f(m)\,e^{-\beta c M_R^\infty m}\,dm.}
#' When the basic reproduction number is less than or equal to one,
#' the function returns \eqn{R_\infty = 0}.
#'
#' @param params A \code{"seair_params"} object.
#' @param f A mobility density, either a function on \eqn{[0,1]} or a
#'   numeric vector of values on \code{m_grid}.
#' @param m_grid Optional mobility grid used when \code{f} is numeric,
#'   or to evaluate the integrals (default: uniform grid of 201 points).
#' @param tol Tolerance for \code{\link{uniroot}}.
#'
#' @return A list with \code{R_inf} (total final size),
#'   \code{MR_inf} (first moment), \code{R0} (basic reproduction
#'   number), and \code{converged} (logical).
#'
#' @examples
#' pars <- seair_params(beta = 1.5, sigma = 0.3, kappa = 0.2,
#'                      gamma_A = 0.1, gamma_I = 0.13,
#'                      alpha = 0.5, delta = 0.3)
#' f <- function(m) 6 * m * (1 - m)
#' final_size(pars, f)
#'
#' @export
final_size <- function(params, f, m_grid = NULL, tol = 1e-10) {
  stopifnot(inherits(params, "seair_params"))
  if (is.null(m_grid)) m_grid <- seq(0, 1, length.out = 201)

  fv <- if (is.function(f)) f(m_grid) else f
  if (length(fv) != length(m_grid))
    stop("Length of f does not match m_grid.")
  if (length(m_grid) < 3 || any(!is.finite(m_grid)) ||
      any(m_grid < 0) || any(m_grid > 1) || any(diff(m_grid) <= 0) ||
      any(!is.finite(fv)) || any(fv < 0)) {
    stop("'m_grid' must be an increasing grid in [0, 1] and 'f' must be finite and nonnegative.")
  }
  w     <- .trap_weights(m_grid)
  total <- sum(w * fv)
  if (total <= 0) stop("f integrates to zero on m_grid.")
  fv    <- fv / total

  R0 <- R0_seair(params, fv, m_grid)
  if (R0 <= 1) {
    return(list(R_inf = 0, MR_inf = 0, R0 = R0, converged = TRUE))
  }

  beta_c <- params$beta * c_effective(params)
  Mf1    <- sum(w * m_grid * fv)   # upper bound on MR_inf

  Phi <- function(x) {
    sum(w * m_grid * fv * (1 - exp(-beta_c * x * m_grid))) - x
  }

  ## Phi(0) = 0 and Phi'(0) = R0 - 1 > 0, Phi(Mf1) < 0, so root in (0, Mf1].
  lo <- 1e-12 * Mf1
  hi <- Mf1
  if (Phi(lo) <= 0 || Phi(hi) >= 0) {
    return(list(R_inf = NA_real_, MR_inf = NA_real_,
                R0 = R0, converged = FALSE))
  }
  rt <- stats::uniroot(Phi, lower = lo, upper = hi, tol = tol)
  MR_inf <- rt$root
  R_inf  <- 1 - sum(w * fv * exp(-beta_c * MR_inf * m_grid))

  list(R_inf = R_inf, MR_inf = MR_inf, R0 = R0, converged = TRUE)
}

#' Final epidemic size under general proportional initial conditions
#'
#' Solves the scalar equation of Theorem 4.3 (final-size relation),
#' \deqn{M_R^\infty + \int_0^1 m S_0(m)\,
#'     \exp\!\bigl(-\beta m [ c\,(M_R^\infty - M_R^0)
#'                         + (\delta/\gamma_I - c)\,M_I^0 ]\bigr) dm
#'     \;=\; \int_0^1 m f(m)\,dm,}
#' allowing a nonzero initial symptomatic fraction
#' \eqn{I_0 = I_{\mathrm{seed}} f}, with \eqn{E_0 = A_0 = R_0 = 0}.
#'
#' @inheritParams final_size
#' @param I_seed Initial symptomatic fraction.
#'
#' @return A list with \code{R_inf}, \code{MR_inf}, \code{R0}, and
#'   \code{converged}.
#'
#' @examples
#' pars <- seair_params(beta = 1.5, sigma = 0.3, kappa = 0.2,
#'                      gamma_A = 0.1, gamma_I = 0.13,
#'                      alpha = 0.5, delta = 0.3)
#' f <- function(m) 6 * m * (1 - m)
#' final_size_general(pars, f, I_seed = 0.001)
#'
#' @export
final_size_general <- function(params, f, m_grid = NULL,
                               I_seed = 1e-4, tol = 1e-10) {
  stopifnot(inherits(params, "seair_params"),
            is.numeric(I_seed), length(I_seed) == 1,
            I_seed >= 0, I_seed < 1)

  if (is.null(m_grid)) m_grid <- seq(0, 1, length.out = 201)
  fv <- if (is.function(f)) f(m_grid) else f
  if (length(fv) != length(m_grid))
    stop("Length of f does not match m_grid.")
  if (length(m_grid) < 3 || any(!is.finite(m_grid)) ||
      any(m_grid < 0) || any(m_grid > 1) || any(diff(m_grid) <= 0) ||
      any(!is.finite(fv)) || any(fv < 0)) {
    stop("'m_grid' must be an increasing grid in [0, 1] and 'f' must be finite and nonnegative.")
  }
  w     <- .trap_weights(m_grid)
  total <- sum(w * fv)
  if (total <= 0) stop("f integrates to zero on m_grid.")
  fv    <- fv / total

  S0 <- (1 - I_seed) * fv
  I0 <- I_seed * fv
  MR_0 <- 0
  MI_0 <- sum(w * m_grid * I0)
  Mf_1 <- sum(w * m_grid * fv)

  ce    <- c_effective(params)
  nu    <- params$delta / params$gamma_I - ce
  beta  <- params$beta
  R0    <- R0_seair(params, fv, m_grid)

  ## Vanishing-seed: I_0 = 0, R_0 = 0, S_0 = f  => reduce to final_size().
  if (I_seed == 0) return(final_size(params, fv, m_grid, tol = tol))

  Fx <- function(x) {
    expo <- -beta * m_grid * (ce * (x - MR_0) + nu * MI_0)
    x - Mf_1 + sum(w * m_grid * S0 * exp(expo))
  }

  lo <- MR_0 + MI_0          # x_ell
  hi <- Mf_1                 # x_r
  if (hi <= lo) {
    return(list(R_inf = 0, MR_inf = lo, R0 = R0, converged = TRUE))
  }
  Fl <- Fx(lo); Fh <- Fx(hi)
  if (!(is.finite(Fl) && is.finite(Fh)) || Fl >= 0 || Fh <= 0) {
    return(list(R_inf = NA_real_, MR_inf = NA_real_,
                R0 = R0, converged = FALSE))
  }
  rt <- stats::uniroot(Fx, lower = lo, upper = hi, tol = tol)
  MR_inf <- rt$root

  expo   <- -beta * m_grid * (ce * (MR_inf - MR_0) + nu * MI_0)
  R_inf  <- 1 - sum(w * S0 * exp(expo))

  list(R_inf = R_inf, MR_inf = MR_inf, R0 = R0, converged = TRUE)
}
