# approach A: truncate residuals to floor_eps, smooth log, then exponentiate to ensure nonnegative

ARCHmXAdditive_backfitting <- function(target_sum_nonpar, X, bandwidth = NULL,
                        max_iter = 50, tol = 1e-4, floor_eps = 1e-8) {
  # target_sum_nonpar: vector of length n with the "total" nonparametric component
  #                    i.e., sigma2_resid = eps^2 - arch_component   (can be < 0)
  # We fit m_ell >= 0 via m_ell(x) = exp(g_ell(x)), smoothing log of partial residuals.
  
  n <- nrow(X) # number of samples
  d <- ncol(X) # number of covariates
  
  # Initialize m values strictly positive (tiny), and storage for functions
  m_values <- matrix(floor_eps, nrow = n, ncol = d)
  m_funcs  <- vector("list", d)
  
  # bandwidths
  if (is.null(bandwidth)) {
    bandwidth <- rep(NA_real_, d)
    for (ell in 1:d) {
      x_range <- diff(range(X[, ell]))
      bandwidth[ell] <- 0.5 * x_range * n^(-1/5)
    }
  }
  
  # backfitting
  for (iter in 1:max_iter) {
    m_old <- m_values
    
    for (ell in 1:d) {
      # Partial residual for component ell
      partial_resid <- target_sum_nonpar - rowSums(m_values[, -ell, drop = FALSE])
      
      # Enforce positivity before log
      y_pos <- pmax(partial_resid, floor_eps)
      # Smooth log(y_pos) vs X[,ell]
      fit <- locpoly(X[, ell], log(y_pos), bandwidth = bandwidth[ell],
                     degree = 1, gridsize = max(100, n), range.x = range(X[, ell])) # from KernSmooth
      
      # if (any(!is.finite(fit$y))) {
      #   print("non-finite fit$y:")
      #   print(fit$y[!is.finite(fit$y)])
      # }
      good <- is.finite(fit$x) & is.finite(fit$y)
      
      # Predict on observed X and exponentiate back
      # ghat_at_x  <- approx(fit$x, fit$y, xout = X[, ell], rule = 2)$y
      ghat_at_x <- approx(fit$x[good], fit$y[good], xout = X[, ell], rule = 2)$y
      ghat_at_x[!is.finite(ghat_at_x)] <- log(floor_eps) # log because we exponentiate later
      cap <- log(.Machine$double.xmax) - 1
      ghat_at_x <- pmin(ghat_at_x, cap)

      # if (any(!is.finite(ghat_at_x))) {
      #   print("non-finite ghat_at_x:")
      #   print(ghat_at_x[!is.finite(ghat_at_x)])
      # }
      
      # m_values[, ell] <- pmax(replace(exp(ghat_at_x), !is.finite(exp(ghat_at_x)), floor_eps), floor_eps)
      m_values[, ell] <- pmax(exp(ghat_at_x), floor_eps)
  
      # Store the function on a grid
      m_funcs[[ell]] <- (function(xgrid, ygrid) {
        # returns m_ell(x) = exp(g_ell(x))
        f_lin <- approxfun(xgrid, ygrid, rule = 2)
        function(z){
          val <- f_lin(z)
          val[!is.finite(val)] <- log(floor_eps)
          val <- pmin(val, cap)
          pmax(exp(val), floor_eps)
          # pmax(exp(f_lin(z)), floor_eps) # 
        } # return m_hat (max(g_hat, floor_eps))
      })(fit$x, fit$y)
    }
    
    # if (max(abs(m_values - m_old)) < tol) break
    change <- m_values - m_old
    change[!is.finite(change)] <- Inf
    if (max(abs(change)) < tol) break
  }
  
  list(m_values = m_values, m_funcs = m_funcs, bandwidth = bandwidth)
}
