# ============================================================================
# PIT (Probability Integral Transform) Diagnostics
# ============================================================================

utils::globalVariables("count")

#' Extract PIT values from a fitted model
#'
#' Computes approximate Probability Integral Transform values using posterior
#' mean residuals and posterior mean margin parameters. Large departures from
#' uniformity can indicate model misfit, but these are not exact posterior
#' predictive PIT values.
#'
#' @param object A fitted model object.
#' @param ... Additional arguments (unused).
#'
#' @return A data frame with columns `time`, `variable`, `pit`.
#'
#' @details
#' PIT values are computed from posterior mean residuals and posterior mean
#' margin parameters. Treat them as a fast plug-in diagnostic rather than an
#' exact posterior predictive transform that integrates over full posterior
#' uncertainty. PIT diagnostics are currently implemented for the three core
#' single-level fit classes only.
#' @export
pit_values <- function(object, ...) {
  UseMethod("pit_values")
}

#' @rdname pit_values
#' @export
pit_values.default <- function(object, ...) {
  cli_abort("{.fun pit_values} is not defined for objects of class {.cls {class(object)[[1]]}}.")
}

#' @rdname pit_values
#' @export
pit_values.dcvar_hmm_fit <- function(object, ...) {
  if (isTRUE(object$switching)) {
    cli_abort(c(
      "{.fun pit_values} is not yet available for state-specific HMM fits (residuals are per-state).",
      "i" = "Use {.fun hmm_states} / {.fun hmm_state_params} for state-specific summaries."
    ))
  }
  NextMethod()
}

#' @rdname pit_values
#' @export
pit_values.dcvar_model_fit <- function(object, ...) {
  margins <- object$margins %||% "normal"

  # Extract posterior mean residuals and parameters
  eps_draws <- posterior::as_draws_matrix(.fit_draws(object$fit, "eps", backend = object$backend))
  n_time_eff <- object$stan_data$n_time - 1L
  D <- object$stan_data$D

  # Posterior mean residuals
  eps_mean <- matrix(NA_real_, n_time_eff, D)
  for (d in seq_len(D)) {
    cols <- paste0("eps[", seq_len(n_time_eff), ",", d, "]")
    eps_mean[, d] <- colMeans(eps_draws[, cols, drop = FALSE])
  }

  # Compute PIT values based on margin type
  pit_mat <- .pit_compute(object, eps_mean, margins)

  # Convert to long format
  var_names <- object$vars
  time_values <- .observed_time_values(object$stan_data, drop_first = TRUE)
  data.frame(
    time = rep(time_values, D),
    variable = rep(var_names, each = n_time_eff),
    pit = as.vector(pit_mat)
  )
}


#' Internal: posterior means for a margin parameter group, validated
#'
#' Mirrors the per-name validation in `.pit_mixed()`: a margin parameter that
#' is absent from the Stan output must abort with a clear message instead of
#' silently yielding all-NA PIT values.
#' @noRd
.pit_margin_means <- function(summ, pattern, expected, label) {
  rows <- grep(pattern, summ$variable)
  if (length(rows) != expected) {
    cli_abort(c(
      "PIT values require the Stan output {.val {label}}, which is missing from the fit.",
      "i" = "Custom Stan files must preserve the expected parameter names."
    ))
  }
  summ$mean[rows]
}


#' Internal: compute PIT values for given residuals and margin type
#' @noRd
.pit_compute <- function(object, eps_mean, margins) {
  T_eff <- nrow(eps_mean)
  D <- ncol(eps_mean)

  if (.is_mixed_margins(margins)) {
    return(.pit_mixed(object, eps_mean, margins))
  }

  switch(margins[[1L]],
    normal = {
      summ <- .fit_summary(object$fit, backend = object$backend)
      sigma_mean <- .pit_margin_means(summ, "^sigma_eps\\[", D, "sigma_eps")
      pit_mat <- matrix(NA_real_, T_eff, D)
      for (d in seq_len(D)) {
        pit_mat[, d] <- stats::pnorm(eps_mean[, d] / sigma_mean[d])
      }
      pit_mat
    },
    exponential = {
      .pit_exponential(object, eps_mean)
    },
    skew_normal = {
      .pit_skew_normal(object, eps_mean)
    },
    gamma = {
      .pit_gamma(object, eps_mean)
    },
    cli_abort("Unknown margin type: {.val {margins[[1L]]}}")
  )
}


