#' RTMB model object
#'
#' @name RTMB_Model-class
#' @rdname RTMB_Model-class
#'
#' @description
#' An R6 class representing a Bayesian model built with RTMB.
#' This class stores model components and provides methods for
#' building the automatic differentiation object, optimizing the
#' posterior, and drawing posterior samples.
#'
#'
#' @field data Named list containing observation data.
#' @field par_list Evaluated list of parameter definitions.
#' @field log_prob Executable function for likelihood and prior calculations.
#' @field transform Function for calculating transformed parameters.
#' @field generate Function for calculating generated quantities.
#' @field par_names List of variable names corresponding to parameter dimensions.
#' @field pl_full Complete parameter list with expanded constraints and dimension information used internally.
#' @field formula Formula used for model generation (if created via a wrapper function).
#' @field requested_contrasts Character; user-requested contrast setting (e.g. "treatment").
#' @field contrasts Character; internal contrast setting used for fitting (e.g. "sum").
#' @field raw_data Data frame; the original data used for the model.
#' @field family Character string of the distribution family (if created via a wrapper function).
#' @field init List or numeric vector of initial values.
#' @field view Character vector of variable names to prioritize in summary.
#' @field code Original model AST generated by `rtmb_code()`.
#' @field map A list specifying parameter mappings for fixing values (e.g., via \code{fixed_model()}).
#' @field prior_correction Numeric value for legacy or low-level marginal-likelihood correction.
#' Fixed-parameter prior removal is handled during AD objective construction.
#' @field fixed_prior_specs Internal specification for removing fixed-parameter prior contributions during AD evaluation.
#' @field type Character; the type of the model (e.g. 'lmer', 'glm', 'table', 'ttest', 'corr', 'mediation').
#' @field extra List; used to store auxiliary data (like term assignments, factors, test results) needed by methods such as classic().
#'
#' @import RTMB
#' @export
RTMB_Model <- R6::R6Class(
  classname = "RTMB_Model",
  public = list(
    data = NULL,
    par_list = NULL,
    log_prob = NULL,
    transform = NULL,
    generate = NULL,
    par_names = NULL,
    pl_full = NULL,
    formula = NULL,
    raw_data = NULL,
    family = NULL,
    init = NULL,
    view = NULL,
    code = NULL,
    map = NULL,
    type = NULL,
    extra = NULL,
    contrasts = NULL,
    requested_contrasts = NULL,
    prior_correction = 0,
    fixed_prior_specs = NULL,

    # 1. Constructor
    #' @description Create a new `RTMB_Model` object.
    #' @param data Named list containing observation data.
    #' @param par_list Evaluated list of parameter definitions.
    #' @param log_prob Executable function for likelihood and prior calculations.
    #' @param transform Function for calculating transformed parameters (optional).
    #' @param generate Function for calculating generated quantities (optional).
    #' @param par_names List of variable names corresponding to parameter dimensions.
    #' @param init List or numeric vector of initial values.
    #' @param view Character vector of variable names to prioritize in summary.
    #' @param code Original model AST generated by `rtmb_code()`.
    initialize = function(data, par_list, log_prob,
                          transform = NULL, generate = NULL, par_names = NULL,
                          init = NULL, view = NULL, code = NULL) {
      self$data <- data
      self$par_names <- par_names
      self$init <- init
      self$view <- view
      self$code <- code
      self$par_list <- lapply(par_list, function(x) {
        if (typeof(x) %in% c("language", "symbol")) {
          eval(x, envir = self$data, enclos = parent.frame())
        } else {
          x
        }
      })
      self$log_prob <- log_prob
      self$transform <- transform
      self$generate <- generate

      self$pl_full <- parse_parameters(self$par_list, self$par_names)
      names(self$pl_full$init) <- self$pl_full$names

      # Changed to go through prepare_init
      init_vec <- self$prepare_init(self$init)
      self$init <- init_vec
      test_para <- constrained_vector_to_list(init_vec, self$par_list)

      test_val <- tryCatch(
        {
          if (!is.null(self$transform)) {
            test_tran <- self$transform(self$data, test_para)
            test_para <- c(test_para, test_tran)
          }
          self$log_prob(self$data, test_para)
        },
        error = function(e) {
          stop("An error occurred during the test execution of the R function (log_prob).\n[Error]: ",
            e$message,
            "\n  * Possible cause: Inappropriate initial parameter values may lead to uncomputable likelihoods (Inf or NaN).",
            "\n  * Solution: Set appropriate initial values within the domain (e.g., >0 for variance) via the 'init' argument, or reconsider the data scaling.",
            call. = FALSE
          )
        }
      )

      if (!is.numeric(test_val) || length(test_val) != 1) {
        stop("log_prob must return a single numeric value (scalar).")
      }

      if (!is.null(self$transform)) {
        test_tran <- tryCatch(
          {
            self$transform(self$data, test_para)
          },
          error = function(e) {
            stop("An error occurred during the test execution of the R function (transform).\n[Error]: ",
              e$message,
              call. = FALSE
            )
          }
        )

        if (!is.null(test_tran)) {
          if (!is.list(test_tran)) {
            stop("transform must return a list.", call. = FALSE)
          }
          if (length(test_tran) > 0 && (is.null(names(test_tran)) || any(names(test_tran) == ""))) {
            stop("transform must return a named list.", call. = FALSE)
          }
        }
      }

      if (!getOption("BayesRTMB.silent", FALSE)) cat("Checking RTMB setup...\n")
      test_ad <- self$build_ad_obj(init = init_vec, laplace = FALSE, jacobian_target = "all")
      test_gr <- tryCatch(test_ad$ad_obj$gr(test_ad$ad_obj$par), error = function(e) e)
      if (inherits(test_gr, "error")) {
        stop("An error occurred during gradient calculation in MakeADFun.\n[Error]: ",
          test_gr$message,
          call. = FALSE
        )
      }
    },

    # Helper method to format and complement initial values (uses self$init if no arguments are provided)
    #' @description Prepare and format initial values for the model parameters.
    #' @param init_arg Optional list or numeric vector of initial values. If NULL, defaults to `self$init` or random generation.
    #' @return A flat numeric vector of constrained initial values.
    prepare_init = function(init_arg) {
      target_init <- if (!is.null(init_arg)) init_arg else self$init

      if (is.null(target_init)) {
        return(generate_random_init(self$pl_full, self$par_list, range = 2))
      }

      # When a named list is provided (partial initial value specification)
      if (is.list(target_init)) {
        # 1. First, initialize everything with random values and convert to list format
        base_vec <- generate_random_init(self$pl_full, self$par_list, range = 2)
        base_list <- constrained_vector_to_list(base_vec, self$par_list)

        # 2. Overwrite only the parameters specified by the user
        for (name in names(target_init)) {
          if (name %in% names(base_list)) {
            base_list[[name]] <- as.numeric(target_init[[name]])
          } else {
            warning(sprintf("The initial value specified for '%s' does not exist in the model.", name), call. = FALSE)
          }
        }
        # 3. Return to a complete flat vector and return
        return(unlist(base_list, use.names = FALSE))
      }

      # When a vector is provided (full specification)
      if (is.numeric(target_init)) {
        return(target_init)
      }

      stop("init must be specified as a named list or a numeric vector.")
    },

    #' @description Get the current parameters as a named list.
    #' @param init Optional numeric vector or list of initial values. If NULL, uses the model's current initial values.
    #' @return A named list of constrained parameters.
    get_par_list = function(init = NULL) {
      init_vec <- self$prepare_init(init)
      return(constrained_vector_to_list(init_vec, self$par_list))
    },

    #' @description Get the RTMB automatic differentiation object.
    #' @param ... Arguments passed to \code{build_ad_obj}.
    #' @return A TMB objective object (ad_obj).
    get_ad_obj = function(...) {
      return(self$build_ad_obj(...)$ad_obj)
    },

    #' @description Convert a list of constrained parameters to unconstrained space.
    #' @param par_list_constrained A named list of parameters in their natural space.
    #' @return A named list of parameters in unconstrained space.
    to_unconstrained = function(par_list_constrained) {
      return(to_unconstrained(par_list_constrained, self$par_list))
    },

    #' @description Convert a list of unconstrained parameters to constrained (natural) space.
    #' @param par_list_unconstrained A named list of parameters in unconstrained space.
    #' @return A named list of parameters in constrained space.
    to_constrained = function(par_list_unconstrained) {
      return(to_constrained(par_list_unconstrained, self$par_list))
    },

    #' @description Convert a flat unconstrained vector to a named list.
    #' @param vec A numeric vector in unconstrained space.
    #' @return A named list of unconstrained parameters.
    unconstrained_vector_to_list = function(vec) {
      return(unconstrained_vector_to_list(vec, self$par_list))
    },

    #' @description Convert a flat constrained vector to a named list.
    #' @param vec A numeric vector in constrained space.
    #' @return A named list of constrained parameters.
    constrained_vector_to_list = function(vec) {
      return(constrained_vector_to_list(vec, self$par_list))
    },

    #' @description Calculate Satterthwaite degrees of freedom for parameters.
    #' @param ad_obj An RTMB objective object at the optimum.
    #' @param idx_fix_active Integer vector; indices of active fixed parameters in the full unconstrained vector.
    #' @param L_u_total Integer; total length of the full unconstrained parameter vector.
    #' @param opt_par Optional numeric vector of optimized parameters.
    #' @param max_df Numeric; maximum allowed degrees of freedom. Default is NULL.
    #' @param silent Logical; whether to suppress informational messages. Default is FALSE.
    #' @param return_sensitivities Logical; if TRUE, returns Hessian sensitivities and V for delta method.
    #' @return A numeric vector of estimated degrees of freedom (length = L_u_total). Inf for random effects.
    calculate_satterthwaite_df = function(ad_obj, idx_fix_active = NULL, L_u_total = NULL, opt_par = NULL, max_df = NULL, silent = FALSE, return_sensitivities = FALSE) {
      .calculate_satterthwaite_df_impl(self, private, ad_obj, idx_fix_active, L_u_total, opt_par, max_df, silent, return_sensitivities)
    },

    #' @description Calculate Satterthwaite degrees of freedom for integrated fixed effects (REML).
    #' @param ad_obj An RTMB objective object.
    #' @param opt_par Numeric vector of optimized variance components.
    #' @param beta_idx Integer vector; indices of fixed effects within the random effects vector.
    #' @param max_df Numeric; maximum allowed degrees of freedom. Default is NULL.
    #' @param silent Logical; whether to suppress informational messages. Default is FALSE.
    #' @param return_sensitivities Logical; if TRUE, returns Hessian sensitivities and delta-method components.
    #' @param ... Additional arguments.
    #' @return A numeric vector of estimated degrees of freedom for the fixed effects.
    calculate_reml_satterthwaite_df = function(ad_obj, opt_par, beta_idx, max_df = NULL, silent = FALSE, return_sensitivities = FALSE) {
      .calculate_reml_satterthwaite_df_impl(
        self = self,
        private = private,
        ad_obj = ad_obj,
        opt_par = opt_par,
        beta_idx = beta_idx,
        max_df = max_df,
        silent = silent,
        return_sensitivities = return_sensitivities
      )
    },

    #' @description Build the RTMB automatic differentiation object.
    #' @param init Optional numeric vector or list of initial values for the parameters. Default is NULL.
    #' @param laplace Logical; whether to use Laplace approximation to marginalize random effects. Default is FALSE.
    #' @param jacobian_target Character string specifying which parameters to apply Jacobian adjustments to:
    #' \itemize{
    #'   \item \code{"all"}: Apply to all parameters (standard for Bayesian MCMC).
    #'   \item \code{"random"}: Apply only to parameters with \code{random = TRUE} (standard for Laplace approximation/REML).
    #'   \item \code{"none"}: No Jacobian adjustment (standard for Maximum Likelihood estimation).
    #' }
    #' Default is "none".
    #' @param map Optional list specifying parameters to fix. Passed directly to MakeADFun. Default is NULL.
    #' @param .marginal_vars Internal use for profiling marginalized parameters.
    #' @return An RTMB objective object (ad_obj).
    build_ad_obj = function(init = NULL, laplace = FALSE, jacobian_target = "none", map = NULL, .marginal_vars = NULL) {
      target_map <- if (is.null(map)) {
        self$map
      } else {
        utils::modifyList(as.list(self$map), as.list(map))
      }

      random_effs <- if (!is.null(.marginal_vars)) {
        .marginal_vars
      } else {
        names(self$par_list)[sapply(self$par_list, function(x) isTRUE(x$random))]
      }

      # Exclude random components where all elements are NA in fixed/map
      if (!is.null(target_map)) {
        random_effs <- random_effs[vapply(random_effs, function(nm) {
          m <- target_map[[nm]]
          if (is.null(m)) return(self$par_list[[nm]]$unc_length > 0)
          any(!is.na(m))
        }, logical(1))]
        
        target_map <- lapply(target_map, private$.normalize_map_factor)
      }

      use_random <- if (laplace && length(random_effs) > 0) random_effs else NULL

      current_init <- self$prepare_init(init)
      current_init_list <- constrained_vector_to_list(current_init, self$par_list)
      init_unc_list <- to_unconstrained(current_init_list, self$par_list)

      pl_full_local <- self$pl_full
      par_list_local <- self$par_list
      log_prob_local <- self$log_prob
      data_local <- self$data

      # Pre-calculate shadow environment for fixed-prior removal
      # This identifies the full names of all parameters and transformations
      env_shadow <- list2env(self$data, parent = asNamespace("BayesRTMB"))
      idx <- 1
      for (nm in names(self$par_list)) {
        p <- self$par_list[[nm]]
        shadow_vec <- self$pl_full$names[idx:(idx + p$length - 1)]
        if (length(p$dim) > 1) dim(shadow_vec) <- p$dim
        env_shadow[[nm]] <- shadow_vec
        idx <- idx + p$length
      }
      if (!is.null(self$transform)) {
        par_only <- mget(names(self$par_list), envir = env_shadow, inherits = FALSE)
        tran_res_shadow <- tryCatch(suppressMessages(self$transform(self$data, par_only)), error = function(e) list())
        for (nm in names(tran_res_shadow)) env_shadow[[nm]] <- tran_res_shadow[[nm]]
      }
      shadow_list_local <- as.list(env_shadow)

      f_ad <- private$.build_f_ad(
        data_local = data_local,
        par_list_local = par_list_local,
        log_prob_local = log_prob_local,
        transform_local = self$transform,
        jacobian_target = jacobian_target,
        adreport = TRUE,
        fixed_prior_specs_local = self$fixed_prior_specs,
        code_model_local = self$code$model,
        shadow_list_local = shadow_list_local
      )

      ad_obj <- tryCatch(
        {
          RTMB::MakeADFun(
            func = f_ad,
            parameters = init_unc_list,
            random = use_random,
            map = target_map,
            silent = TRUE
          )
        },
        error = function(e) {
          stop("Failed to setup MakeADFun.\n[Error]: ", e$message, call. = FALSE)
        }
      )

      return(list(
        ad_obj = ad_obj,
        pl_full = pl_full_local,
        use_random = use_random
      ))
    },

    # 3. Maximum Likelihood / MAP Estimation Method
    #' @description Perform Maximum Likelihood or Maximum A Posteriori (MAP) estimation.
    #' @param laplace Logical; whether to use Laplace approximation for random effects. Default is TRUE.
    #' @param init Optional initial values for parameters (numeric vector or list).
    #' @param num_estimate Integer; Number of multi-start optimization runs. Default is 1.
    #' @param control A list of control settings passed to the optimizer.
    #' @param optimizer Character; The optimizer to use: "nlminb" (default) or "optim".
    #' @param method Character; The method for "optim" (e.g., "BFGS", "L-BFGS-B").
    #' @param map Optional list specifying parameters to fix (factors).
    #' @param fixed Optional list specifying parameter values to fix.
    #' @param se_method Character; method for uncertainty estimation.
    #' \itemize{
    #'   \item \code{"wald"}: Wald standard errors and Wald confidence intervals.
    #'   \item \code{"sampling"}: Simulation-based error propagation.
    #'   \item \code{"none"}: Do not compute standard errors, confidence intervals, or degrees of freedom.
    #' }
    #' @param num_samples Integer; number of samples to draw when se_method is "sampling". Default is 1000.
    #' @param seed Integer; random seed for sampling.
    #' @param marginal Character vector or "auto" or "none"; specifies which parameters to marginalize out via Laplace approximation.
    #' \itemize{
    #'   \item \code{"none"} (default): No additional marginalization.
    #'   \item \code{"auto"}: Uses variables specified by the model wrapper (e.g., fixed effects in lmer).
    #'   \item \code{"all"}: Marginalize all parameters (advanced use).
    #'   \item \code{character vector}: Specific parameter names to marginalize.
    #' }
    #' @param df_method Character or numeric; method for finite degrees-of-freedom approximation.
    #' Used only when `marginal` resolves to one or more parameters and `se_method != "none"`.
    #' \itemize{
    #'   \item \code{"auto"}: Uses Satterthwaite approximation when marginalization is active.
    #'   \item \code{"satterthwaite"}: Satterthwaite approximation.
    #'   \item \code{"inf"} or \code{"none"}: Use infinite degrees of freedom.
    #'   \item numeric: Use the supplied fixed degrees of freedom.
    #' }
    #' @param view Character vector of parameter names to prioritize in summary display.
    #' @param ... Additional arguments.
    #' @return A fitted `MAP_Fit` object.
    optimize = function(laplace = TRUE, init = NULL, num_estimate = 1, control = list(),
                        optimizer = "nlminb", method = "BFGS", map = NULL, fixed = NULL,
                        se_method = c("wald", "sampling", "none"), 
                        num_samples = 1000, seed = 123,
                        marginal = "none", df_method = "auto", view = NULL, ...) {
      
      se_method <- match.arg(se_method)
      se_flag <- !identical(se_method, "none")

      # --- Strict check for deprecated arguments ---
      dot_args <- list(...)
      reserved_old_args <- c(
        "empirical", "REML", "df", "df_pars",
        "target_vars", "is_classic", ".return_object",
        "ci_method", "se_sampling", "views", "se"
      )
      bad_args <- intersect(names(dot_args), reserved_old_args)
      if (length(bad_args) > 0L) {
        stop(
          "Unsupported argument(s) in optimize(): ",
          paste(bad_args, collapse = ", "), "\n",
          "Use `se_method = 'none'`, `se_method = 'wald'`, or `se_method = 'sampling'`.",
          call. = FALSE
        )
      }

      if (isTRUE(self$extra$two_stage) && identical(self$type, "lrt")) {
        return(.lrt_two_stage_optimize(
          self = self,
          laplace = laplace,
          init = init,
          num_estimate = num_estimate,
          control = control,
          optimizer = optimizer,
          method = method,
          map = map,
          fixed = fixed,
          se_method = se_method,
          num_samples = num_samples,
          seed = seed,
          marginal = marginal,
          df_method = df_method,
          view = view,
          ...
        ))
      }

      .inference_pipeline(
        self = self,
        private = private,
        laplace = laplace,
        init = init,
        num_estimate = num_estimate,
        control = control,
        optimizer = optimizer,
        method = method,
        map = map,
        fixed = fixed,
        se = se_flag,
        se_method = se_method,
        num_samples = num_samples,
        seed = seed,
        marginal = marginal,
        df_method = df_method,
        view = view,
        ...
      )
    },


    #' @description Perform frequentist inference (REML/ML) with model-appropriate degrees of freedom.
    #' @param df_method Character; Method for calculating degrees of freedom. Default is "auto".
    #' \itemize{
    #'   \item \code{"auto"}: Determined based on model type (e.g., Satterthwaite for mixed models).
    #'   \item \code{"residual"}: Residual degrees of freedom.
    #'   \item \code{"satterthwaite"}: Satterthwaite approximation (robust for mixed models).
    #'   \item \code{"inf"}: Use z-distribution (infinite degrees of freedom).
    #' }
    #' @param se_method Character; The method for CI estimation: "wald" (default) or "sampling".
    #' @param num_samples Integer; number of samples to draw when se_method is "sampling". Default is 1000.
    #' @param seed Integer; random seed for sampling.
    #' @param view Character vector of parameters to prioritize in the summary output.
    #' @param map Optional list specifying parameters to fix.
    #' @param fixed Optional list specifying parameter values to fix.
    #' @param ... Additional arguments.
    #' @return A `Classic_Fit` object.
    classic = function(df_method = "auto", se_method = c("wald", "sampling"),
                       num_samples = 1000, seed = 123,
                       view = NULL, map = NULL, fixed = NULL, ...) {
      
      # --- Check for unsupported/deprecated arguments ---
      dot_args <- list(...)
      reserved_classic_args <- c(
        "df", "REML", "marginal", "df_pars", "views",
        "empirical", "target_vars", "is_classic", ".return_object",
        "ci_method", "se_sampling",
        "robust", "robust_se",
        "bootstrap", "n_boot", "B",
        "cluster", "vcov", "se_type"
      )
      bad_args <- intersect(names(dot_args), reserved_classic_args)
      if (length(bad_args) > 0L) {
        stop(
          "Unsupported argument(s) in classic(): ",
          paste(bad_args, collapse = ", "), "\n",
          "Use fit$robust_se() or fit$bootstrap() after classic().",
          call. = FALSE
        )
      }

      private$.classic_pipeline(
        df_method = df_method,
        se_method = se_method,
        num_samples = num_samples,
        seed = seed,
        view = view,
        map = map,
        fixed = fixed,
        ...
      )
    },

    # 4. MCMC NUTS Sampling Method
    #' @description Draw posterior samples from the model.
    #' @param sampling Number of sampling iterations. Default is 1000.
    #' @param warmup Number of warmup iterations. Default is 1000.
    #' @param chains Number of MCMC chains. Default is 4.
    #' @param thin Thinning interval. Default is 1.
    #' @param seed Random seed.
    #' @param delta Target acceptance rate for HMC/NUTS. Default is 0.8.
    #' @param max_treedepth Maximum tree depth for HMC/NUTS. Default is 10.
    #' @param parallel Logical; whether to run chains in parallel. Default is FALSE.
    #' @param laplace Logical; whether to use Laplace approximation. Default is FALSE.
    #' @param init Optional initial values for parameters.
    #' @param init_jitter sd of randomize initial values for parameters.
    #' @param save_csv Optional list for saving MCMC results. e.g., list(name = "model", dir = "BayesRTMB_mcmc").
    #' @param map Optional list specifying parameters to fix (factors).
    #' @param fixed Optional list specifying parameter values to fix.
    #' @return A fitted `MCMC_Fit` object.
    sample = function(sampling = 1000, warmup = 1000, chains = 4,
                      thin = 1, seed = sample.int(1e6, 1),
                      delta = 0.8, max_treedepth = 10,
                      parallel = FALSE, laplace = FALSE,
                      init = NULL, init_jitter = 0.1, save_csv = NULL,
                      map = NULL, fixed = NULL) {
      if (isTRUE(self$extra$two_stage) && identical(self$type, "lrt")) {
        stop(
          "two_stage estimation for rtmb_lrt() is currently implemented only for optimize(). ",
          "MCMC sampling is not implemented for two_stage models yet.",
          call. = FALSE
        )
      }
      .sample_impl(self, private, sampling, warmup, chains, thin, seed, delta, 
                   max_treedepth, parallel, laplace, init, init_jitter, 
                   save_csv, map, fixed)
    },

    #' @description Run Automatic Differentiation Variational Inference (ADVI).
    #' @param iter Integer; fixed number of iterations for the optimization. Default is 3000.
    #' @param tol_rel_obj Numeric; relative tolerance for the ELBO change to determine convergence. Default is 0.005.
    #' @param window_size Integer; window size for median smoothing in the convergence check. Default is 100.
    #' @param num_samples Integer; number of posterior samples to generate from the fitted variational distribution. Default is 1000.
    #' @param num_estimate Integer; number of times to run the VB estimation (treated as chains). Default is 4.
    #' @param alpha Numeric; learning rate for the Adam optimizer. Default is 0.01.
    #' @param laplace Logical; whether to use Laplace approximation to marginalize random effects. Default is FALSE.
    #' @param print_freq Integer; iterations interval for progress output. Set to 0 to disable. Default is 1000.
    #' @param method Character; method of Variational Inference. Default is "meanfield".
    #' @param parallel Logical; whether to run estimations in parallel. Default is FALSE.
    #' @param seed Integer; random seed for reproducibility.
    #' @param init Optional numeric vector or list for initial parameter values. Default is NULL.
    #' @param save_csv Optional list for saving VB results. e.g., list(name = "model", dir = "BayesRTMB_vb").
    #' @param map Optional list specifying parameters to fix (factors).
    #' @param fixed Optional list specifying parameter values to fix.
    #' @return A fitted `VB_Fit` object containing posterior samples and diagnostic information.
    variational = function(iter = 3000,
                           tol_rel_obj = 0.005, window_size = 100,
                           num_samples = 1000, num_estimate = 4, alpha = 0.01,
                           laplace = FALSE, print_freq = 1000,
                           method = c("meanfield", "fullrank", "hybrid"), parallel = FALSE,
                           seed = sample.int(1e6, 1), init = NULL, save_csv = NULL,
                           map = NULL, fixed = NULL) {
      if (isTRUE(self$extra$two_stage) && identical(self$type, "lrt")) {
        stop(
          "two_stage estimation for rtmb_lrt() is currently implemented only for optimize(). ",
          "Variational inference is not implemented for two_stage models yet.",
          call. = FALSE
        )
      }
      .variational_impl(self, private, iter, tol_rel_obj, window_size, num_samples, 
                        num_estimate, alpha, laplace, print_freq, method, 
                        parallel, seed, init, save_csv, map, fixed)
    },


    # 5. Print model code method
    #' @description Print model code or model structure.
    #' @return The object itself, invisibly.
    print_code = function() {
      .print_code_impl(self, private)
    },

    #' @description Print the internal log-probability function.
    #' @return The object itself, invisibly.
    print_log_prob = function() {
      if (is.null(self$log_prob)) {
        cat("log_prob is not available.\n")
        return(invisible(self))
      }

      # Extract the internal function from the error-handling wrapper if it exists
      fn <- self$log_prob
      if (exists("comp_model", envir = environment(fn))) {
        fn <- get("comp_model", envir = environment(fn))
      }

      cat("=== RTMB Model: log_prob (Internal Function) ===\n")
      print(fn)

      if (!is.null(self$fixed_prior_specs) && length(self$fixed_prior_specs) > 0) {
        cat("\n=== Fixed-prior removal ===\n")
        cat("The following fixed parameters have their sampling-statement priors removed during AD evaluation:\n")
        cat(paste0("- ", names(self$fixed_prior_specs), collapse = "\n"), "\n")
      }

      return(invisible(self))
    },

    #' @description Print the internal transformation function.
    #' @return The object itself, invisibly.
    print_transform = function() {
      if (is.null(self$transform)) {
        cat("transform is not available.\n")
        return(invisible(self))
      }

      # Extract the internal function from the error-handling wrapper if it exists
      fn <- self$transform
      if (exists("comp_transform", envir = environment(fn))) {
        fn <- get("comp_transform", envir = environment(fn))
      }

      cat("=== RTMB Model: transform (Internal Function) ===\n")
      print(fn)
      return(invisible(self))
    },

    #' @description Print the internal generation function.
    #' @return The object itself, invisibly.
    print_generate = function() {
      if (is.null(self$generate)) {
        cat("generate is not available.\n")
        return(invisible(self))
      }

      # Extract the internal function from the error-handling wrapper if it exists
      fn <- self$generate
      if (exists("comp_generate", envir = environment(fn))) {
        fn <- get("comp_generate", envir = environment(fn))
      }

      cat("=== RTMB Model: generate (Internal Function) ===\n")
      print(fn)
      return(invisible(self))
    },

    #' @description Automatically detect and count the number of independent data points (observations).
    #' @return Integer; total number of observations, or NULL if model code is missing.
    get_n_obs = function() {
      .get_n_obs_impl(self, private)
    },

    #' @description Create a model with fixed parameters.
    #' @param fixed A named list of parameter values to fix.
    #' @param silent Logical; whether to suppress informational messages. Default is FALSE.
    #' @return A new RTMB_Model object.
    fixed_model = function(fixed = list(), silent = FALSE) {
      return(private$.apply_constraints(fixed, silent = silent))
    }
  ),

  private = list(
    # --- Classic pipeline implementation ---
    .validate_classic_model = function() {
      .validate_classic_model(self, private)
    },
    .resolve_classic_settings = function(df_method = "auto") {
      .resolve_classic_settings(self, private, df_method)
    },
    .classic_pipeline = function(...) {
      .classic_pipeline(self, private, ...)
    },
    .build_classic_components = function(...) {
      .build_classic_components(self, private, ...)
    },
    .get_unc_indices = function(name) {
      idx_curr <- 1L
      for (nm in names(self$par_list)) {
        L_u <- self$par_list[[nm]]$unc_length
        if (nm == name) {
          if (L_u == 0) return(integer(0))
          return(idx_curr:(idx_curr + L_u - 1L))
        }
        idx_curr <- idx_curr + L_u
      }
      integer(0)
    },

    # --- Internal helper: merge self$map and user supplied map consistently ---
    .merge_map_arg = function(map = NULL) {
      if (is.null(self$map) && is.null(map)) return(NULL)
      if (is.null(self$map)) return(as.list(map))
      if (is.null(map)) return(self$map)
      utils::modifyList(as.list(self$map), as.list(map))
    },

    # --- Internal helper: ensure map factor levels follow first-seen order ---
    .normalize_map_factor = function(m) {
      if (is.null(m)) return(NULL)
      vals <- as.character(m)
      factor(vals, levels = unique(vals[!is.na(vals)]))
    },

    # --- Internal helper: compute full unconstrained indices for active fixed and random parameters ---
    .compute_rtmb_index_map = function(ad_obj, target_map = NULL) {
      idx_ran <- ad_obj$env$random
      if (is.null(idx_ran)) idx_ran <- integer(0)

      idx_fix_active <- integer(0)
      idx_ran_full <- integer(0)
      factor_levels_seen <- list()

      opt_par_curr <- 1L
      idx_mapped_curr <- 1L
      idx_full_curr <- 1L

      for (name in names(self$par_list)) {
        L_u <- self$par_list[[name]]$unc_length
        if (L_u > 0) {
          map_f <- if (!is.null(target_map[[name]])) target_map[[name]] else NULL
          for (i in seq_len(L_u)) {
            pos_full <- idx_full_curr + i - 1L
            if (!is.null(map_f) && is.na(map_f[i])) next

            if (idx_mapped_curr %in% idx_ran) {
              idx_ran_full <- c(idx_ran_full, pos_full)
              idx_mapped_curr <- idx_mapped_curr + 1L
              next
            }

            if (!is.null(map_f)) {
              lvl <- as.character(map_f[i])
              if (!(lvl %in% names(factor_levels_seen))) {
                factor_levels_seen[[lvl]] <- opt_par_curr
                idx_fix_active <- c(idx_fix_active, pos_full)
                opt_par_curr <- opt_par_curr + 1L
              }
            } else {
              idx_fix_active <- c(idx_fix_active, pos_full)
              opt_par_curr <- opt_par_curr + 1L
            }
            idx_mapped_curr <- idx_mapped_curr + 1L
          }
          idx_full_curr <- idx_full_curr + L_u
        }
      }

      list(idx_fix_active = idx_fix_active, idx_ran_full = idx_ran_full, 
           idx_ran = idx_ran, factor_levels_seen = factor_levels_seen)
    },

    # --- Internal helper: copy SEs to duplicated mapped parameters ---
    .copy_mapped_se = function(unc_se_vec, se_fix, target_map, factor_levels_seen, idx_ran_full) {
      idx_curr <- 1L
      for (name in names(self$par_list)) {
        L_u <- self$par_list[[name]]$unc_length
        if (L_u > 0) {
          map_f <- if (!is.null(target_map[[name]])) target_map[[name]] else NULL
          for (i in seq_len(L_u)) {
            pos_full <- idx_curr + i - 1L
            if (pos_full %in% idx_ran_full) next
            if (!is.null(map_f) && !is.na(map_f[i])) {
              lvl <- as.character(map_f[i])
              mapped_opt_idx <- factor_levels_seen[[lvl]]
              if (!is.null(mapped_opt_idx) && mapped_opt_idx <= length(se_fix)) {
                unc_se_vec[pos_full] <- se_fix[mapped_opt_idx]
              }
            }
          }
          idx_curr <- idx_curr + L_u
        }
      }
      unc_se_vec
    },

    # --- Internal helper: resolve marginalization targets ---
    # --- Internal helper: resolve marginalization targets ---
    .resolve_marginal_vars = function(marginal = "none") {
      if (is.null(marginal)) marginal <- "none"
      if (!is.character(marginal)) stop("`marginal` must be 'none', 'auto', 'all', or a character vector of parameter names.", call. = FALSE)
      if (length(marginal) == 0L) return(character(0))
      if (length(marginal) == 1L && identical(marginal, "none")) return(character(0))

      # 1. Resolve keywords to parameter names
      if (length(marginal) == 1L && identical(marginal, "auto")) {
        target_vars <- self$extra$marginal
        if (is.null(target_vars) || length(target_vars) == 0L) return(character(0))
        if (!is.character(target_vars)) stop("Wrapper-defined marginal variables must be a character vector.", call. = FALSE)
        
        missing <- setdiff(target_vars, names(self$par_list))
        if (length(missing) > 0L) stop("Wrapper-defined marginal variables are not model parameters: ", paste(missing, collapse = ", "), call. = FALSE)
        target_vars <- unique(target_vars)
      } else if (length(marginal) == 1L && identical(marginal, "all")) {
        target_vars <- names(self$par_list)
      } else {
        reserved <- c("none", "auto", "all")
        if (any(marginal %in% reserved)) stop("`marginal` must be exactly one of 'none', 'auto', 'all', or a character vector of parameter names.", call. = FALSE)
        target_vars <- unique(marginal)
        missing <- setdiff(target_vars, names(self$par_list))
        if (length(missing) > 0L) stop("The following variables specified in `marginal` are not model parameters: ", paste(missing, collapse = ", "), call. = FALSE)
      }

      # 2. Prune fully fixed parameters
      if (!is.null(self$map) && length(target_vars) > 0L) {
        if (identical(marginal, "auto")) {
          # Silently remove fully fixed parameters in 'auto' mode
          target_vars <- target_vars[vapply(target_vars, function(nm) {
            m <- self$map[[nm]]
            is.null(m) || any(!is.na(m))
          }, logical(1))]
        } else if (!identical(marginal, "all")) {
          # Error if explicitly requested parameters are fully fixed
          fully_fixed_targets <- vapply(target_vars, function(nm) {
            m <- self$map[[nm]]
            !is.null(m) && all(is.na(m))
          }, logical(1))
          if (any(fully_fixed_targets)) {
            stop(sprintf("`marginal` contains parameters that are fixed and therefore cannot be marginalized: %s", 
                 paste(target_vars[fully_fixed_targets], collapse = ", ")), call. = FALSE)
          }
        }
      }

      return(target_vars)
    },

    # --- Internal helper: temporary marginal classification by changing random flags ---
    .with_marginal_par_list = function(target_vars = NULL, map = NULL) {
      original_par_list <- self$par_list
      modified_par_list <- self$par_list
      target_map <- private$.merge_map_arg(map)

      for (name in names(modified_par_list)) {
        is_mapped_out <- !is.null(target_map[[name]]) && all(is.na(target_map[[name]]))
        is_target <- !is.null(target_vars) && name %in% target_vars
        is_random_original <- isTRUE(modified_par_list[[name]]$random)

        if (!is_mapped_out && (is_target || is_random_original)) {
          modified_par_list[[name]]$random <- TRUE
        } else {
          modified_par_list[[name]]$random <- FALSE
        }
      }

      list(original = original_par_list, modified = modified_par_list)
    },

    # --- Low-level RTMB fitting engine shared by optimize() and classic() ---
    .fit_rtmb = function(laplace = TRUE, init = NULL, num_estimate = 1, control = list(),
                         optimizer = "nlminb", method = "BFGS", map = NULL, se = TRUE,
                         reml = FALSE, target_vars = NULL, jacobian_target = NULL,
                         apply_prior_correction = TRUE,
                         verbose = !getOption("BayesRTMB.silent", FALSE)) {

      if (is.null(control)) control <- list()
      optimizer <- match.arg(optimizer, c("nlminb", "optim"))
      if (!is.numeric(num_estimate) || length(num_estimate) != 1L || num_estimate < 1) stop("num_estimate must be a positive integer.", call. = FALSE)
      se_messages <- character(0)
      add_se_message <- function(msg) {
        if (!msg %in% se_messages) se_messages <<- c(se_messages, msg)
      }
      data_na_summary <- function(dat) {
        if (is.null(dat)) return(character(0))
        out <- character(0)
        for (nm in names(dat)) {
          x <- dat[[nm]]
          if (is.atomic(x) || is.matrix(x) || is.array(x) || is.data.frame(x)) {
            n_na <- sum(is.na(x))
            if (n_na > 0L) out <- c(out, sprintf("%s (%d NA)", nm, n_na))
          }
        }
        out
      }
      optimizer_error_messages <- character(0)
      format_optimizer_error <- function(msg) {
        na_vars <- data_na_summary(self$data)
        if (grepl("NA/NaN|NaN|NA", msg) && length(na_vars) > 0L) {
          paste0(
            "Optimization error: ", msg,
            ". The model data contain missing values in: ",
            paste(na_vars, collapse = ", "),
            ". If the model does not handle missing values explicitly, remove/impute them or use a wrapper/missing-data option that supports NA."
          )
        } else {
          paste0("Optimization error: ", msg)
        }
      }

      reml_state <- NULL
      if (isTRUE(reml)) {
        reml_state <- private$.with_marginal_par_list(target_vars = target_vars, map = map)
        self$par_list <- reml_state$modified
        on.exit({ self$par_list <- reml_state$original }, add = TRUE)
        laplace <- TRUE
      }

      # Decide jacobian target after potential REML -> laplace conversion
      if (is.null(jacobian_target)) jacobian_target <- if (isTRUE(laplace)) "random" else "none"
      jacobian_target <- match.arg(jacobian_target, c("all", "random", "none"))

      target_map <- private$.merge_map_arg(map)
      if (!is.null(target_map)) {
        target_map <- lapply(target_map, private$.normalize_map_factor)
      }

      if (verbose) cat("Starting RTMB optimization...\n")

      ad_setup <- if (isTRUE(self$silent) || getOption("BayesRTMB.silent", FALSE)) {
        suppressMessages(suppressWarnings(self$build_ad_obj(init = init, laplace = laplace, jacobian_target = jacobian_target, map = map)))
      } else {
        self$build_ad_obj(init = init, laplace = laplace, jacobian_target = jacobian_target, map = map)
      }
      base_ad_obj <- ad_setup$ad_obj

      opt_results <- vector("list", num_estimate); obj_vals <- rep(NA_real_, num_estimate)
      conv_codes <- rep(NA_integer_, num_estimate); messages <- rep(NA_character_, num_estimate)

      for (i in seq_len(num_estimate)) {
        if (num_estimate > 1L && verbose) cat(sprintf("Optimization run %d/%d...\r", i, num_estimate))
        res <- tryCatch({
          if (i > 1L && is.null(init) && length(base_ad_obj$par) > 0L) base_ad_obj$par <- base_ad_obj$par + rnorm(length(base_ad_obj$par), 0, 0.5)
          if (length(base_ad_obj$par) == 0L) {
            opt <- list(par = numeric(0), objective = base_ad_obj$fn(numeric(0)), convergence = 0L, message = "all parameters fixed")
          } else if (optimizer == "nlminb") {
            if (is.null(control$iter.max)) control$iter.max <- 5000
            opt <- suppressWarnings(nlminb(base_ad_obj$par, base_ad_obj$fn, base_ad_obj$gr, control = control))
          } else {
            if (is.null(control$maxit)) control$maxit <- 5000
            opt <- optim(base_ad_obj$par, base_ad_obj$fn, base_ad_obj$gr, method = method, control = control); opt$objective <- opt$value
          }
          list(opt = opt, ad_obj = base_ad_obj)
        }, error = function(e) {
          msg <- conditionMessage(e)
          optimizer_error_messages <<- c(optimizer_error_messages, msg)
          NULL
        })

        if (!is.null(res)) {
          opt_results[[i]] <- res; obj_vals[i] <- res$opt$objective
          conv_codes[i] <- res$opt$convergence; messages[i] <- if (!is.null(res$opt$message)) as.character(res$opt$message) else NA_character_
        }
      }
      if (verbose) cat("\n")
      valid_idx <- which(!is.na(obj_vals))
      if (length(valid_idx) == 0L) {
        na_vars <- data_na_summary(self$data)
        if (length(na_vars) > 0L && any(grepl("NA/NaN|NaN|NA", optimizer_error_messages))) {
          stop(
            "All optimization attempts failed. RTMB reported an NA/NaN-related optimization error, and the model data contain missing values in: ",
            paste(na_vars, collapse = ", "),
            ". If these NA values are not handled explicitly in the model code, remove/impute them or use a wrapper/missing-data option that supports NA.",
            call. = FALSE
          )
        }
        stop("All optimization attempts failed.", call. = FALSE)
      }
      if (length(optimizer_error_messages) > 0L) {
        warning(
          paste(unique(vapply(optimizer_error_messages, format_optimizer_error, character(1L))), collapse = "\n"),
          call. = FALSE
        )
      }
      converged_idx <- valid_idx[!is.na(conv_codes[valid_idx]) & conv_codes[valid_idx] == 0L]
      if (length(converged_idx) > 0L) {
        best_idx <- converged_idx[which.min(obj_vals[converged_idx])]
      } else {
        best_idx <- valid_idx[which.min(obj_vals[valid_idx])]
        best_code <- conv_codes[best_idx]
        best_message <- messages[best_idx]
        detail <- paste0(
          " convergence code = ", if (is.na(best_code)) "NA" else best_code,
          if (!is.na(best_message) && nzchar(best_message)) paste0("; message = ", best_message) else ""
        )
        if (num_estimate > 1L) {
          warning(
            "No optimization run converged; BEST is selected by the lowest objective among non-converged runs (",
            detail,
            ").",
            call. = FALSE
          )
        } else {
          warning(
            "Optimization did not converge (",
            detail,
            "). Estimates may be unreliable; consider increasing num_estimate, changing initial values, or adjusting optimizer control settings.",
            call. = FALSE
          )
        }
      }
      best_res <- opt_results[[best_idx]]

      if (isTRUE(apply_prior_correction) && !is.na(self$prior_correction) && self$prior_correction != 0) {
        for (i in seq_along(obj_vals)) if (!is.na(obj_vals[i])) obj_vals[i] <- obj_vals[i] + self$prior_correction
        best_res$opt$objective <- best_res$opt$objective + self$prior_correction
      }

      ad_obj <- best_res$ad_obj; opt <- best_res$opt; ad_obj$fn(opt$par)

      opt_history <- data.frame(
        estimate = seq_len(num_estimate),
        objective = obj_vals,
        code = conv_codes,
        message = messages,
        stringsAsFactors = FALSE
      )

      if (num_estimate > 1L && isTRUE(verbose)) {
        cat("\nConvergence Diagnostics per estimate:\n")
        for (i in seq_len(num_estimate)) {
          status <- if (!is.na(conv_codes[i]) && conv_codes[i] == 0L) "Converged" else "Not Converged"
          best_marker <- if (i == best_idx) "  <-- BEST" else ""
          cat(sprintf("  est%d: objective = %10.2f (%s)%s\n", i, obj_vals[i], status, best_marker))
        }
      }

      sd_rep <- if (isTRUE(se)) {
        tryCatch(
          withCallingHandlers(
            RTMB::sdreport(ad_obj, getJointPrecision = isTRUE(reml)),
            warning = function(w) {
              add_se_message(paste0("SE warning from sdreport(): ", conditionMessage(w)))
              invokeRestart("muffleWarning")
            }
          ),
          error = function(e) {
            add_se_message(paste0("SE calculation via sdreport() failed: ", conditionMessage(e)))
            NULL
          }
        )
      } else NULL

      unc_est_vec <- unlist(ad_obj$env$parList(), use.names = FALSE); L_u_total <- length(unc_est_vec)
      index_map <- private$.compute_rtmb_index_map(ad_obj, target_map)
      idx_fix_active <- index_map$idx_fix_active; idx_ran_full <- index_map$idx_ran_full
      if (length(opt$par) > 0L && length(idx_fix_active) == length(opt$par)) unc_est_vec[idx_fix_active] <- opt$par

      if (isTRUE(laplace) && !is.null(sd_rep) && length(idx_ran_full) > 0L) {
        smry_ran <- tryCatch(summary(sd_rep, select = "random"), error = function(e) NULL)
        if (!is.null(smry_ran) && nrow(smry_ran) == length(idx_ran_full)) unc_est_vec[idx_ran_full] <- smry_ran[, "Estimate"]
      }

      unc_se_vec <- rep(0, L_u_total); Cov_u <- matrix(0, L_u_total, L_u_total); fallback_needed <- FALSE
      if (isTRUE(se) && !is.null(sd_rep) && !is.null(sd_rep$cov.fixed)) {
        se_fix <- sqrt(pmax(diag(sd_rep$cov.fixed), 0))
        if (!isTRUE(sd_rep$pdHess) || any(!is.finite(se_fix)) || length(idx_fix_active) != length(se_fix)) {
          fallback_needed <- TRUE
          if (!isTRUE(sd_rep$pdHess)) {
            add_se_message("SE warning: sdreport() returned pdHess = FALSE; Hessian-based fallback will be attempted.")
          }
          if (any(!is.finite(se_fix))) {
            add_se_message("SE warning: sdreport() produced non-finite standard errors; Hessian-based fallback will be attempted.")
          }
          if (length(idx_fix_active) != length(se_fix)) {
            add_se_message("SE warning: sdreport() covariance size did not match the active parameter vector; Hessian-based fallback will be attempted.")
          }
        }
        else {
          unc_se_vec[idx_fix_active] <- se_fix; Cov_u[idx_fix_active, idx_fix_active] <- sd_rep$cov.fixed
          unc_se_vec <- private$.copy_mapped_se(unc_se_vec, se_fix, target_map, index_map$factor_levels_seen, idx_ran_full)
        }
      } else if (isTRUE(se)) {
        fallback_needed <- TRUE
        add_se_message("SE warning: sdreport() did not return cov.fixed; Hessian-based fallback will be attempted.")
      }

      if (isTRUE(se) && isTRUE(reml) && !is.null(sd_rep) && !is.null(sd_rep$jointPrecision) && length(idx_ran_full) > 0L) {
        V_joint <- tryCatch(
          solve(as.matrix(sd_rep$jointPrecision)),
          error = function(e) {
            add_se_message("SE warning: joint precision matrix was singular; using MASS::ginv() for random-effect covariance.")
            if (requireNamespace("MASS", quietly = TRUE)) MASS::ginv(as.matrix(sd_rep$jointPrecision)) else NULL
          }
        )
        if (!is.null(V_joint) && nrow(V_joint) == length(idx_ran_full) && ncol(V_joint) == length(idx_ran_full)) {
          Cov_u[idx_ran_full, idx_ran_full] <- as.matrix(V_joint)
          unc_se_vec[idx_ran_full] <- sqrt(pmax(diag(as.matrix(V_joint)), 0))
        }
      }

      if (isTRUE(se) && isTRUE(laplace) && !is.null(sd_rep) && length(idx_ran_full) > 0L) {
        smry_ran <- tryCatch(summary(sd_rep, select = "random"), error = function(e) NULL)
        if (!is.null(smry_ran) && nrow(smry_ran) == length(idx_ran_full) && "Std. Error" %in% colnames(smry_ran)) unc_se_vec[idx_ran_full] <- smry_ran[, "Std. Error"]
      }

      if (isTRUE(se) && isTRUE(fallback_needed) && length(opt$par) > 0L) {
        H <- tryCatch(ad_obj$he(opt$par), error = function(e) tryCatch(ad_obj$he(opt$par + 1e-6), error = function(e2) NULL))
        if (!is.null(H)) {
          Cov_pseudo <- tryCatch(
            solve(H),
            error = function(e) {
              add_se_message("SE warning: Hessian matrix was singular; using MASS::ginv() to approximate the covariance matrix.")
              if (requireNamespace("MASS", quietly = TRUE)) MASS::ginv(H) else NULL
            }
          )
          if (!is.null(Cov_pseudo) && length(idx_fix_active) == length(sqrt(pmax(diag(Cov_pseudo), 0)))) {
            se_pseudo <- sqrt(pmax(diag(Cov_pseudo), 0)); unc_se_vec[idx_fix_active] <- se_pseudo; Cov_u[idx_fix_active, idx_fix_active] <- Cov_pseudo
            unc_se_vec <- private$.copy_mapped_se(unc_se_vec, se_pseudo, target_map, index_map$factor_levels_seen, idx_ran_full)
          } else {
            add_se_message("SE warning: Hessian-based fallback failed; standard errors may be unavailable or unreliable.")
          }
        } else {
          add_se_message("SE warning: Hessian could not be evaluated; standard errors may be unavailable or unreliable.")
        }
      }
      if (isTRUE(se) && length(se_messages) > 0L && isTRUE(verbose) && !getOption("BayesRTMB.silent", FALSE)) {
        cat(paste0(se_messages, collapse = "\n"), "\n", sep = "")
      }

      diag(Cov_u) <- pmax(diag(Cov_u), unc_se_vec^2); Cov_u[is.na(Cov_u)] <- 0
      unc_est_list <- unconstrained_vector_to_list(unc_est_vec, self$par_list)
      unc_se_list <- unconstrained_vector_to_list(unc_se_vec, self$par_list)
      con_est_list <- to_constrained(unc_est_list, self$par_list)

      list(ad_setup = ad_setup, ad_obj = ad_obj, opt = opt, best_index = best_idx, opt_history = opt_history,
           objective = opt$objective, log_ml = -opt$objective, convergence = opt$convergence,
           sd_rep = sd_rep, par_unc = unc_est_vec, se_unc = unc_se_vec, vcov_unc = Cov_u,
           par_unc_list = unc_est_list, se_unc_list = unc_se_list, par = con_est_list,
           idx_fix_active = idx_fix_active, idx_ran_full = idx_ran_full, idx_ran = index_map$idx_ran,
           factor_levels_seen = index_map$factor_levels_seen, L_u_total = L_u_total,
           fallback_needed = fallback_needed,
           map = target_map, laplace = laplace, reml = reml, jacobian_target = jacobian_target)
    },

    .simple_jacobian = function(f, x, h = 1e-4) {
      p <- length(x)
      f0 <- f(x)
      res <- matrix(0, nrow = length(f0), ncol = p)
      for (i in 1:p) {
        h_i <- max(abs(x[i]), 1e-4) * h
        x_plus <- x
        x_minus <- x
        x_plus[i] <- x[i] + h_i
        x_minus[i] <- x[i] - h_i
        res[, i] <- (f(x_plus) - f(x_minus)) / (2 * h_i)
      }
      return(res)
    },
    # --- Common patterns for REML and parameter classification ---
    fix_patterns = c("Intercept", "Intercept_c", "b", "mean", "prob", "beta", "mu", "delta", "diff", "mean_diff"),
    var_patterns = c("sigma", "sd", "rho", "rho_resid", "L_resid", "corr", "CF_corr", "cov", "CF_cov", "cutpoints", "shape", "phi", "nu", "z", "lambda", "tau", "theta", "alpha", "omega", "kappa"),

    # --- Internal helper for index computation ---
    .compute_active_indices = function(target_map = NULL) {
      idx_fix_active <- integer(0)
      factor_levels_seen <- list()
      opt_par_curr <- 1
      
      map_to_use <- if (!is.null(target_map)) target_map else self$map
      
      for (name in names(self$par_list)) {
        p <- self$par_list[[name]]
        L_u <- p$unc_length
        m <- if (!is.null(map_to_use[[name]])) map_to_use[[name]] else factor(1:L_u)
        
        # Identify indices that are not NA (i.e., estimated)
        idx_active_local <- which(!is.na(m))
        if (length(idx_active_local) > 0) {
          # Check for shared levels (identical factors)
          lvl_names <- as.character(m[idx_active_local])
          for (lvl in unique(lvl_names)) {
            pos_in_active <- which(lvl_names == lvl)
            # The first occurrence of this level gets the next free optimization index
            if (!(lvl %in% names(factor_levels_seen))) {
              factor_levels_seen[[lvl]] <- opt_par_curr
              opt_par_curr <- opt_par_curr + 1
            }
            # Add to global unconstrained index list
            pos_last <- sum(self$pl_full$unc_length[seq_len(which(names(self$par_list) == name) - 1)]) + idx_active_local[pos_in_active]
            idx_fix_active <- c(idx_fix_active, pos_last)
          }
        }
      }
      return(list(idx_active = sort(unique(idx_fix_active)), n_opt = opt_par_curr - 1))
    },

    # --- Internal helper for derived parameter calculation (delta method) ---
    .build_derived_df_table = function(func, base_out, is_generate = FALSE, se_sampling = FALSE, num_samples = 1000) {
      if (is.null(func) || is.null(base_out)) return(NULL)
      
      # This replaces the 150-line local function in optimize() and classic()
      # Integration of both logic versions
      # ... (Implementation uses logic from build_derived_df)
      return(NULL) # Placeholder for now, will implement in next chunk
    },

    # --- Internal helper for fixed parameter recursion ---
    .dispatch_fixed = function(.method_to_call, ...) {
      args <- list(...)
      fixed <- args$fixed
      args$fixed <- NULL
      new_model <- self$fixed_model(fixed)
      # Ensure the method is retrieved as a bound function from the R6 object
      func <- new_model[[.method_to_call]]
      if (!is.function(func)) stop(sprintf("Method '%s' not found in RTMB_Model.", .method_to_call))
      return(do.call(func, args))
    },

    # --- Internal helper for AD function setup ---
    .build_f_ad = function(data_local, par_list_local, log_prob_local, transform_local = NULL, 
                           jacobian_target = "all", adreport = FALSE,
                           fixed_prior_specs_local = NULL, code_model_local = NULL, 
                           shadow_list_local = NULL) {
      function(y_unc_list) {
        para <- to_constrained(y_unc_list, par_list_local)

        if (!is.null(transform_local)) {
          tran_res <- transform_local(data_local, para)
          para <- c(para, tran_res)

          # transform ADREPORT (only if requested, e.g. for sdreport)
          if (adreport && length(tran_res) > 0) {
            for (n in names(tran_res)) {
              # Inject into parent environment so ADREPORT can find it by name
              assign(n, tran_res[[n]])
              eval(substitute(RTMB::ADREPORT(V), list(V = as.name(n))))
            }
          }
        }

        lp <- log_prob_local(data_local, para)

        # Apply dynamic fixed-prior removal if specs are present
        if (!is.null(fixed_prior_specs_local) && length(fixed_prior_specs_local) > 0) {
          lp <- lp - .eval_fixed_prior_contribution_impl(
            code_model = code_model_local,
            shadow_list = shadow_list_local,
            dat = data_local,
            para = para,
            fixed = fixed_prior_specs_local
          )
        }

        if (jacobian_target == "all") {
          lj <- calc_log_jacobian(y_unc_list, par_list_local, only_random = FALSE)
          return(-(lp + lj))
        } else if (jacobian_target == "random") {
          lj <- calc_log_jacobian(y_unc_list, par_list_local, only_random = TRUE)
          return(-(lp + lj))
        } else {
          return(-lp)
        }
      }
    },

    # --- Logic for applying parameter fixing (fixed argument) ---
    .has_multivariate_prior_on = function(expr, target_base) {
      if (is.null(expr)) return(FALSE)

      get_base_name <- function(x) {
        if (is.name(x)) return(as.character(x))
        if (is.call(x) && identical(x[[1]], as.name("["))) {
          return(get_base_name(x[[2]]))
        }
        NULL
      }

      is_multi_dist <- function(rhs) {
        if (!is.call(rhs)) return(FALSE)
        fn <- as.character(rhs[[1]])
        length(fn) == 1 && fn %in% c("multi_normal", "multi_student_t", "multi_cauchy")
      }

      walk <- function(x) {
        if (!is.call(x)) return(FALSE)

        if (identical(x[[1]], as.name("~"))) {
          lhs_base <- get_base_name(x[[2]])
          rhs <- x[[3]]
          if (identical(lhs_base, target_base) && is_multi_dist(rhs)) {
            return(TRUE)
          }
        }

        any(vapply(as.list(x)[-1], walk, logical(1)))
      }

      walk(expr)
    },

    .apply_constraints = function(fixed = list(), silent = FALSE) {
      new_model <- self$clone()
      if (length(fixed) == 0) return(new_model)

      new_map <- as.list(new_model$map)
      new_init <- new_model$init 

      if (is.null(names(new_init))) {
        names(new_init) <- self$pl_full$names
      }

      normalize_map <- function(m) {
        if (is.null(m)) return(NULL)
        factor(as.character(m))
      }

      for (name in names(fixed)) {
        val <- fixed[[name]]
        p_base <- gsub("\\[.*\\]$", "", name)
        has_index <- grepl("\\[", name)

        if (!(p_base %in% names(self$par_list))) stop("Parameter not found: ", p_base)

        if (has_index && !is.null(self$code$model)) {
          if (private$.has_multivariate_prior_on(self$code$model, p_base)) {
            stop(
              sprintf(
                "Cannot use fixed = list('%s' = ...) because '%s' has a multivariate prior. ",
                name, p_base
              ),
              "Partial fixing of a multivariate prior is not supported. ",
              "Use a reduced model instead, e.g. remove the corresponding predictor from the formula.",
              call. = FALSE
            )
          }
        }

        p_info <- self$par_list[[p_base]]
        L_u <- p_info$unc_length

        if (is.null(new_map[[p_base]])) {
          new_map[[p_base]] <- if (L_u > 0) factor(seq_len(L_u)) else factor(integer(0))
        }
        curr_map <- new_map[[p_base]]

        start_idx <- 1
        for (n in names(self$par_list)) {
          if (n == p_base) break
          start_idx <- start_idx + self$par_list[[n]]$length
        }

        if (!grepl("\\[", name)) {
          # Fix the entire parameter
          val_vec <- if (length(val) == 1) rep(val, p_info$length) else val
          new_init[start_idx:(start_idx + p_info$length - 1)] <- as.numeric(val_vec)

          if (L_u > 0) {
            # IMPORTANT: Recreate as a factor with no levels instead of curr_map[] <- NA
            curr_map <- factor(rep(NA_character_, L_u))
          }

        } else {
          # Fix specific elements
          # Guard: element-wise fixing is not supported for structurally constrained types
          # where constrained length != unconstrained length (simplex, corr_matrix, etc.)
          constrained_types <- c("simplex", "centered", "corr_matrix", "cov_matrix",
                                 "CF_corr", "CF_cov", "centered_matrix",
                                 "centered_tri", "positive_centered_tri")
          if (p_info$length != L_u && (p_info$type %in% constrained_types)) {
            stop(sprintf(
              "Element-wise fixing is not supported for parameter '%s' (type = '%s'). ",
              p_base, p_info$type
            ),
            "Fix the entire parameter instead with fixed = list('", p_base, "' = ...).",
            call. = FALSE)
          }

          idx_str <- gsub(".*\\[(.*)\\].*", "\\1", name)
          p_names <- self$par_names[[p_base]]
          idx <- suppressWarnings(as.integer(idx_str))
          if (is.na(idx) && !is.null(p_names)) idx <- which(p_names == idx_str)

          if (length(idx) > 0) {
            new_init[start_idx + idx - 1] <- as.numeric(val)

            if (L_u > 0) {
              # For types where length == unc_length, constrained index maps directly
              if (idx <= L_u) {
                curr_map[idx] <- NA
              }
              # Drop unused levels even in partial fixing
              curr_map <- normalize_map(curr_map)
            }
          }
        }

        new_map[[p_base]] <- curr_map
      }

      new_model$init <- new_init
      new_model$map <- new_map
      
      # Dynamic prior correction for Bayes Factor (log-density of fixed parameters)
      # Instead of wrapping log_prob (which is messy for printing), we store the specs
      # and handle removal during AD objective construction in .build_f_ad.
      new_model$prior_correction <- 0 # Explicitly reset scalar correction
      
      # Merge with existing fixed_prior_specs if any
      merged_specs <- self$fixed_prior_specs %||% list()
      merged_specs[names(fixed)] <- fixed
      new_model$fixed_prior_specs <- merged_specs

      # Validate to catch manual lp increments involving fixed parameters early
      # Align with runtime by including transformed parameters if they exist
      test_para <- constrained_vector_to_list(new_init, self$par_list)
      if (!is.null(self$transform)) {
        test_tran <- tryCatch(suppressMessages(self$transform(self$data, test_para)), error = function(e) list())
        if (length(test_tran) > 0) {
          missing <- setdiff(names(test_tran), names(test_para))
          # Exclude variables that are already in data
          missing <- setdiff(missing, names(self$data))
          if (length(missing) > 0) test_para <- c(test_para, test_tran[missing])
        }
      }
      
      # Perform validation using the internal implementation function
      # We need a temporary shadow list for validation
      env_shadow_tmp <- list2env(self$data, parent = asNamespace("BayesRTMB"))
      idx_tmp <- 1
      for (nm in names(self$par_list)) {
        p <- self$par_list[[nm]]
        shadow_vec <- self$pl_full$names[idx_tmp:(idx_tmp + p$length - 1)]
        if (length(p$dim) > 1) dim(shadow_vec) <- p$dim
        env_shadow_tmp[[nm]] <- shadow_vec
        idx_tmp <- idx_tmp + p$length
      }
      if (!is.null(self$transform)) {
        par_only <- mget(names(self$par_list), envir = env_shadow_tmp, inherits = FALSE)
        tran_res_shadow <- tryCatch(suppressMessages(self$transform(self$data, par_only)), error = function(e) list())
        for (nm in names(tran_res_shadow)) env_shadow_tmp[[nm]] <- tran_res_shadow[[nm]]
      }
      shadow_list_tmp <- as.list(env_shadow_tmp)

      suppressWarnings(.eval_fixed_prior_contribution_impl(
        code_model = self$code$model,
        shadow_list = shadow_list_tmp,
        dat = self$data,
        para = test_para,
        fixed = merged_specs
      ))

      # Prune fully fixed parameters from extra$marginal if they exist
      if (!is.null(new_model$extra$marginal)) {
        fully_fixed <- names(new_map)[vapply(new_map, function(m) !is.null(m) && all(is.na(m)), logical(1))]
        if (length(fully_fixed) > 0) {
          new_model$extra$marginal <- setdiff(new_model$extra$marginal, fully_fixed)
        }
      }

      return(new_model)
    }
  )
)

