# ============================================================================
# S3 methods for class "npf_seir"
# ============================================================================

#' Print method for npf_seir
#' @param x An \code{npf_seir} object.
#' @param ... Ignored.
#' @return Returns \code{x} invisibly. Called primarily for its side effect
#'   of printing a concise summary of the fitted model to the console,
#'   including the number of observations and particles, and the posterior
#'   means of the Ornstein-Uhlenbeck hyperparameters at the final time step.
#' @export
print.npf_seir <- function(x, ...) {
  n_obs <- length(x$observations)
  tm <- x$theta_mean[max(which(!is.na(x$theta_mean[, 1]))), ]
  cat("Nested Particle Filter - stochastic SEIR with latent OU transmission\n")
  cat(sprintf("  n = %d observations   K = %d outer   M = %d inner particles\n",
              n_obs, x$K, x$M))
  cat("  Posterior means of OU hyperparameters (at final step):\n")
  cat(sprintf("    kappa     = %.4f  (timescale ~%.1f days)\n",
              tm[1], 1/tm[1]))
  cat(sprintf("    sigma_Psi = %.4f\n", tm[2]))
  cat(sprintf("    mu_Psi    = %.4f   exp(mu_Psi) = %.4f\n",
              tm[3], exp(tm[3])))
  invisible(x)
}


#' Summary method for npf_seir
#'
#' Summarises posterior parameter estimates and the range of \eqn{R_t}.
#'
#' @param object An \code{npf_seir} object.
#' @param burn   Integer. Number of initial steps to exclude from summaries.
#'   Default 0.
#' @param ... Ignored.
#' @return A named list of class \code{"npf_seir_summary"} returned invisibly,
#'   with components:
#'   \describe{
#'     \item{\code{theta}}{Named numeric vector of length 3 giving the weighted
#'       posterior mean of the OU hyperparameters
#'       \eqn{(\kappa, \sigma_\Psi, \mu_\Psi)} at the final time step.}
#'     \item{\code{ou}}{Named list with elements \code{alpha} and \code{tau},
#'       the exact discrete-time OU transition parameters derived from
#'       \code{theta}; see \code{\link{ou_params}}.}
#'   }
#'   The function is also called for its side effect of printing a formatted
#'   summary to the console.
#' @importFrom stats rnorm
#' @export
summary.npf_seir <- function(object, burn = 0L, ...) {
  if (!is.numeric(burn) || length(burn) != 1L || is.na(burn) || burn < 0 || burn != as.integer(burn)) {
    stop("burn must be a non-negative integer.", call. = FALSE)
  }

  x    <- object
  burn <- as.integer(burn)
  if (burn >= length(x$Rt_mean)) {
    stop("burn must be smaller than the number of time points.", call. = FALSE)
  }

  keep <- seq(burn + 1L, length(x$Rt_mean))

  # Posterior means weighted by final outer weights
  tm <- x$theta_mean[max(which(!is.na(x$theta_mean[, 1]))), ]
  ou <- ou_params(tm[1], tm[2])

  cat("=== npf_seir summary ===\n\n")
  cat("Call:\n"); print(x$call); cat("\n")

  cat("Posterior parameter estimates (weighted mean at final step):\n")
  cat(sprintf("  kappa     = %.4f  [alpha = %.4f]\n", tm[1], ou$alpha))
  cat(sprintf("  sigma_Psi = %.4f  [tau   = %.4f]\n", tm[2], ou$tau))
  cat(sprintf("  mu_Psi    = %.4f  [exp(mu_Psi) = %.4f]\n",
              tm[3], exp(tm[3])))
  cat(sprintf("  Implied reproduction scale: exp(mu_Psi)/gamma = %.3f\n",
              exp(tm[3]) / x$fixed$gamma))

  rt  <- x$Rt_mean[keep]
  rt  <- rt[!is.na(rt)]
  cat(sprintf("\nR_t range (steps %d-%d): min=%.3f  median=%.3f  max=%.3f\n",
              burn + 1L, length(x$Rt_mean),
              min(rt), median(rt), max(rt)))
  cat(sprintf("Days with R_t > 1: %d / %d\n",
              sum(rt > 1), length(rt)))

  invisible(structure(list(theta = tm, ou = ou), class = "npf_seir_summary"))
}