#' Internal: PIT for exponential margins
#' @noRd
.pit_exponential <- function(object, eps_mean) {
  T_eff <- nrow(eps_mean)
  D <- ncol(eps_mean)
  summ <- .fit_summary(object$fit, backend = object$backend)
  skew_dir <- object$skew_direction

  sigma_exp_mean <- .pit_margin_means(summ, "^sigma_exp\\[", D, "sigma_exp")
  rate_exp <- 1.0 / sigma_exp_mean

  pit_mat <- matrix(NA_real_, T_eff, D)
  for (d in seq_len(D)) {
    x_shifted <- sigma_exp_mean[d] + skew_dir[d] * eps_mean[, d]
    u <- stats::pexp(x_shifted, rate = rate_exp[d])
    if (skew_dir[d] < 0) u <- 1.0 - u
    pit_mat[, d] <- u
  }
  pit_mat
}


#' Internal: PIT for skew-normal margins
#' @noRd
.pit_skew_normal <- function(object, eps_mean) {
  if (!requireNamespace("sn", quietly = TRUE)) {
    cli_abort("Package {.pkg sn} is required for skew-normal PIT values.")
  }
  T_eff <- nrow(eps_mean)
  D <- ncol(eps_mean)
  summ <- .fit_summary(object$fit, backend = object$backend)

  omega_mean <- .pit_margin_means(summ, "^omega\\[", D, "omega")
  delta_mean <- .pit_margin_means(summ, "^delta\\[", D, "delta")

  alpha_mean <- delta_mean / sqrt(1 - delta_mean^2)
  xi_mean <- -omega_mean * delta_mean * sqrt(2 / pi)

  pit_mat <- matrix(NA_real_, T_eff, D)
  for (d in seq_len(D)) {
    pit_mat[, d] <- sn::psn(eps_mean[, d],
      xi = xi_mean[d], omega = omega_mean[d], alpha = alpha_mean[d]
    )
  }
  pit_mat
}


#' Internal: PIT for gamma margins
#' @noRd
.pit_gamma <- function(object, eps_mean) {
  T_eff <- nrow(eps_mean)
  D <- ncol(eps_mean)
  summ <- .fit_summary(object$fit, backend = object$backend)
  skew_dir <- object$skew_direction

  sigma_gam_mean <- .pit_margin_means(summ, "^sigma_gam\\[", D, "sigma_gam")
  shape_gam_mean <- .pit_margin_means(summ, "^shape_gam$", 1L, "shape_gam")

  sqrt_shape <- sqrt(shape_gam_mean)
  rate_gam <- sqrt_shape / sigma_gam_mean

  pit_mat <- matrix(NA_real_, T_eff, D)
  for (d in seq_len(D)) {
    mean_x <- sqrt_shape * sigma_gam_mean[d]
    x_shifted <- mean_x + skew_dir[d] * eps_mean[, d]
    u <- stats::pgamma(x_shifted, shape = shape_gam_mean, rate = rate_gam[d])
    if (skew_dir[d] < 0) u <- 1.0 - u
    pit_mat[, d] <- u
  }
  pit_mat
}


#' Internal: PIT for mixed (per-variable) margins
#'
#' Dispatches each dimension to its own family's plug-in PIT using posterior
#' mean residuals and posterior mean margin parameters.
#' @noRd
.pit_mixed <- function(object, eps_mean, margins) {
  T_eff <- nrow(eps_mean)
  D <- ncol(eps_mean)
  summ <- .fit_summary(object$fit, backend = object$backend)
  skew_dir <- object$skew_direction
  gm <- function(name) {
    row <- match(name, summ$variable)
    if (is.na(row)) {
      cli_abort("PIT for mixed margins requires Stan output {.val {name}}.")
    }
    summ$mean[row]
  }

  needs_sn <- any(margins == "skew_normal")
  if (needs_sn && !requireNamespace("sn", quietly = TRUE)) {
    cli_abort("Package {.pkg sn} is required for skew-normal PIT values.")
  }

  pit_mat <- matrix(NA_real_, T_eff, D)
  for (d in seq_len(D)) {
    fam <- margins[[d]]
    if (identical(fam, "normal")) {
      sigma_d <- gm(paste0("sigma_eps[", d, "]"))
      pit_mat[, d] <- stats::pnorm(eps_mean[, d] / sigma_d)
    } else if (identical(fam, "exponential")) {
      sigma_exp_d <- gm(paste0("sigma_exp[", d, "]"))
      x_shifted <- sigma_exp_d + skew_dir[d] * eps_mean[, d]
      u <- stats::pexp(x_shifted, rate = 1.0 / sigma_exp_d)
      if (skew_dir[d] < 0) u <- 1.0 - u
      pit_mat[, d] <- u
    } else if (identical(fam, "skew_normal")) {
      omega_d <- gm(paste0("omega[", d, "]"))
      delta_d <- gm(paste0("delta[", d, "]"))
      alpha_d <- delta_d / sqrt(1 - delta_d^2)
      xi_d <- -omega_d * delta_d * sqrt(2 / pi)
      pit_mat[, d] <- sn::psn(eps_mean[, d], xi = xi_d, omega = omega_d, alpha = alpha_d)
    } else {
      sigma_gam_d <- gm(paste0("sigma_gam[", d, "]"))
      shape_gam_d <- gm(paste0("shape_gam[", d, "]"))
      sqrt_shape <- sqrt(shape_gam_d)
      mean_x <- sqrt_shape * sigma_gam_d
      x_shifted <- mean_x + skew_dir[d] * eps_mean[, d]
      u <- stats::pgamma(x_shifted, shape = shape_gam_d, rate = sqrt_shape / sigma_gam_d)
      if (skew_dir[d] < 0) u <- 1.0 - u
      pit_mat[, d] <- u
    }
  }
  pit_mat
}


