#' Solve the mobility-based SEAIR system
#'
#' Discretises the mobility-based SEAIR system on the user-supplied
#' mobility grid (using the trapezoidal rule for the non-local
#' integral \eqn{\mathcal{K}}) and integrates the resulting coupled
#' ODE system over \code{times} with \pkg{deSolve}.
#'
#' @param init A \code{"seair_init"} object from \code{\link{seair_init}}.
#' @param params A \code{"seair_params"} object from
#'   \code{\link{seair_params}}.
#' @param times Numeric vector of times at which the solution is
#'   reported (must be increasing and include \code{0}).
#' @param method Integration method passed to \code{\link[deSolve]{ode}}
#'   (default \code{"lsoda"}).
#' @param ... Additional arguments passed to \code{\link[deSolve]{ode}}.
#'
#' @return A list of class \code{"seair_solution"} with entries
#'   \code{times}, \code{m_grid}, \code{S}, \code{E}, \code{A},
#'   \code{I}, \code{R} (each a matrix of dimension
#'   \code{length(times)} by \code{length(m_grid)}),
#'   together with the input \code{params} and \code{init}.
#'
#' @examples
#' m <- seq(0, 1, length.out = 31)
#' f <- dbeta(m, 2, 2)
#' init <- seair_init(m, f, I_seed = 1e-4)
#' 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)
#' sol <- seair_solve(init, pars, times = seq(0, 30, by = 2))
#'
#' @export
seair_solve <- function(init, params, times,
                        method = "lsoda", ...) {
  stopifnot(inherits(init, "seair_init"),
            inherits(params, "seair_params"),
            is.numeric(times), length(times) >= 2,
            all(is.finite(times)), all(diff(times) > 0), times[1] == 0)

  m  <- init$m_grid
  n  <- length(m)
  w  <- .trap_weights(m)   # quadrature weights for the trapezoidal rule

  beta    <- params$beta
  sigma   <- params$sigma
  kappa   <- params$kappa
  gamma_A <- params$gamma_A
  gamma_I <- params$gamma_I
  alpha   <- params$alpha
  delta   <- params$delta

  ## State vector layout: [S(1..n), E(1..n), A(1..n), I(1..n), R(1..n)]
  y0 <- c(init$S, init$E, init$A, init$I, init$R)

  rhs <- function(t, y, parms) {
    S <- y[1:n]
    E <- y[(n + 1):(2 * n)]
    A <- y[(2 * n + 1):(3 * n)]
    I <- y[(3 * n + 1):(4 * n)]
    R <- y[(4 * n + 1):(5 * n)]

    K <- sum(w * m * (alpha * A + delta * I))

    dS <- -beta * m * S * K
    dE <-  beta * m * S * K - sigma * E
    dA <-  sigma * E - (kappa + gamma_A) * A
    dI <-  kappa * A - gamma_I * I
    dR <-  gamma_A * A + gamma_I * I

    list(c(dS, dE, dA, dI, dR))
  }

  out <- deSolve::ode(y = y0, times = times, func = rhs,
                      parms = NULL, method = method, ...)

  mat <- unclass(out)
  tt  <- mat[, 1]
  S   <- mat[, 1 + (1:n),           drop = FALSE]
  E   <- mat[, 1 + ( (n + 1):(2 * n) ),     drop = FALSE]
  A   <- mat[, 1 + ( (2 * n + 1):(3 * n) ), drop = FALSE]
  I   <- mat[, 1 + ( (3 * n + 1):(4 * n) ), drop = FALSE]
  R   <- mat[, 1 + ( (4 * n + 1):(5 * n) ), drop = FALSE]

  sol <- list(times = tt, m_grid = m,
              S = S, E = E, A = A, I = I, R = R,
              params = params, init = init)
  class(sol) <- "seair_solution"
  sol
}

#' Aggregate a SEAIR solution over mobility
#'
#' Computes the aggregate time series
#' \eqn{\int_0^1 S(m,t)\,dm}, \eqn{\int_0^1 E(m,t)\,dm}, etc., using
#' the trapezoidal rule on the solver's mobility grid.
#'
#' @param sol A \code{"seair_solution"} object returned by
#'   \code{\link{seair_solve}}.
#'
#' @return A data frame with columns \code{time}, \code{S}, \code{E},
#'   \code{A}, \code{I}, \code{R} (aggregate ratios).
#'
#' @examples
#' m <- seq(0, 1, length.out = 31)
#' f <- dbeta(m, 2, 2)
#' init <- seair_init(m, f, I_seed = 1e-4)
#' 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)
#' sol <- seair_solve(init, pars, times = seq(0, 30, by = 2))
#' agg <- seair_aggregate(sol)
#' head(agg)
#'
#' @export
seair_aggregate <- function(sol) {
  stopifnot(inherits(sol, "seair_solution"))
  w <- .trap_weights(sol$m_grid)
  data.frame(
    time = sol$times,
    S = as.numeric(sol$S %*% w),
    E = as.numeric(sol$E %*% w),
    A = as.numeric(sol$A %*% w),
    I = as.numeric(sol$I %*% w),
    R = as.numeric(sol$R %*% w)
  )
}

## Trapezoidal quadrature weights on a (possibly non-uniform) grid.
.trap_weights <- function(x) {
  n <- length(x)
  w <- numeric(n)
  w[1]     <- (x[2] - x[1]) / 2
  w[n]     <- (x[n] - x[n - 1]) / 2
  if (n > 2) {
    w[2:(n - 1)] <- (x[3:n] - x[1:(n - 2)]) / 2
  }
  w
}
