#' Transform pmf model output into a forecast object
#'
#' @param model_out_tbl Model output tibble with predictions
#' @param oracle_output Predictions that would have been generated by an oracle
#' model that knew the observed target data values in advance
#' @param output_type_id_order For ordinal variables in pmf format, this is a
#' vector of levels for pmf forecasts, in increasing order of the levels. The
#' order of the values for the output_type_id can be found by referencing the
#' hub's tasks.json configuration file. For all output types other than pmf,
#' this is ignored.
#'
#' @return A `forecast_ordinal` object (when `output_type_id_order` is provided)
#'   or a `forecast_nominal` object (when `output_type_id_order` is `NULL`).
#'
#' @examplesIf requireNamespace("hubExamples", quietly = TRUE)
#' # Nominal pmf forecast (no output_type_id_order provided)
#' pmf_outputs <- hubExamples::forecast_outputs |>
#'   dplyr::filter(.data[["output_type"]] == "pmf")
#'
#' nominal_forecast <- transform_pmf_model_out(
#'   model_out_tbl = pmf_outputs,
#'   oracle_output = hubExamples::forecast_oracle_output
#' )
#' nominal_forecast
#'
#' # Ordinal pmf forecast (output_type_id_order provided)
#' ordinal_forecast <- transform_pmf_model_out(
#'   model_out_tbl = pmf_outputs,
#'   oracle_output = hubExamples::forecast_oracle_output,
#'   output_type_id_order = c("low", "moderate", "high", "very high")
#' )
#' ordinal_forecast
#'
#' @importFrom rlang .data
#' @export
transform_pmf_model_out <- function(
  model_out_tbl,
  oracle_output,
  output_type_id_order = NULL
) {
  model_out_tbl <- validate_model_oracle_out(model_out_tbl, oracle_output)

  # subset both model_out_tbl and oracle_output to output_type == "pmf"
  model_out_tbl <- model_out_tbl |>
    dplyr::filter(.data[["output_type"]] == "pmf")

  if (c("output_type") %in% colnames(oracle_output)) {
    oracle_output <- oracle_output |>
      dplyr::filter(.data[["output_type"]] == "pmf") |>
      dplyr::select(-c("output_type"))
  }

  # validate or set output_type_id_order
  if (!is.null(output_type_id_order)) {
    output_type_id_order <- validate_output_type_id_order(
      output_type_id_order,
      model_out_tbl
    )
    is_ordinal <- TRUE
  } else {
    is_ordinal <- FALSE
    output_type_id_order <- unique(model_out_tbl$output_type_id)
  }

  task_id_cols <- get_task_id_cols(model_out_tbl)

  # assemble data in scoringutils format
  model_out_tbl <- model_out_tbl |>
    dplyr::rename(model = "model_id")

  data <- dplyr::left_join(
    model_out_tbl,
    oracle_output,
    by = c(
      task_id_cols[task_id_cols %in% colnames(oracle_output)],
      "output_type_id"
    ),
    relationship = "many-to-one"
  ) |>
    dplyr::group_by(dplyr::across(dplyr::all_of(c(task_id_cols, "model")))) |>
    dplyr::mutate(
      observation = .data[["output_type_id"]][.data[["oracle_value"]] == 1],
      observation = factor(
        .data[["observation"]],
        levels = output_type_id_order,
        ordered = is_ordinal
      ),
      output_type_id = factor(
        .data[["output_type_id"]],
        levels = output_type_id_order,
        ordered = is_ordinal
      )
    ) |>
    dplyr::ungroup() |>
    dplyr::select(-dplyr::all_of("oracle_value"))

  if (is_ordinal) {
    forecast_pmf <- scoringutils::as_forecast_ordinal(
      data,
      forecast_unit = c("model", task_id_cols),
      observed = "observation",
      predicted = "value",
      predicted_label = "output_type_id"
    )
  } else {
    forecast_pmf <- scoringutils::as_forecast_nominal(
      data,
      forecast_unit = c("model", task_id_cols),
      observed = "observation",
      predicted = "value",
      predicted_label = "output_type_id"
    )
  }

  forecast_pmf
}


#' Validate `output_type_id_order` for ordinal variables:
#' - Must be a vector
#' - Must be (set-)equal to the set of all unique `output_type_id` values in `model_out_tbl`
#'
#' This function expects that the model_out_tbl has already been filtered to
#' output_type == "pmf". This is done in `transform_pmf_model_out`.
#'
#' @noRd
validate_output_type_id_order <- function(output_type_id_order, model_out_tbl) {
  if (!is.vector(output_type_id_order)) {
    cli::cli_abort("`output_type_id_order` must be a vector.")
  }

  present_levels <- unique(model_out_tbl$output_type_id)
  extra_order_levels <- setdiff(output_type_id_order, present_levels)
  missing_order_levels <- setdiff(present_levels, output_type_id_order)
  if (length(extra_order_levels) > 0 || length(missing_order_levels) > 0) {
    cli::cli_abort(
      c(
        "x" = "`output_type_id_order` must align with the set of all unique
               `output_type_id` values in `model_out_tbl`.",
        "!" = ifelse(
          length(extra_order_levels) == 0,
          NULL,
          "The following levels were present in `output_type_id_order` but not in `model_out_tbl`:
          {.val {extra_order_levels}}."
        ),
        "!" = ifelse(
          length(missing_order_levels) == 0,
          NULL,
          "The following levels were present in `model_out_tbl` but not in `output_type_id_order`:
          {.val {missing_order_levels}};"
        )
      )
    )
  }

  output_type_id_order
}