#' KS test for PIT uniformity
#'
#' Runs a Kolmogorov-Smirnov test per variable to assess whether PIT values
#' are approximately uniform. This is a heuristic check on the plug-in PIT
#' values returned by [pit_values()], not an exact posterior predictive test.
#'
#' @param object A fitted model object.
#' @param ... Additional arguments (unused).
#'
#' @return A data frame with columns `variable`, `ks_statistic`, `p_value`, `n`.
#'
#' @details
#' This applies a Kolmogorov-Smirnov test to the approximate PIT values
#' returned by [pit_values()]. The result is a heuristic check and does not
#' account for serial dependence or full posterior uncertainty. PIT tests are
#' currently implemented for single-level fits only.
#' @export
pit_test <- function(object, ...) {
  UseMethod("pit_test")
}

#' @rdname pit_test
#' @export
pit_test.default <- function(object, ...) {
  cli_abort("{.fun pit_test} is not defined for objects of class {.cls {class(object)[[1]]}}.")
}

#' @rdname pit_test
#' @export
pit_test.dcvar_model_fit <- function(object, ...) {
  pit_df <- pit_values(object)

  results <- lapply(unique(pit_df$variable), function(v) {
    pit_v <- pit_df$pit[pit_df$variable == v]
    ks <- stats::ks.test(pit_v, "punif")
    data.frame(
      variable = v,
      ks_statistic = unname(ks$statistic),
      p_value = ks$p.value,
      n = length(pit_v)
    )
  })

  do.call(rbind, results)
}


#' @export
pit_values.dcvar_multilevel_fit <- function(object, ...) {
  cli_abort("PIT values are not yet supported for multilevel models.")
}

#' @export
pit_values.dcvar_sem_fit <- function(object, ...) {
  cli_abort("PIT values are not yet supported for SEM models.")
}

#' @export
pit_test.dcvar_multilevel_fit <- function(object, ...) {
  cli_abort("PIT tests are not yet supported for multilevel models.")
}

#' @export
pit_test.dcvar_sem_fit <- function(object, ...) {
  cli_abort("PIT tests are not yet supported for SEM models.")
}


#' Plot PIT histograms
#'
#' Creates faceted histograms of the approximate PIT values returned by
#' [pit_values()]. Under good model fit, these histograms should be roughly
#' uniform, but they remain plug-in diagnostics rather than exact posterior
#' predictive checks.
#'
#' @param object A fitted model object.
#' @param bins Number of histogram bins (default: 20).
#' @param ... Additional arguments (unused).
#'
#' @return A ggplot object.
#'
#' @details PIT histograms visualize the approximate plug-in PIT values
#'   returned by [pit_values()]. They are currently implemented for the three
#'   core single-level fit classes only.
#' @export
plot_pit <- function(object, bins = 20, ...) {
  pit_df <- pit_values(object)

  ggplot2::ggplot(pit_df, ggplot2::aes(x = .data$pit)) +
    ggplot2::geom_histogram(
      bins = bins,
      binwidth = 1 / bins,
      boundary = 0,
      ggplot2::aes(y = ggplot2::after_stat(count / sum(count) * bins)),
      fill = "steelblue", alpha = 0.7, color = "white"
    ) +
    ggplot2::geom_hline(yintercept = 1, linetype = "dashed", color = "red", linewidth = 0.8) +
    ggplot2::facet_wrap(~ variable) +
    ggplot2::labs(
      x = "PIT Value",
      y = "Density",
      title = "Approximate PIT Diagnostics",
      subtitle = "Red dashed line: uniform plug-in reference"
    ) +
    ggplot2::theme_minimal() +
    ggplot2::theme(
      plot.title = ggplot2::element_text(face = "bold"),
      strip.text = ggplot2::element_text(face = "bold")
    )
}


