###############################################################################
# Reconciliation with Bottom-Up Importance Sampling (BUIS)
###############################################################################

# Checks that there is no bottom continuous variable child of a
# discrete upper variable
.check_hierfamily_rel <- function(sh.res, distr, debug = FALSE) {
  for (bi in seq_along(distr[sh.res$bottom_idxs])) {
    distr_bottom <- distr[sh.res$bottom_idxs][[bi]]
    rel_upper_i <- sh.res$A[, bi]
    rel_distr_upper <- unlist(distr[sh.res$upper_idxs])[rel_upper_i == 1]
    err_message <- "A continuous bottom distribution cannot be child of a discrete one."
    if (distr_bottom == "continuous") {
      if (sum(rel_distr_upper == "discrete") | sum(rel_distr_upper %in% .DISCR_DISTR)) {
        if (debug) {
          return(-1)
        } else {
          stop(err_message)
        }
      }
    }
    if (distr_bottom %in% .CONT_DISTR) {
      if (sum(rel_distr_upper == "discrete") | sum(rel_distr_upper %in% .DISCR_DISTR)) {
        if (debug) {
          return(-1)
        } else {
          stop(err_message)
        }
      }
    }
  }
  if (debug) {
    return(0)
  }
}

.emp_pmf <- function(l, density_samples) {
  empirical_pmf <- PMF_from_samples(density_samples)
  w <- sapply(l, function(i) empirical_pmf[i + 1])
  return(w)
}

.compute_weights <- function(b, u, in_type_, distr_) {
  if (in_type_ == "samples") {
    if (distr_ == "discrete") {
      # Discrete samples
      w <- .emp_pmf(b, u)
    } else if (distr_ == "continuous") {
      # KDE
      d <- stats::density(u, bw = "SJ", n = 2**16)
      df <- stats::approxfun(d)
      w <- df(b)
    }
    # be sure no NA are returned, if NA, we want 0:
    # for the discrete branch:   if b_i !in u        --> NA
    # for the continuous branch: if b_i !in range(u) --> NA
    w[is.na(w)] <- 0
  } else if (in_type_ == "params") {
    w <- .distr_pmf(b, u, distr_) # this never returns NA
  }
  # be sure not to return all 0 weights, return ones instead
  # if (sum(w) == 0) { w = w + 1 }
  return(w)
}