# --- Internal helper: calculate log-density of fixed parameters dynamically ---
# Note: Only priors specified as sampling statements (x ~ dist(...)) are automatically removed.
# Manual increments to 'lp' involving fixed parameters will trigger an error to avoid incorrect Bayes Factors.
.eval_fixed_prior_contribution_impl <- function(code_model, shadow_list, dat, para, fixed) {
  if (length(fixed) == 0 || is.null(code_model)) return(0)

  # 1. Prepare environment with CURRENT parameter values (including transformations)
  env_val <- list2env(dat, parent = asNamespace("BayesRTMB"))
  for (nm in names(para)) env_val[[nm]] <- para[[nm]]
  
  # 2. Shadow environment for identifying fixed components
  # Use the provided shadow_list (pre-calculated to avoid redundant transform calls)
  env_shadow <- list2env(shadow_list, parent = asNamespace("BayesRTMB"))
  fixed_keys <- names(fixed)
  fixed_bases <- gsub("\\[.*\\]$", "", fixed_keys)
  tainted_vars <- character(0)
  lp_corr <- 0
  
  evaluate_node <- function(x) {
    if (is.call(x)) {
      if (identical(x[[1]], as.name("~"))) {
        lhs_expr <- x[[2]]
        rhs_call <- x[[3]]
        dist_name <- as.character(rhs_call[[1]])
        
        # Evaluate LHS in shadow env to find which elements are fixed
        # We use tryCatch because transformed variables might not be representable as strings
        shadow_names <- tryCatch(eval(lhs_expr, envir = env_shadow), error = function(e) NULL)
        
        # Identify fixed elements
        fixed_mask <- if (!is.null(shadow_names)) {
          (shadow_names %in% fixed_keys) | (gsub("\\[.*\\]$", "", shadow_names) %in% fixed_keys)
        } else {
          FALSE
        }
        
        if (any(fixed_mask)) {
          val_lhs <- eval(lhs_expr, envir = env_val)
          val_args <- lapply(as.list(rhs_call[-1]), eval, envir = env_val)
          
          # Call distribution function
          lpdf_name <- paste0(dist_name, "_lpdf")
          lpmf_name <- paste0(dist_name, "_lpmf")
          
          # Find the function in BayesRTMB namespace
          pkg_ns <- asNamespace("BayesRTMB")
          dist_fn <- if (exists(lpdf_name, envir = pkg_ns)) get(lpdf_name, envir = pkg_ns)
                     else if (exists(lpmf_name, envir = pkg_ns)) get(lpmf_name, envir = pkg_ns)
                     else NULL
          
          if (!is.null(dist_fn)) {
            # If multivariate, we take the whole thing (we already checked partial fixing for multi)
            is_multi <- dist_name %in% c("multi_normal", "multi_student_t", "multi_cauchy")
            
            if (is_multi) {
              # Full parameter must be fixed
              lp_corr <<- lp_corr + do.call(dist_fn, c(list(val_lhs), val_args))
            } else {
              # Univariate - might be partial
              # Get element-wise density (force sum=FALSE if supported)
              densities <- tryCatch(
                do.call(dist_fn, c(list(val_lhs), val_args, list(sum = FALSE))),
                error = function(e) do.call(dist_fn, c(list(val_lhs), val_args))
              )
              
              # Check if element-wise was needed but only a scalar was returned
              if (length(densities) == 1L && length(fixed_mask) > 1L && !all(fixed_mask)) {
                stop(
                  sprintf("Cannot partially remove a vectorized prior for '%s' because element-wise log densities are unavailable.", dist_name),
                  call. = FALSE
                )
              }
              
              # Flatten to handle matrix/array accurately
              densities_v <- as.vector(densities)
              mask_v <- as.vector(fixed_mask)
              
              if (length(densities_v) == length(mask_v)) {
                lp_corr <<- lp_corr + sum(densities_v[mask_v])
              } else if (length(densities_v) == 1L) {
                # Scalar fallback (only if all elements are fixed, already checked above)
                lp_corr <<- lp_corr + densities_v
              }
            }
          } else {
            stop(
              sprintf("Cannot remove fixed prior contribution for distribution '%s': corresponding lpdf/lpmf function was not found.", dist_name),
              call. = FALSE
            )
          }
        }
      } else if (identical(x[[1]], as.name("for"))) {
        # Execute the loop
        var_name <- as.character(x[[2]])
        seq_val <- eval(x[[3]], envir = env_val)
        loop_body <- x[[4]]
        
        for (v in seq_val) {
          env_val[[var_name]] <- v
          env_shadow[[var_name]] <- v
          evaluate_node(loop_body)
        }
      } else if (identical(x[[1]], as.name("{"))) {
        lapply(as.list(x)[-1], evaluate_node)
      } else if (identical(x[[1]], as.name("if"))) {
        cond <- tryCatch(eval(x[[2]], envir = env_val), error = function(e) FALSE)
        if (isTRUE(cond)) evaluate_node(x[[3]]) else if (length(x) == 4) evaluate_node(x[[4]])
      } else if (identical(x[[1]], as.name("<-")) || identical(x[[1]], as.name("="))) {
        # Check for manual log-prob increments or assignments involving fixed parameters
        lhs_vars <- all.vars(x[[2]])
        lhs_name <- if (length(lhs_vars) > 0) lhs_vars[1] else as.character(x[[2]])[1]
        rhs_vars <- all.vars(x[[3]])
        
        # If the RHS involves fixed or already tainted variables
        if (any(setdiff(rhs_vars, "lp") %in% c(fixed_bases, tainted_vars))) {
          if (identical(lhs_name, "lp")) {
            stop("Manual log-probability increments (lp <- ...) involving fixed parameters or derived variables are not supported for automatic prior removal. Please use sampling statements (~) instead.", call. = FALSE)
          } else if (!is.na(lhs_name)) {
            # Taint the LHS variable so that subsequent use in lp <- ... can be detected
            tainted_vars <<- unique(c(tainted_vars, lhs_name))
          }
        }
      } else {
        # Recurse into all other calls
        lapply(as.list(x)[-1], evaluate_node)
      }
    }
  }

  evaluate_node(code_model)
  return(lp_corr)
}

