#' Beta-mixture density on [0, 1]
#'
#' Evaluates a weighted mixture of Beta densities
#' \deqn{f(m) \;=\; \sum_{k=1}^{K} w_k\,\mathrm{Beta}(m; a_k, b_k),}
#' renormalised on \eqn{[0,1]}.  Weights are coerced to be nonnegative
#' and to sum to one.
#'
#' @param m Numeric vector of evaluation points in \eqn{[0,1]}.
#' @param weights Numeric vector of mixture weights (length \eqn{K}).
#' @param shape1 Numeric vector of Beta \eqn{a_k} parameters.
#' @param shape2 Numeric vector of Beta \eqn{b_k} parameters.
#'
#' @return Numeric vector of densities at \code{m}.
#'
#' @examples
#' m <- seq(0, 1, length.out = 101)
#' beta_mixture_density(m,
#'                      weights = c(0.5, 0.3, 0.2),
#'                      shape1  = c(2, 5, 9),
#'                      shape2  = c(8, 3, 2))
#'
#' @export
beta_mixture_density <- function(m, weights, shape1, shape2) {
  K <- length(weights)
  stopifnot(is.numeric(m),
            all(is.finite(m)), all(m >= 0), all(m <= 1),
            is.numeric(weights), K >= 1,
            is.numeric(shape1), is.numeric(shape2),
            length(shape1) == K, length(shape2) == K,
            all(is.finite(weights)),
            all(is.finite(shape1)), all(is.finite(shape2)),
            all(shape1 > 0), all(shape2 > 0))
  w <- pmax(weights, 0)
  if (sum(w) <= 0) stop("All mixture weights are zero.")
  w <- w / sum(w)
  out <- numeric(length(m))
  for (k in seq_len(K)) {
    out <- out + w[k] * stats::dbeta(m, shape1[k], shape2[k])
  }
  out
}