#' Plot method for npf_seir
#'
#' Produces posterior summary plots for \eqn{R_t}, \eqn{\beta_t}, or both.
#'
#' @param x      An \code{npf_seir} object.
#' @param burn   Integer. Number of initial steps to exclude from the plot.
#' @param which  Character. \code{"Rt"} (default), \code{"beta"}, or
#'   \code{"both"} to show both trajectories in a two-panel display.
#' @param dates  Optional Date vector with one entry per observation to use as
#'   the x-axis.
#' @param ...    Additional arguments passed to \code{plot()}.
#' @return Returns \code{x} invisibly. Called primarily for its side effect
#'   of producing a base-graphics plot. When \code{which = "Rt"} or
#'   \code{"beta"}, a single panel is drawn showing the posterior mean
#'   trajectory and 95\% credible band. When \code{which = "both"}, a
#'   two-panel figure is produced showing \eqn{R_t} (top) and \eqn{\beta_t}
#'   (bottom).
#' @importFrom graphics plot
#' @export
plot.npf_seir <- function(x, burn = 0L, which = "Rt", dates = NULL, ...) {
  if (!is.numeric(burn) || length(burn) != 1L || is.na(burn) || burn < 0 || burn != as.integer(burn)) {
    stop("burn must be a non-negative integer.", call. = FALSE)
  }

  n_obs <- length(x$observations)
  burn <- as.integer(burn)
  if (burn >= n_obs) {
    stop("burn must be smaller than the number of time points.", call. = FALSE)
  }

  which <- match.arg(which, choices = c("Rt", "beta", "both"))
  keep <- seq(burn + 1L, n_obs)
  xax  <- if (!is.null(dates)) dates[keep] else keep

  # Extract 'main' from ... so it doesn't conflict with hardcoded main
  dots  <- list(...)
  main1 <- if (!is.null(dots$main)) dots$main else "Effective reproduction number"
  dots$main <- NULL

  old_par <- par(no.readonly = TRUE)
  on.exit(par(old_par))

  if (which == "both") {
    par(mfrow = c(2, 1), mar = c(3, 4, 2, 1), oma = c(1, 0, 0, 0))
  } else {
    par(mfrow = c(1, 1), mar = c(4, 4, 2, 1))
  }

  draw_band <- function(mean, lo, hi, ylab, main, col) {
    ylim <- range(c(0, hi), na.rm = TRUE) * c(1, 1.05)
    do.call(plot, c(list(xax, mean, type = "n", ylim = ylim,
      ylab = ylab, xlab = if (which == "both") "" else "Step",
      main = main, las = 1), dots))
    polygon(c(xax, rev(xax)), c(hi, rev(lo)),
      col = adjustcolor(col, 0.25), border = NA)
    lines(xax, mean, col = col, lwd = 2)
    legend("topright", legend = c("Posterior mean", "95% CI"),
      lwd = c(2, NA), fill = c(NA, adjustcolor(col, 0.25)),
      border = c(NA, NA), col = c(col, NA), bty = "n")
  }

  if (which %in% c("Rt", "both")) {
    rt_m  <- x$Rt_mean[keep]
    rt_lo <- x$Rt_lo[keep]
    rt_hi <- x$Rt_hi[keep]
    draw_band(rt_m, rt_lo, rt_hi, expression(R[t]), main1, "steelblue")
    abline(h = 1, lty = 3, col = "gray50")
  }

  if (which %in% c("beta", "both")) {
    beta_m  <- x$beta_mean[keep]
    beta_lo <- x$beta_lo[keep]
    beta_hi <- x$beta_hi[keep]
    beta_main <- if (which == "beta") main1 else expression(paste("Transmission rate ", beta[t]))
    draw_band(beta_m, beta_lo, beta_hi, expression(beta[t]), beta_main, "darkorange3")
  }

  
  invisible(x)
}