#' @rdname pit_values
#' @export
pit_values.dcvar_tv_fit <- function(object, ...) {
  D <- object$stan_data$D
  n_time_eff <- object$stan_data$n_time - 1L
  margins_vec <- rep(object$margins %||% "normal", length.out = D)
  skew_dir <- object$skew_direction %||% rep(1, D)

  needs_sn <- any(margins_vec == "skew_normal")
  if (needs_sn && !requireNamespace("sn", quietly = TRUE)) {
    cli_abort("Package {.pkg sn} is required for skew-normal PIT values.")
  }

  eps_draws <- posterior::as_draws_matrix(.fit_draws(object$fit, "eps", backend = object$backend))
  eps_mean <- matrix(NA_real_, n_time_eff, D)
  for (d in seq_len(D)) {
    cols <- paste0("eps[", seq_len(n_time_eff), ",", d, "]")
    eps_mean[, d] <- colMeans(eps_draws[, cols, drop = FALSE])
  }

  summ <- .fit_summary(object$fit, backend = object$backend)
  gm <- function(name) {
    row <- match(name, summ$variable)
    if (is.na(row)) {
      cli_abort("PIT for time-varying fits requires Stan output {.val {name}}.")
    }
    summ$mean[row]
  }

  # Per-time plug-in scale for each dimension (posterior means): the sigma_t
  # path when scales vary, otherwise the constant baseline tiled over time.
  scale_mean <- matrix(NA_real_, n_time_eff, D)
  for (d in seq_len(D)) {
    if (isTRUE(object$tv_sigma)) {
      rows <- match(paste0("sigma_t[", seq_len(n_time_eff), ",", d, "]"), summ$variable)
      if (anyNA(rows)) {
        cli_abort("PIT for time-varying fits requires Stan output {.val sigma_t}.")
      }
      scale_mean[, d] <- summ$mean[rows]
    } else {
      base <- switch(margins_vec[d],
        normal = gm(paste0("sigma_eps[", d, "]")),
        exponential = gm(paste0("sigma_exp[", d, "]")),
        skew_normal = gm(paste0("omega[", d, "]")),
        gamma = gm(paste0("sigma_gam[", d, "]"))
      )
      scale_mean[, d] <- base
    }
  }

  pit_mat <- matrix(NA_real_, n_time_eff, D)
  for (d in seq_len(D)) {
    fam <- margins_vec[d]
    s <- scale_mean[, d]
    if (identical(fam, "normal")) {
      pit_mat[, d] <- stats::pnorm(eps_mean[, d] / s)
    } else if (identical(fam, "exponential")) {
      u <- stats::pexp(s + skew_dir[d] * eps_mean[, d], rate = 1 / s)
      if (skew_dir[d] < 0) u <- 1 - u
      pit_mat[, d] <- u
    } else if (identical(fam, "skew_normal")) {
      delta_d <- gm(paste0("delta[", d, "]"))
      alpha_d <- delta_d / sqrt(1 - delta_d^2)
      omega_t <- if (isTRUE(object$tv_sigma)) {
        s / .skew_normal_residual_sd_factor(delta_d)
      } else {
        s
      }
      xi_t <- -omega_t * delta_d * sqrt(2 / pi)
      pit_mat[, d] <- sn::psn(eps_mean[, d], xi = xi_t, omega = omega_t, alpha = alpha_d)
    } else {
      shape_d <- gm(paste0("shape_gam[", d, "]"))
      mean_x <- sqrt(shape_d) * s
      u <- stats::pgamma(mean_x + skew_dir[d] * eps_mean[, d],
                         shape = shape_d, rate = sqrt(shape_d) / s)
      if (skew_dir[d] < 0) u <- 1 - u
      pit_mat[, d] <- u
    }
  }

  time_values <- .observed_time_values(object$stan_data, drop_first = TRUE)
  data.frame(
    time = rep(time_values, D),
    variable = rep(object$vars, each = n_time_eff),
    pit = as.vector(pit_mat)
  )
}