#' Fit a mobility distribution to an observed symptomatic time series
#'
#' Fits a Beta-mixture mobility distribution \eqn{f} by minimising the
#' sum of squared errors between an observed aggregate symptomatic
#' time series \eqn{\langle I \rangle(t)} and the simulated output of
#' the mobility-based SEAIR model at the same times.  All transmission
#' and stage parameters are held fixed.
#'
#' The unconstrained optimisation parameter is, for each component
#' \eqn{k = 1, \ldots, K}, a triple
#' \eqn{(\log a_k, \log b_k, \log \tilde{w}_k)} with mixture weights
#' recovered as \eqn{w_k = \tilde{w}_k / \sum_j \tilde{w}_j}.  The
#' first weight's log is fixed to 0 to remove the scale ambiguity,
#' so the parameter vector has length \eqn{3 K - 1}.
#'
#' @param times Numeric vector of observation times (starting at 0).
#' @param I_obs Observed aggregate symptomatic fraction at each time.
#' @param params A \code{"seair_params"} object holding fixed
#'   transmission and stage parameters.
#' @param K Number of Beta-mixture components (default 3).
#' @param m_grid Mobility grid for the solver.
#' @param I_seed Initial symptomatic fraction.
#' @param start Optional user-supplied starting parameter vector of
#'   length \eqn{3K - 1}.
#' @param n_restarts Number of random restarts (default 5).
#' @param control List of \code{\link[stats]{optim}} control options.
#' @param verbose Logical; if \code{TRUE}, print progress.
#'
#' @return A list with \code{weights}, \code{shape1}, \code{shape2}
#'   (the fitted Beta-mixture parameters), \code{f_hat} (density on
#'   \code{m_grid}), \code{m_grid}, \code{loss} (final SSE),
#'   \code{I_fit} (simulated aggregate symptomatic time series at
#'   \code{times}), and \code{optim_result}.
#'
#' @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)
#' m <- seq(0, 1, length.out = 21)
#' f_true <- dbeta(m, 3, 3)
#' init <- seair_init(m, f_true, I_seed = 1e-4)
#' times <- seq(0, 20, by = 5)
#' sol <- seair_solve(init, pars, times)
#' I_obs <- seair_aggregate(sol)$I
#'
#' fit <- fit_mobility(times, I_obs, pars,
#'                     K = 1, m_grid = m,
#'                     start = log(c(3, 3)),
#'                     n_restarts = 0,
#'                     control = list(maxit = 20))
#' fit$loss
#'
#' @export
fit_mobility <- function(times, I_obs, params,
                         K = 3,
                         m_grid = seq(0, 1, length.out = 51),
                         I_seed = 1e-4,
                         start = NULL,
                         n_restarts = 5,
                         control = list(maxit = 500),
                         verbose = FALSE) {
  stopifnot(is.numeric(times), is.numeric(I_obs),
            length(times) == length(I_obs), length(times) >= 2,
            all(is.finite(times)), all(diff(times) > 0), times[1] == 0,
            all(is.finite(I_obs)), all(I_obs >= 0),
            inherits(params, "seair_params"),
            is.numeric(K), length(K) == 1, K >= 1, K == as.integer(K),
            is.numeric(m_grid), length(m_grid) >= 3,
            all(is.finite(m_grid)), all(m_grid >= 0), all(m_grid <= 1),
            all(diff(m_grid) > 0),
            is.numeric(I_seed), length(I_seed) == 1,
            is.finite(I_seed), I_seed >= 0, I_seed < 1,
            is.numeric(n_restarts), length(n_restarts) == 1,
            n_restarts >= 0, n_restarts == as.integer(n_restarts),
            is.list(control),
            is.logical(verbose), length(verbose) == 1, !is.na(verbose))
  K <- as.integer(K)
  n_restarts <- as.integer(n_restarts)

  n_par <- 3 * K - 1
  if (is.null(start) && n_restarts == 0) {
    stop("At least one of 'start' or a positive 'n_restarts' is required.")
  }

  unpack <- function(theta) {
    log_a <- theta[1:K]
    log_b <- theta[(K + 1):(2 * K)]
    if (K == 1) {
      w_tilde <- 1
    } else {
      w_tilde <- c(1, exp(theta[(2 * K + 1):(3 * K - 1)]))
    }
    list(shape1  = exp(log_a),
         shape2  = exp(log_b),
         weights = w_tilde / sum(w_tilde))
  }

  obj <- function(theta) {
    pars_mix <- tryCatch(unpack(theta), error = function(e) NULL)
    if (is.null(pars_mix)) return(.Machine$double.xmax)
    f_vals <- tryCatch(
      beta_mixture_density(m_grid,
                           weights = pars_mix$weights,
                           shape1  = pars_mix$shape1,
                           shape2  = pars_mix$shape2),
      error = function(e) NULL)
    if (is.null(f_vals) || any(!is.finite(f_vals))) {
      return(.Machine$double.xmax)
    }
    init <- seair_init(m_grid, f_vals, I_seed = I_seed)
    sol  <- tryCatch(seair_solve(init, params, times = times),
                     error = function(e) NULL)
    if (is.null(sol)) return(.Machine$double.xmax)
    I_sim <- seair_aggregate(sol)$I
    sum((I_sim - I_obs)^2)
  }

  random_start <- function() {
    log_a <- log(stats::runif(K, 1, 10))
    log_b <- log(stats::runif(K, 1, 10))
    if (K == 1) return(c(log_a, log_b))
    c(log_a, log_b, stats::rnorm(K - 1, sd = 0.5))
  }

  best <- list(par = NULL, value = Inf, convergence = NA_integer_)

  ## User-supplied start, then random restarts.
  starts <- list()
  if (!is.null(start)) {
    if (!is.numeric(start) || any(!is.finite(start)))
      stop("'start' must be a finite numeric vector.")
    if (length(start) != n_par)
      stop(sprintf("'start' must have length %d (got %d).",
                   n_par, length(start)))
    starts[[1]] <- start
  }
  for (r in seq_len(n_restarts)) {
    starts[[length(starts) + 1]] <- random_start()
  }

  for (i in seq_along(starts)) {
    if (verbose) message("restart ", i, " of ", length(starts))
    res <- tryCatch(
      stats::optim(starts[[i]], obj, method = "Nelder-Mead",
                   control = control),
      error = function(e) NULL)
    if (!is.null(res) && is.finite(res$value) && res$value < best$value) {
      best <- res
    }
  }

  if (!is.finite(best$value)) {
    stop("Optimisation failed at all starts; try increasing 'n_restarts'.")
  }

  pars_mix <- unpack(best$par)
  f_hat    <- beta_mixture_density(m_grid,
                                   weights = pars_mix$weights,
                                   shape1  = pars_mix$shape1,
                                   shape2  = pars_mix$shape2)
  init <- seair_init(m_grid, f_hat, I_seed = I_seed)
  I_fit <- seair_aggregate(seair_solve(init, params, times))$I

  list(weights      = pars_mix$weights,
       shape1       = pars_mix$shape1,
       shape2       = pars_mix$shape2,
       f_hat        = f_hat,
       m_grid       = m_grid,
       loss         = best$value,
       I_fit        = I_fit,
       optim_result = best)
}