#' Posterior predictive forecasts from an npf_seir object
#'
#' Generates forward simulations from the filter's final filtered particles to produce
#' a posterior predictive distribution over future case counts.
#'
#' @param object An \code{npf_seir} object.
#' @param horizon Integer. Number of days ahead to forecast. Default 14.
#' @param n_draws Integer. Number of Monte Carlo draws. Default 1000.
#' @param ... Ignored.
#'
#' @return A named list with components: \code{mean} (numeric vector of
#'   length \code{horizon}, posterior predictive mean), \code{lo} and
#'   \code{hi} (95\% predictive interval bounds), and \code{draws}
#'   (\code{n_draws}-by-\code{horizon} matrix of Monte Carlo draws).
#' @export
#' @examples
#' \donttest{
#' fixed <- list(eps = 1/5, gamma = 1/10, delta = 1/(70*365),
#'               b = 1/(70*365), q = 0.3, N = 1e6, sigma_comp = 0.01)
#' x0    <- c(1e6 - 100, 50, 50, 0, log(0.25))
#' set.seed(1)
#' traj  <- seir_simulate(200, 0.05, 0.10, log(0.25), x0, fixed)
#' fit   <- npf_seir(traj$obs[-1], fixed, K=50, M=100, seed=42)
#' fc    <- predict(fit, horizon = 14)
#' print(fc$mean)
#' }
predict.npf_seir <- function(object, horizon = 14L, n_draws = 1000L, ...) {
  x <- object
  if (!is.numeric(horizon) || length(horizon) != 1L || is.na(horizon) ||
      horizon < 1L || horizon != as.integer(horizon)) {
    stop("horizon must be a positive integer.", call. = FALSE)
  }
  if (!is.numeric(n_draws) || length(n_draws) != 1L || is.na(n_draws) ||
      n_draws < 1L || n_draws != as.integer(n_draws)) {
    stop("n_draws must be a positive integer.", call. = FALSE)
  }

  horizon <- as.integer(horizon)
  n_draws <- as.integer(n_draws)
  K <- x$K; M <- x$M
  draws <- matrix(NA_real_, n_draws, horizon)

  for (d in seq_len(n_draws)) {
    k    <- sample.int(K, 1, prob = x$final_w)
    theta <- x$final_theta[k, ]
    ou   <- ou_params(theta[1], theta[2])
    alpha <- ou$alpha; tau <- ou$tau
    mu_psi <- theta[3]

    inner_idx <- sample.int(M, 1)
    state <- x$final_inner[[k]][inner_idx, ]

    for (h in seq_len(horizon)) {
      s <- state[1]; e <- state[2]; i <- state[3]
      r <- state[4]; psi <- state[5]
      beta <- exp(psi)
      foi  <- beta * s * i / x$fixed$N
      sc   <- x$fixed$sigma_comp

      state[1] <- max(s + (x$fixed$b*x$fixed$N - foi - x$fixed$delta*s) + sc*rnorm(1), 0)
      state[2] <- max(e + (foi - (x$fixed$eps+x$fixed$delta)*e)          + sc*rnorm(1), 0)
      state[3] <- max(i + (x$fixed$eps*e - (x$fixed$gamma+x$fixed$delta)*i) + sc*rnorm(1), 0)
      state[4] <- max(r + (x$fixed$gamma*i - x$fixed$delta*r)            + sc*rnorm(1), 0)
      state[5] <- mu_psi + alpha*(psi - mu_psi) + tau*rnorm(1)

      lam       <- max(x$fixed$q * x$fixed$eps * state[2], 1e-10)
      draws[d, h] <- rpois(1, lam)
    }
  }

  list(
    mean  = colMeans(draws),
    lo    = apply(draws, 2, quantile, 0.025),
    hi    = apply(draws, 2, quantile, 0.975),
    draws = draws
  )
}