#' @title BUIS for Probabilistic Reconciliation of forecasts via conditioning
#'
#' @description
#'
#' Uses the Bottom-Up Importance Sampling algorithm to draw samples from the reconciled
#' forecast distribution, obtained via conditioning.
#'
#' @details
#'
#' The parameter `base_fc` is a list containing n = n_upper + n_bottom elements.
#' The first n_upper elements of the list are the upper base forecasts, in the order given by the rows of A.
#' The elements from n_upper+1 until the end of the list are the bottom base forecasts, in the order given by the columns of A.
#'
#' The i-th element depends on the values of `in_type[[i]]` and `distr[[i]]`.
#'
#' If `in_type[[i]]`='samples', then `base_fc[[i]]` is a vector containing samples from the base forecast distribution.
#'
#' If `in_type[[i]]`='params', then `base_fc[[i]]` is a list containing the estimated:
#'
#' * mean and sd for the Gaussian base forecast if `distr[[i]]`='gaussian', see \link[stats]{Normal};
#' * lambda for the Poisson base forecast if `distr[[i]]`='poisson', see \link[stats]{Poisson};
#' * size and prob (or mu) for the negative binomial base forecast if `distr[[i]]`='nbinom', see \link[stats]{NegBinomial}.
#'
#' See the description of the parameters `in_type` and `distr` for more details.
#'
#' Warnings are triggered from the Importance Sampling step if:
#'
#' * weights are all zeros, then the upper is ignored during reconciliation;
#' * the effective sample size is < 200;
#' * the effective sample size is < 1% of the sample size (`num_samples` if `in_type` is 'params' or the size of the base forecast if if `in_type` is 'samples').
#'
#' Note that warnings are an indication that the base forecasts might have issues.
#' Please check the base forecasts in case of warnings.
#'
#' @param A aggregation matrix (n_upper x n_bottom).
#' @param base_fc A list containing the base_forecasts, see details.
#' @param in_type A string or a list of length n_upper + n_bottom. If it is a list the i-th element is a string with two possible values:
#'
#' * 'samples' if the i-th base forecasts are in the form of samples;
#' * 'params'  if the i-th base forecasts are in the form of estimated parameters.
#'
#' If it `in_type` is a string it is assumed that all base forecasts are of the same type.
#'
#' @param distr A string or a list of length n_upper + n_bottom describing the type of base forecasts.
#' If it is a list the i-th element is a string with two possible values:
#'
#' * 'continuous' or 'discrete' if `in_type[[i]]`='samples';
#' * 'gaussian', 'poisson' or 'nbinom' if `in_type[[i]]`='params'.
#'
#' If `distr` is a string it is assumed that all distributions are of the same type.
#'
#' @param num_samples Number of samples drawn from the reconciled distribution.
#'        This is ignored if `bottom_in_type='samples'`; in this case, the number of reconciled samples is equal to
#'        the number of samples of the base forecasts.
#'
#' @param suppress_warnings Logical. If \code{TRUE}, no warnings about effective sample size
#'        are triggered. If \code{FALSE}, warnings are generated. Default is \code{FALSE}. See Details.
#' @param return_upper Logical, whether to return the reconciled parameters for the upper variables (default is TRUE).
#' @param seed Seed for reproducibility.
#'
#' @return A list containing the reconciled forecasts. The list has the following named elements:
#'
#' * `bottom_rec_samples`: a matrix (n_bottom x `num_samples`) containing the reconciled samples for the bottom time series;
#' * `upper_rec_samples`: (only if `return_upper = TRUE`) a matrix (n_upper x `num_samples`) containing the reconciled samples for the upper time series.
#'
#' @examples
#'
#' library(bayesRecon)
#'
#' # Create a minimal hierarchy with 2 bottom and 1 upper variable
#' rec_mat <- get_reconc_matrices(agg_levels = c(1, 2), h = 2)
#' A <- rec_mat$A
#' S <- rec_mat$S
#'
#'
#' # 1) Gaussian base forecasts
#'
#' # Set the parameters of the Gaussian base forecast distributions
#' mu1 <- 2
#' mu2 <- 4
#' muY <- 9
#' mus <- c(muY, mu1, mu2)
#'
#' sigma1 <- 2
#' sigma2 <- 2
#' sigmaY <- 3
#' sigmas <- c(sigmaY, sigma1, sigma2)
#'
#' base_fc <- list()
#' for (i in 1:length(mus)) {
#'   base_fc[[i]] <- list(mean = mus[[i]], sd = sigmas[[i]])
#' }
#'
#'
#' # Sample from the reconciled forecast distribution using the BUIS algorithm
#' buis <- reconc_BUIS(A, base_fc,
#'   in_type = "params",
#'   distr = "gaussian", num_samples = 100000, seed = 42
#' )
#'
#' samples_buis <- rbind(buis$upper_rec_samples, buis$bottom_rec_samples)
#'
#' # In the Gaussian case, the reconciled distribution is still Gaussian and can be
#' # computed in closed form
#' Sigma <- diag(sigmas^2) # transform into covariance matrix
#' analytic_rec <- reconc_gaussian(A,
#'   base_fc_mean = mus,
#'   base_fc_cov = Sigma
#' )
#'
#' # Compare the reconciled means obtained analytically and via BUIS
#' print(c(S %*% analytic_rec$bottom_rec_mean))
#' print(rowMeans(samples_buis))
#'
#'
#' # 2) Poisson base forecasts
#'
#' # Set the parameters of the Poisson base forecast distributions
#' lambda1 <- 2
#' lambda2 <- 4
#' lambdaY <- 9
#' lambdas <- c(lambdaY, lambda1, lambda2)
#'
#' base_fc <- list()
#' for (i in 1:length(lambdas)) {
#'   base_fc[[i]] <- list(lambda = lambdas[i])
#' }
#'
#' # Sample from the reconciled forecast distribution using the BUIS algorithm
#' buis <- reconc_BUIS(A, base_fc,
#'   in_type = "params",
#'   distr = "poisson", num_samples = 100000, seed = 42
#' )
#' samples_buis <- rbind(buis$upper_rec_samples, buis$bottom_rec_samples)
#'
#' # Print the reconciled means
#' print(rowMeans(samples_buis))
#'
#' @references
#' Zambon, L., Azzimonti, D. & Corani, G. (2024).
#' *Efficient probabilistic reconciliation of forecasts for real-valued and count time series*.
#' Statistics and Computing 34 (1), 21.
#' \doi{10.1007/s11222-023-10343-y}.
#'
#'
#' @seealso
#' [reconc_gaussian()]
#'
#' @export
reconc_BUIS <- function(A,
                        base_fc,
                        in_type,
                        distr,
                        num_samples = 2e4,
                        suppress_warnings = FALSE,
                        return_upper = TRUE,
                        seed = NULL) {
  if (!is.null(seed)) set.seed(seed)
  n_upper <- nrow(A)
  n_bottom <- ncol(A)
  n_tot <- length(base_fc)

  # Transform distr and in_type into lists
  if (!is.list(distr)) {
    distr <- rep(list(distr), n_tot)
  }
  if (!is.list(in_type)) {
    in_type <- rep(list(in_type), n_tot)
  }

  # Ensure that data inputs are valid
  .check_input_BUIS(A, base_fc, in_type, distr)

  # Split bottoms, uppers
  # the first nrow(A) elements of base_fc are upper
  # the second ncol(A) elements of base_fc are lower

  split_hierarchy_res <- list(
    A = A,
    upper = base_fc[1:nrow(A)],
    bottom = base_fc[(nrow(A) + 1):n_tot],
    upper_idxs = 1:nrow(A),
    bottom_idxs = (nrow(A) + 1):n_tot
  )
  upper_base_fc <- split_hierarchy_res$upper
  bottom_base_fc <- split_hierarchy_res$bottom

  # Check on continuous/discrete in relationship to the hierarchy
  .check_hierfamily_rel(split_hierarchy_res, distr)

  # H, G
  is_hier <- .check_hierarchical(A)
  # If A is hierarchical we do not solve the integer linear programming problem
  if (is_hier) {
    H <- A
    G <- NULL
    upper_base_fc_H <- upper_base_fc
    upper_base_fc_G <- NULL
    in_typeH <- in_type[split_hierarchy_res$upper_idxs]
    distr_H <- distr[split_hierarchy_res$upper_idxs]
    in_typeG <- NULL
    distr_G <- NULL
  } else {
    get_HG_res <- .get_HG(A, upper_base_fc, distr[split_hierarchy_res$upper_idxs], in_type[split_hierarchy_res$upper_idxs])
    H <- get_HG_res$H
    upper_base_fc_H <- get_HG_res$Hv
    G <- get_HG_res$G
    upper_base_fc_G <- get_HG_res$Gv
    in_typeH <- get_HG_res$Hin_type
    distr_H <- get_HG_res$Hdistr
    in_typeG <- get_HG_res$Gin_type
    distr_G <- get_HG_res$Gdistr
  }

  # Reconciliation using BUIS

  # 1. Bottom samples
  B <- list()
  in_type_bottom <- in_type[split_hierarchy_res$bottom_idxs]
  for (bi in 1:n_bottom) {
    if (in_type_bottom[[bi]] == "samples") {
      B[[bi]] <- unlist(bottom_base_fc[[bi]])
    } else if (in_type_bottom[[bi]] == "params") {
      B[[bi]] <- .distr_sample(
        bottom_base_fc[[bi]],
        distr[split_hierarchy_res$bottom_idxs][[bi]],
        num_samples
      )
    }
  }
  B <- do.call("cbind", B) # B is a matrix (num_samples x n_bottom)

  out <- .core_reconc_BUIS(
    A = A, H = H, G = G, B = B,
    upper_base_fc_H = upper_base_fc_H,
    in_typeH = in_typeH, distr_H = distr_H,
    upper_base_fc_G = upper_base_fc_G,
    in_typeG = in_typeG, distr_G = distr_G,
    .comp_w = .compute_weights,
    suppress_warnings = suppress_warnings,
    return_upper = return_upper
  )

  # # Bottom-Up IS on the hierarchical part
  # for (hi in 1:nrow(H)) {
  #   c = H[hi, ]
  #   b_mask = (c != 0)
  #   weights = .compute_weights(
  #     b = (B %*% c),
  #     # (num_samples x 1)
  #     u = upper_base_fc_H[[hi]],
  #     in_type_ = in_typeH[[hi]],
  #     distr_ = distr_H[[hi]]
  #   )
  #   check_weights_res = .check_weights(weights)
  #   if (check_weights_res$warning & !suppress_warnings) {
  #     warning_msg = check_weights_res$warning_msg
  #     # add information to the warning message
  #     upper_fromA_i = which(lapply(seq_len(nrow(A)), function(i) sum(abs(A[i,] - c))) == 0)
  #     for (wmsg in warning_msg) {
  #       wmsg = paste(wmsg, paste0("Check the upper forecast at index: ", upper_fromA_i,"."))
  #       warning(wmsg)
  #     }
  #   }
  #   if(check_weights_res$warning & (1 %in% check_weights_res$warning_code)){
  #     next
  #   }
  #   B[, b_mask] = .resample(B[, b_mask], weights)
  # }

  # if (!is.null(G)) {
  #   # Plain IS on the additional constraints
  #   weights = matrix(1, nrow = nrow(B))
  #   for (gi in 1:nrow(G)) {
  #     c = G[gi, ]
  #     weights = weights * .compute_weights(
  #       b = (B %*% c),
  #       u = upper_base_fc_G[[gi]],
  #       in_type_ = in_typeG[[gi]],
  #       distr_ = distr_G[[gi]]
  #     )
  #   }
  #   check_weights_res = .check_weights(weights)
  #   if (check_weights_res$warning & !suppress_warnings) {
  #     warning_msg = check_weights_res$warning_msg
  #     # add information to the warning message
  #     upper_fromA_i = c()
  #     for (gi in 1:nrow(G)) {
  #       c = G[gi, ]
  #       upper_fromA_i = c(upper_fromA_i,
  #                         which(lapply(seq_len(nrow(A)), function(i) sum(abs(A[i,] - c))) == 0))
  #     }
  #     for (wmsg in warning_msg) {
  #       wmsg = paste(wmsg, paste0("Check the upper forecasts at index: ", paste0("{",paste(upper_fromA_i, collapse = ","), "}.")))
  #       warning(wmsg)
  #     }
  #   }
  #   if(!(check_weights_res$warning & (1 %in% check_weights_res$warning_code))){
  #     B = .resample(B, weights)
  #   }

  # }

  return(out)
}

