qmle_loglik <- function(params, eps, X, p, delta, m_funcs, start_idx = 1) {
  n <- length(eps)
  d <- ncol(X)
  
  
  alphas <- params[1:(p)]
  
  if ( any(alphas < 0) || sum(alphas) >= 1) {
    return(-1e10)
  }
  
  sigma2 <- numeric(n)
  sigma2[1:p] <- var(eps[start_idx:n])
  
  for (t in max(p+1, start_idx):n) {
    arch_component <-  sum(alphas * eps[(t-1):(t-p)]^2)
    nonpar_component <- 0
    for (ell in 1:d) {
      x_val <- X[t-1, ell]
      nonpar_component <- nonpar_component + m_funcs[[ell]](x_val)
    }
    sigma2[t] <- arch_component + nonpar_component
    if (sigma2[t] <= 0) sigma2[t] <- 1e-6
  }
  
  valid_idx <- max(p+1, start_idx):n
  loglik <- -0.5 * sum(log(2 * pi * sigma2[valid_idx]) + 
                         eps[valid_idx]^2 / sigma2[valid_idx])
  
  return(loglik)
}


######################
ARCHmXAdditive_estimate <- function(eps, X, p, delta, train_end = NULL, n_burn = 10, 
                                   max_outer_iter = 30, tol = 1e-4, 
                                   backfit.method = "smooth_means_Positive_Constraint_Exp") 
{
  n <- length(eps)
  d <- ncol(X)
  if (is.null(train_end)) 
    train_end <- n
  start_idx <- n_burn + 1
  
  # init m's as zero functions (will become positive)
  m_funcs <- vector("list", d)
  for (ell in 1:d) m_funcs[[ell]] <- function(x) 0
  
  
  alphas_init <- rep(0.05, p) # rep(0.6 / p, p)
  params_init <- c( alphas_init)
  converged_outer <- FALSE
  
  eps_train <- eps[1:train_end]
  X_train <- X[1:train_end, , drop = FALSE]
  
  for (outer_iter in 1:max_outer_iter) {
    params_old <- params_init
    
    optim_result <- optim(
      par = params_init,
      fn = function(params) -qmle_loglik(params, eps_train, X_train, p, delta, m_funcs, start_idx),
      method = "L-BFGS-B",
      lower = c(1e-6, rep(0, p)),
      upper = c(Inf, rep(0.99, p)),
      control = list(maxit = 2000, factr = 1e7)
    )
    
    params_new <- optim_result$par
    alphas <- params_new[1:(p)]
    
    # ----------- CHANGED: build target sum for nonparam part (can be negative) -----------
    sigma2_resid <- numeric(train_end)
    sigma2_resid[1:p] <- NA
    for (t in (p+1):train_end) {
      arch_component <- sum(alphas * eps_train[(t-1):(t-p)]^2)
      sigma2_resid[t] <- eps_train[t]^2 - arch_component
    }
    # Valid training slice and lagged X for m_ell(X_{t-1,ell})
    valid_idx <- start_idx:train_end
    X_lagged <- X_train[valid_idx - 1, , drop = FALSE]
    resid_valid <- sigma2_resid[valid_idx]
    
    # ----------- CHANGED: positive backfitting on log scale -----------------------------
    if (backfit.method == "smooth_means_Positive_Constraint_Exp")
      backfit_result <- ARCHmXAdditive_backfitting(resid_valid, X_lagged, max_iter = 100)
    else if (backfit.method == "smooth_means_lowes")
      backfit_result <- ARCHmXAdditive_backfitting_lowes(resid_valid, X_lagged, max_iter = 100)
    else if (backfit.method == "smooth_variance")
      backfit_result <- ARCHmXAdditive_backfitting_variance_smooth(resid_valid, X_lagged, max_iter = 100)
    m_funcs <- backfit_result$m_funcs
    
    # Check parameter change
    if (max(abs(params_new - params_old)) < tol) {
      converged_outer <- TRUE
      break
    }
    params_init <- params_new
  }
  
  optim_result <- optim(
    par = params_init,
    fn = function(params) -qmle_loglik(params, eps_train, X_train, p, delta, m_funcs, start_idx),
    method = "L-BFGS-B",
    lower = c(1e-6, rep(0, p)),
    upper = c(Inf, rep(0.99, p)),
    control = list(maxit = 2000, factr = 1e7)
  )
  params_final <- optim_result$par
  alphas <- params_final[1:(p)]
  
  # Fitted sigma2 over ALL data
  sigma2_hat <- numeric(n)
  sigma2_hat[1:p] <- NA
  for (t in (p+1):n) {
    arch_component <- sum(alphas * eps[(t-1):(t-p)]^2)
    nonpar_component <- 0
    for (ell in 1:d) nonpar_component <- nonpar_component + m_funcs[[ell]](X[t-1, ell])
    sigma2_hat[t] <- max(arch_component + nonpar_component, 1e-6)
  }
  
  sigma_hat <- sqrt(sigma2_hat)
  sigma_realized <- abs(eps)
  
  # Metrics (unchanged) ...
  
  # Export m-hats on a grid
  m_hats <- list()
  for (ell in 1:d) {
    x_grid <- seq(min(X[, ell]), max(X[, ell]), length.out = 100)
    m_grid <- sapply(x_grid, m_funcs[[ell]])
    m_hats[[ell]] <- list(x = x_grid, m = m_grid)
  }
  
  list(
    alphas = alphas,
    m_hats = m_hats,
    m_funcs = m_funcs,
    sigma_hat = sigma_hat,
    #sigma_realized = sigma_realized,
    #MAE_train = mean(abs(sigma_hat[(start_idx):train_end] - sigma_realized[(start_idx):train_end]), na.rm=TRUE),
    #MAPE_train = mean(abs((sigma_hat[(start_idx):train_end] - sigma_realized[(start_idx):train_end]) / 
    #                        sigma_realized[(start_idx):train_end]), na.rm=TRUE) * 100,
    #RMSE_train = sqrt(mean((sigma_hat[(start_idx):train_end] - sigma_realized[(start_idx):train_end])^2, na.rm=TRUE)),
    #MAE_test = if (train_end < n) mean(abs(sigma_hat[(train_end+1):n] - sigma_realized[(train_end+1):n]), na.rm=TRUE) else NA,
    #MAPE_test = if (train_end < n) mean(abs((sigma_hat[(train_end+1):n] - sigma_realized[(train_end+1):n]) / 
    #                                          sigma_realized[(train_end+1):n]), na.rm=TRUE) * 100 else NA,
    #RMSE_test = if (train_end < n) sqrt(mean((sigma_hat[(train_end+1):n] - sigma_realized[(train_end+1):n])^2, na.rm=TRUE)) else NA,
    loglik = -optim_result$value,
    converged = converged_outer,
    #train_end = train_end,
    n_burn = n_burn
  )
}