#' Core Reconciliation via Bayesian Universality Information Sharing
#'
#' Internal function that performs the core reconciliation logic for the BUIS method, which
#' reconciles forecasts using importance sampling based on both hierarchical and equality
#' constraints through a Bayesian framework.
#'
#' @param A Matrix defining the overall hierarchy.
#' @param H Matrix defining hierarchical constraints.
#' @param G Matrix defining general linear constraints.
#' @param B Matrix of bottom level base forecast samples.
#' @param upper_base_fc_H List of upper base forecasts for hierarchical constraints.
#' @param in_typeH Character string specifying input type for H forecasts ('pmf', 'samples', or 'params').
#' @param distr_H Character string specifying distribution type for H forecasts ('poisson' or 'nbinom').
#' @param upper_base_fc_G List of upper base forecasts for general constraints.
#' @param in_typeG Character string specifying input type for G forecasts ('pmf', 'samples', or 'params').
#' @param distr_G Character string specifying distribution type for G forecasts ('poisson' or 'nbinom').
#' @param .comp_w Function to compute weights for importance sampling. Default is `.compute_weights`.
#' @param suppress_warnings Logical. If TRUE, suppresses warnings about sample quality. Default is FALSE.
#'
#' @return A list containing:
#'   \itemize{
#'     \item `bottom_rec`: List with reconciled bottom forecasts (pmf and/or samples).
#'     \item `upper_rec_H`: List with reconciled upper forecasts for H constraints.
#'     \item `upper_rec_G`: List with reconciled upper forecasts for G constraints.
#'   }
#'
#' @keywords internal
#' @export
.core_reconc_BUIS <- function(A,
                              H, G,
                              B,
                              upper_base_fc_H,
                              in_typeH,
                              distr_H,
                              upper_base_fc_G,
                              in_typeG,
                              distr_G,
                              .comp_w = .compute_weights,
                              suppress_warnings = FALSE,
                              return_upper = TRUE) {
  # Hierarchical part
  for (hi in 1:nrow(H)) {
    c <- H[hi, ]
    b_mask <- (c != 0)
    weights <- .comp_w(
      b = (B %*% c),
      # (num_samples x 1)
      u = upper_base_fc_H[[hi]],
      in_type_ = in_typeH[[hi]],
      distr_ = distr_H[[hi]]
    )
    check_weights_res <- .check_weights(weights)
    if (check_weights_res$warning & !suppress_warnings) {
      warning_msg <- check_weights_res$warning_msg
      # add information to the warning message
      upper_fromA_i <- which(lapply(seq_len(nrow(A)), function(i) sum(abs(A[i, ] - c))) == 0)
      for (wmsg in warning_msg) {
        wmsg <- paste(wmsg, paste0("Check the upper forecast at index: ", upper_fromA_i, "."))
        warning(wmsg)
      }
    }
    if (check_weights_res$warning & (1 %in% check_weights_res$warning_code)) {
      next
    }
    B[, b_mask] <- .resample(B[, b_mask], weights)
  }

  # Non-hierarchical part
  if (!is.null(G)) {
    # Plain IS on the additional constraints
    weights <- matrix(1, nrow = nrow(B))
    for (gi in 1:nrow(G)) {
      c <- G[gi, ]
      weights <- weights * .comp_w(
        b = (B %*% c),
        u = upper_base_fc_G[[gi]],
        in_type_ = in_typeG[[gi]],
        distr_ = distr_G[[gi]]
      )
    }
    check_weights_res <- .check_weights(weights)
    if (check_weights_res$warning & !suppress_warnings) {
      warning_msg <- check_weights_res$warning_msg
      # add information to the warning message
      upper_fromA_i <- c()
      for (gi in 1:nrow(G)) {
        c <- G[gi, ]
        upper_fromA_i <- c(
          upper_fromA_i,
          which(lapply(seq_len(nrow(A)), function(i) sum(abs(A[i, ] - c))) == 0)
        )
      }
      for (wmsg in warning_msg) {
        wmsg <- paste(wmsg, paste0("Check the upper forecasts at index: ", paste0("{", paste(upper_fromA_i, collapse = ","), "}.")))
        warning(wmsg)
      }
    }
    if (!(check_weights_res$warning & (1 %in% check_weights_res$warning_code))) {
      B <- .resample(B, weights)
    }
  }
  B <- t(B)
  U <- A %*% B
  out <- list(bottom_rec_samples = B)
  if (return_upper) {
    out$upper_rec_samples <- U
  }
  return(out)
}
