#' Generalized Neural Network Trainer
#'
#' @description
#' `r lifecycle::badge("experimental")`
#'
#' `train_nn()` is a generic function for training neural networks with a
#' user-defined architecture via [nn_arch()]. Dispatch is based on the class
#' of `x`.
#'
#' Recommended workflow:
#' 1. Define architecture with [nn_arch()] (optional).
#' 2. Train with `train_nn()`.
#' 3. Predict with [predict.nn_fit()].
#'
#' All methods delegate to a shared implementation core after preprocessing.
#' When `architecture = NULL`, the model falls back to a plain feed-forward neural network
#' (`nn_linear`) architecture.
#'
#' @param x Dispatch is based on its current class:
#'   - `matrix`: used directly, no preprocessing applied.
#'   - `data.frame`: preprocessed via `hardhat::mold()`. `y` may be a vector /
#'     factor / matrix of outcomes, or a formula describing the outcome–predictor
#'     relationship within `x`.
#'   - `formula`: combined with `data` and preprocessed via `hardhat::mold()`.
#'   - `dataset`: a `torch` dataset object; batched via `torch::dataloader()`.
#'     This is the recommended interface for sequence/time-series and image data.
#' @param y Outcome data. Interpretation depends on the method:
#'   - For the `matrix` and `data.frame` methods: a numeric vector, factor, or
#'     matrix of outcomes.
#'   - For the `data.frame` method only: alternatively a formula of the form
#'     `outcome ~ predictors`, evaluated against `x`.
#'   - Ignored when `x` is a formula (outcome is taken from the formula) or a
#'     `dataset` (labels come from the dataset itself).
#' @param data A data frame. Required when `x` is a formula.
#' @param hidden_neurons Integer vector specifying the number of neurons in each
#'   hidden layer, e.g. `c(128, 64)` for two hidden layers. When `NULL` or missing,
#'   no hidden layers are used and the model reduces to a single linear mapping
#'   from inputs to outputs.
#' @param activations Activation function specification(s) for the hidden layers.
#'   See [act_funs()] for supported values. Recycled if a single value is given.
#' @param output_activation Optional activation function for the output layer.
#'   Defaults to `NULL` (no activation / linear output).
#' @param bias Logical. Whether to include bias terms in each layer. Default `TRUE`.
#' @param architecture An [nn_arch()] object specifying a custom architecture. Default
#'   `NULL`, which falls back to a standard feed-forward network.
#' @param arch Backward-compatible alias for `architecture`. If both are supplied,
#'   they must be identical.
#' @param flatten_input Logical or `NULL` (dataset method only). Controls whether
#'   each batch/sample is flattened to 2D before entering the model. `NULL`
#'   (default) auto-selects: `TRUE` when `architecture = NULL`, otherwise `FALSE`.
#' @param early_stopping An [early_stop()] object specifying early stopping
#'   behaviour, or `NULL` (default) to disable early stopping. When supplied,
#'   training halts if the monitored metric does not improve by at least
#'   `min_delta` for `patience` consecutive epochs.
#'   Example: `early_stopping = early_stop(patience = 10)`.
#' @param epochs Positive integer. Number of full passes over the training data.
#'   Default `100`.
#' @param batch_size Positive integer. Number of samples per mini-batch. Default `32`.
#' @param penalty Non-negative numeric. L1/L2 regularization strength (lambda).
#'   Default `0` (no regularization).
#' @param mixture Numeric in \[0, 1\]. Elastic net mixing parameter: `0` = pure
#'   ridge (L2), `1` = pure lasso (L1). Default `0`.
#' @param learn_rate Positive numeric. Step size for the optimizer. Default `0.001`.
#' @param optimizer Character. Optimizer algorithm. One of `"adam"` (default),
#'   `"sgd"`, or `"rmsprop"`.
#' @param optimizer_args Named list of additional arguments forwarded to the
#'   optimizer constructor (e.g. `list(momentum = 0.9)` for SGD). Default `list()`.
#' @param loss Character or function. Loss function used during training. Built-in
#'   options: `"mse"` (default), `"mae"`, `"cross_entropy"`, or `"bce"`. For
#'   classification tasks with a scalar label, `"cross_entropy"` is set
#'   automatically. Alternatively, supply a custom function or formula with
#'   signature `function(input, target)` returning a scalar `torch_tensor`.
#' @param validation_split Numeric in \[0, 1). Proportion of training data held
#'   out for validation. Default `0` (no validation set).
#' @param device Character. Compute device: `"cpu"`, `"cuda"`, or `"mps"`.
#'   Default `NULL`, which auto-detects the best available device.
#' @param verbose Logical. If `TRUE`, prints loss (and validation loss) at regular
#'   intervals during training. Default `FALSE`.
#' @param cache_weights Logical. If `TRUE`, stores a copy of the trained weight
#'   matrices in the returned object under `$cached_weights`. Default `FALSE`.
#' @param ... Additional arguments passed to specific methods.
#'
#' @return An object of class `"nn_fit"`, or one of its subclasses:
#'   - `c("nn_fit_tab", "nn_fit")` — returned by the `data.frame` and `formula` methods
#'   - `c("nn_fit_ds", "nn_fit")` — returned by the `dataset` method
#'
#'   All subclasses share a common structure. See **Details** for the list of
#'   components.
#'
#' @details
#' The returned `"nn_fit"` object is a named list with the following components:
#'
#' - `model` — the trained `torch::nn_module` object
#' - `fitted` — fitted values on the training data (or `NULL` for dataset fits)
#' - `loss_history` — numeric vector of per-epoch training loss, trimmed to
#'   actual epochs run (relevant when early stopping is active)
#' - `val_loss_history` — per-epoch validation loss, or `NULL` if
#'   `validation_split = 0`
#' - `n_epochs` — number of epochs actually trained
#' - `stopped_epoch` — epoch at which early stopping triggered, or `NA` if
#'   training ran to completion
#' - `hidden_neurons`, `activations`, `output_activation` — architecture spec
#' - `penalty`, `mixture` — regularization settings
#' - `feature_names`, `response_name` — variable names (tabular methods only)
#' - `no_x`, `no_y` — number of input features and output nodes
#' - `is_classification` — logical flag
#' - `y_levels`, `n_classes` — class labels and count (classification only)
#' - `device` — device the model is on
#' - `cached_weights` — list of weight matrices, or `NULL`
#' - `arch` — the `nn_arch` object used, or `NULL`
#'
#' @section Supported tasks and input formats:
#' `train_nn()` is task-agnostic by design (no explicit `task` argument).
#' Task behavior is determined by your input interface and architecture:
#' - **Tabular data**: use `matrix`, `data.frame`, or `formula` methods.
#' - **Time series**: use the `dataset` method with per-item tensors shaped as
#'   `[time, features]` (or your preferred convention) and a recurrent
#'   architecture via [nn_arch()].
#' - **Image classification**: use the `dataset` method with per-item tensors
#'   shaped for your first layer (commonly `[channels, height, width]` for
#'   `torch::nn_conv2d`). If your source arrays are channel-last, reorder in the
#'   dataset or via `input_transform`.
#'
#' @examples
#' \donttest{
#' if (torch::torch_is_installed()) {
#'     # Matrix method — no preprocessing
#'     model = train_nn(
#'         x = as.matrix(iris[, 2:4]),
#'         y = iris$Sepal.Length,
#'         hidden_neurons = c(64, 32),
#'         activations = "relu",
#'         epochs = 50
#'     )
#'
#'     # Data frame method — y as a vector
#'     model = train_nn(
#'         x = iris[, 2:4],
#'         y = iris$Sepal.Length,
#'         hidden_neurons = c(64, 32),
#'         activations = "relu",
#'         epochs = 50
#'     )
#'
#'     # Data frame method — y as a formula evaluated against x
#'     model = train_nn(
#'         x = iris,
#'         y = Sepal.Length ~ . - Species,
#'         hidden_neurons = c(64, 32),
#'         activations = "relu",
#'         epochs = 50
#'     )
#'
#'     # Formula method — outcome derived from formula
#'     model = train_nn(
#'         x = Sepal.Length ~ .,
#'         data = iris[, 1:4],
#'         hidden_neurons = c(64, 32),
#'         activations = "relu",
#'         epochs = 50
#'     )
#'
#'     # No hidden layers — linear model
#'     model = train_nn(
#'         x = Sepal.Length ~ .,
#'         data = iris[, 1:4],
#'         epochs = 50
#'     )
#'
#'     # Architecture object (nn_arch -> train_nn)
#'     mlp_arch = nn_arch(nn_name = "mlp_model")
#'     model = train_nn(
#'         x = Sepal.Length ~ .,
#'         data = iris[, 1:4],
#'         hidden_neurons = c(64, 32),
#'         activations = "relu",
#'         architecture = mlp_arch,
#'         epochs = 50
#'     )
#'
#'     # Custom layer architecture
#'     custom_linear = torch::nn_module(
#'         "CustomLinear",
#'         initialize = function(in_features, out_features, bias = TRUE) {
#'             self$layer = torch::nn_linear(in_features, out_features, bias = bias)
#'         },
#'         forward = function(x) self$layer(x)
#'     )
#'     custom_arch = nn_arch(
#'         nn_name = "custom_linear_mlp",
#'         nn_layer = ~ custom_linear
#'     )
#'     model = train_nn(
#'         x = Sepal.Length ~ .,
#'         data = iris[, 1:4],
#'         hidden_neurons = c(16, 8),
#'         activations = "relu",
#'         architecture = custom_arch,
#'         epochs = 50
#'     )
#'
#'     # With early stopping
#'     model = train_nn(
#'         x = Sepal.Length ~ .,
#'         data = iris[, 1:4],
#'         hidden_neurons = c(64, 32),
#'         activations = "relu",
#'         epochs = 200,
#'         validation_split = 0.2,
#'         early_stopping = early_stop(patience = 10)
#'     )
#' }
#' }
#'
#' @seealso [predict.nn_fit()], [nn_arch()], [act_funs()], [early_stop()]
#' @name gen-nn-train
#' @export
train_nn = function(x, ...) {
    UseMethod("train_nn")
}


#' Normalize architecture argument for train_nn()
#' @noRd
.resolve_train_architecture = function(architecture = NULL, arch = NULL) {
    if (!is.null(architecture) && !is.null(arch) && !identical(architecture, arch)) {
        cli::cli_abort(c(
            "{.arg architecture} and {.arg arch} were both supplied with different values.",
            i = "Supply only {.arg architecture}, or provide identical values for both."
        ))
    }

    out = architecture %||% arch

    if (!is.null(out) && !inherits(out, "nn_arch")) {
        cli::cli_abort("{.arg architecture} must be an {.cls nn_arch} object created with {.fn nn_arch}.")
    }

    out
}


#' Normalize predict newdata argument
#' @noRd
.resolve_predict_newdata = function(newdata = NULL, new_data = NULL) {
    if (!is.null(new_data) && is.null(newdata)) {
        cli::cli_warn("{.arg new_data} is a legacy alias. Prefer {.arg newdata}.")
        return(new_data)
    }

    newdata
}

#' @rdname gen-nn-train
#'
#' @section Matrix method:
#' When `x` is supplied as a raw numeric matrix, no preprocessing is applied.
#' Data is passed directly to the shared `train_nn_impl` core.
#'
#' @export
train_nn.matrix =
    function(
        x,
        y,
        hidden_neurons = NULL,
        activations = NULL,
        output_activation = NULL,
        bias = TRUE,
        arch = NULL,
        architecture = NULL,
        early_stopping = NULL,
        epochs = 100,
        batch_size = 32,
        penalty = 0,
        mixture = 0,
        learn_rate = 0.001,
        optimizer = "adam",
        optimizer_args = list(),
        loss = "mse",
        validation_split = 0,
        device = NULL,
        verbose = FALSE,
        cache_weights = FALSE,
        ...
    ) {
        arch = .resolve_train_architecture(architecture = architecture, arch = arch)

        act_specs = eval_act_funs({{ activations }}, {{ output_activation }})
        activations = act_specs$activations
        output_activation = act_specs$output_activation

        train_nn_impl(
            x = x,
            y = y,
            hidden_neurons = hidden_neurons,
            activations = activations,
            output_activation = output_activation,
            bias = bias,
            arch = arch,
            early_stopping = early_stopping,
            epochs = epochs,
            batch_size = batch_size,
            penalty = penalty,
            mixture = mixture,
            learn_rate = learn_rate,
            optimizer = optimizer,
            optimizer_args = optimizer_args,
            loss = loss,
            validation_split = validation_split,
            device = device,
            verbose = verbose,
            cache_weights = cache_weights,
            fit_class = "nn_fit"
        )
    }


#' @rdname gen-nn-train
#'
#' @section Data frame method:
#' When `x` is a data frame, `y` can be either a vector / factor / matrix of
#' outcomes, or a formula of the form `outcome ~ predictors` evaluated against
#' `x`. Preprocessing is handled by `hardhat::mold()`.
#'
#' @export
train_nn.data.frame =
    function(
        x,
        y,
        hidden_neurons = NULL,
        activations = NULL,
        output_activation = NULL,
        bias = TRUE,
        arch = NULL,
        architecture = NULL,
        early_stopping = NULL,
        epochs = 100,
        batch_size = 32,
        penalty = 0,
        mixture = 0,
        learn_rate = 0.001,
        optimizer = "adam",
        optimizer_args = list(),
        loss = "mse",
        validation_split = 0,
        device = NULL,
        verbose = FALSE,
        cache_weights = FALSE,
        ...
    ) {
        arch = .resolve_train_architecture(architecture = architecture, arch = arch)

        act_specs = eval_act_funs({{ activations }}, {{ output_activation }})
        activations = act_specs$activations
        output_activation = act_specs$output_activation

        processed = if (rlang::is_formula(y)) hardhat::mold(y, x) else hardhat::mold(x, y)

        .train_nn_tab_impl(
            processed = processed,
            formula = NULL,
            hidden_neurons = hidden_neurons,
            activations = activations,
            output_activation = output_activation,
            bias = bias,
            arch = arch,
            early_stopping = early_stopping,
            epochs = epochs,
            batch_size = batch_size,
            penalty = penalty,
            mixture = mixture,
            learn_rate = learn_rate,
            optimizer = optimizer,
            optimizer_args = optimizer_args,
            loss = loss,
            validation_split = validation_split,
            device = device,
            verbose = verbose,
            cache_weights = cache_weights
        )
    }


#' @rdname gen-nn-train
#'
#' @section Formula method:
#' When `x` is a formula, `data` must be supplied as the data frame against
#' which the formula is evaluated. Preprocessing is handled by `hardhat::mold()`.
#'
#' @export
train_nn.formula =
    function(
        x,
        data,
        hidden_neurons = NULL,
        activations = NULL,
        output_activation = NULL,
        bias = TRUE,
        arch = NULL,
        architecture = NULL,
        early_stopping = NULL,
        epochs = 100,
        batch_size = 32,
        penalty = 0,
        mixture = 0,
        learn_rate = 0.001,
        optimizer = "adam",
        optimizer_args = list(),
        loss = "mse",
        validation_split = 0,
        device = NULL,
        verbose = FALSE,
        cache_weights = FALSE,
        ...
    ) {
        arch = .resolve_train_architecture(architecture = architecture, arch = arch)

        if (missing(data) || is.null(data)) {
            cli::cli_abort("{.arg data} must be provided when using the formula interface.")
        }

        act_specs = eval_act_funs({{ activations }}, {{ output_activation }})
        activations = act_specs$activations
        output_activation = act_specs$output_activation

        processed = hardhat::mold(x, data)

        .train_nn_tab_impl(
            processed = processed,
            formula = x,
            hidden_neurons = hidden_neurons,
            activations = activations,
            output_activation = output_activation,
            bias = bias,
            arch = arch,
            early_stopping = early_stopping,
            epochs = epochs,
            batch_size = batch_size,
            penalty = penalty,
            mixture = mixture,
            learn_rate = learn_rate,
            optimizer = optimizer,
            optimizer_args = optimizer_args,
            loss = loss,
            validation_split = validation_split,
            device = device,
            verbose = verbose,
            cache_weights = cache_weights
        )
    }


#' @rdname gen-nn-train
#' @export
train_nn.default = function(x, ...) {
    cli::cli_abort(c(
        "No {.fn train_nn} method for class {.cls {class(x)}}.",
        i = "Supported classes: {.cls matrix}, {.cls data.frame}, {.cls formula}, {.cls dataset}."
    ))
}


#' Preprocessing bridge for data.frame and formula methods
#' @keywords internal
.train_nn_tab_impl =
    function(
        processed,
        formula,
        hidden_neurons,
        activations,
        output_activation,
        bias,
        arch,
        early_stopping,
        epochs,
        batch_size,
        penalty,
        mixture,
        learn_rate,
        optimizer,
        optimizer_args,
        loss,
        validation_split,
        device,
        verbose,
        cache_weights
    ) {
        predictors = processed$predictors
        outcomes = processed$outcomes

        if (!is.matrix(predictors)) predictors = as.matrix(predictors)

        if (is.data.frame(outcomes)) {
            outcomes = if (ncol(outcomes) == 1L) outcomes[[1L]] else as.matrix(outcomes)
        } else if (!is.null(ncol(outcomes)) && ncol(outcomes) == 1L) {
            outcomes = outcomes[[1L]]
        }

        fit = train_nn_impl(
            x = predictors,
            y = outcomes,
            hidden_neurons = hidden_neurons,
            activations = activations,
            output_activation = output_activation,
            bias = bias,
            arch = arch,
            early_stopping = early_stopping,
            epochs = epochs,
            batch_size = batch_size,
            penalty = penalty,
            mixture = mixture,
            learn_rate = learn_rate,
            optimizer = optimizer,
            optimizer_args = optimizer_args,
            loss = loss,
            validation_split = validation_split,
            device = device,
            verbose = verbose,
            cache_weights = cache_weights,
            fit_class = "nn_fit_tab"
        )

        fit$blueprint = processed$blueprint
        if (!is.null(formula)) fit$formula = formula

        fit
    }


#' Shared core implementation
#' @keywords internal
train_nn_impl =
    function(
        x,
        y,
        hidden_neurons,
        activations = NULL,
        output_activation = NULL,
        bias = TRUE,
        arch = NULL,
        early_stopping = NULL,
        epochs = 100,
        batch_size = 32,
        penalty = 0,
        mixture = 0,
        learn_rate = 0.001,
        optimizer = "adam",
        optimizer_args = list(),
        loss = "mse",
        validation_split = 0,
        device = NULL,
        verbose = FALSE,
        cache_weights = FALSE,
        fit_class = "nn_fit"
    ) {
        if (!requireNamespace("torch", quietly = TRUE)) {
            cli::cli_abort("Package {.pkg torch} is required but not installed.")
        }

        if (missing(hidden_neurons) || is.null(hidden_neurons) || length(hidden_neurons) == 0L) {
            hidden_neurons = integer(0L)
        }

        # ---- Device ----
        device = if (is.null(device)) get_default_device() else validate_device(device)
        if (verbose) cli::cli_alert_info("Using device: {device}")

        validate_regularization(penalty, mixture)

        # ---- Input transform ----
        input_fn = if (!is.null(arch) && !is.null(arch$input_transform)) {
            rlang::as_function(arch$input_transform)
        } else {
            identity
        }

        # ---- Metadata ----
        feature_names = colnames(x) %||% paste0("V", seq_len(ncol(x)))
        response_name = names(y)[1L] %||% "y"
        is_classification = is.factor(y) || is.character(y)

        # ---- y encoding ----
        if (is_classification) {
            if (is.character(y)) y = as.factor(y)
            y_levels = levels(y)
            n_classes = length(y_levels)
            y_numeric = as.integer(y)
            no_y = n_classes
        } else {
            y_levels = NULL
            n_classes = NULL
            y_numeric = if (is.matrix(y)) y else as.numeric(y)
            no_y = if (is.matrix(y)) ncol(y) else 1L
        }

        no_x = ncol(x)
        n_obs = nrow(x)

        # ---- Validation split ----
        if (validation_split > 0 && validation_split < 1) {
            n_val = floor(n_obs * validation_split)
            val_idx = sample(n_obs, n_val)
            train_idx = setdiff(seq_len(n_obs), val_idx)

            x_train = x[train_idx, , drop = FALSE]
            y_train = y_numeric[train_idx]
            x_val = x[val_idx, , drop = FALSE]
            y_val = y_numeric[val_idx]
        } else {
            x_train = x
            y_train = y_numeric
            x_val = NULL
            y_val = NULL
        }

        # ---- Build model ----
        arch_args = if (!is.null(arch)) {
            args = unclass(arch)
            args$input_transform = NULL
            args
        } else {
            list()
        }

        arch_env = if (!is.null(arch)) attr(arch, "env") else parent.frame()
        # model_expr = do.call(
        #     nn_module_generator,
        #     c(
        #         list(
        #             hd_neurons = hidden_neurons,
        #             no_x = no_x,
        #             no_y = no_y,
        #             activations = activations,
        #             output_activation = output_activation,
        #             bias = bias
        #         ),
        #         arch_args,
        #         .env = arch_env
        #     )
        # )
        model_expr = rlang::exec(
            nn_module_generator,

            # Arguments
            hd_neurons = hidden_neurons,
            no_x = no_x,
            no_y = no_y,
            activations = activations,
            output_activation = output_activation,
            bias = bias,

            # Executing arguments from `nn_arch()`
            # To sync up with the current environment
            !!!arch_args,
            .env = arch_env
        )

        model = rlang::eval_tidy(model_expr)()
        model$to(device = device)

        # ---- Early stopping state ----
        es_state = if (!is.null(early_stopping)) {
            if (!inherits(early_stopping, "early_stop_spec")) {
                cli::cli_abort(c(
                    "{.arg early_stopping} must be created with {.fn early_stop}.",
                    i = "Example: {.code early_stopping = early_stop(patience = 10)}."
                ))
            }
            if (early_stopping$monitor == "val_loss" && validation_split == 0) {
                cli::cli_abort(c(
                    "Early stopping on {.val val_loss} requires a validation set.",
                    i = "Set {.arg validation_split} > 0, or use {.code monitor = 'train_loss'}."
                ))
            }
            list(
                spec = early_stopping,
                best_loss = Inf,
                best_weights = NULL,
                wait = 0L,
                stopped_epoch = NA_integer_
            )
        } else {
            NULL
        }

        # ---- Tensors ----
        .make_tensor = function(mat, is_y = FALSE) {
            t = torch::torch_tensor(
                mat,
                device = device,
                dtype = if (is_y && is_classification) torch::torch_long()
                else torch::torch_float32()
            )
            if (!is_y) input_fn(t) else t
        }

        x_train_t = .make_tensor(x_train)
        y_train_t = .make_tensor(
            if (is_classification || is.matrix(y_train)) y_train else matrix(y_train, ncol = 1L),
            is_y = TRUE
        )

        if (!is.null(x_val)) {
            x_val_t = .make_tensor(x_val)
            y_val_t = .make_tensor(
                if (is_classification || is.matrix(y_val)) y_val else matrix(y_val, ncol = 1L),
                is_y = TRUE
            )
        }

        # ---- Optimizer ----
        validate_optimizer(tolower(optimizer))
        optimizer_fn = get(paste0("optim_", tolower(optimizer)), envir = asNamespace("torch"))
        # opt = do.call(optimizer_fn, c(list(params = model$parameters, lr = learn_rate), optimizer_args))
        opt = rlang::exec(
            optimizer_fn,

            # Arguments
            params = model$parameters,
            lr = learn_rate,

            # Extra arguments from `optimizer_args`
            !!!optimizer_args
        )

        # ---- Loss function ----
        loss_fn = if (is.function(loss)) {
            .validate_loss_fn(loss)
        } else if (rlang::is_formula(loss)) {
            .validate_loss_fn(rlang::as_function(loss))
        } else {
            loss_name = tolower(loss)

            if (is_classification && loss_name == "mse") {
                loss_name = "cross_entropy"
                if (verbose) cli::cli_alert("Auto-detected classification task. Using {.val cross_entropy} loss.")
            }

            switch(
                loss_name,
                mse = function(input, target) torch::nnf_mse_loss(input, target),
                mae = function(input, target) torch::nnf_l1_loss(input, target),
                cross_entropy = function(input, target) torch::nnf_cross_entropy(input, target),
                bce = function(input, target) torch::nnf_binary_cross_entropy_with_logits(input, target),
                cli::cli_abort(
                    "Unknown loss: {.val {loss_name}}.",
                    i = "Use one of {.val mse}, {.val mae}, {.val cross_entropy}, {.val bce}, or supply a {.cls function}."
                )
            )
        }

        # ---- Training loop ----
        loss_history = numeric(epochs)
        val_loss_history = if (!is.null(x_val)) numeric(epochs) else NULL
        n_batches = ceiling(nrow(x_train) / batch_size)

        for (epoch in seq_len(epochs)) {
            model$train()
            epoch_loss = 0
            idx = sample(nrow(x_train))

            for (batch in seq_len(n_batches)) {
                start_i = (batch - 1L) * batch_size + 1L
                end_i = min(batch * batch_size, nrow(x_train))
                batch_idx = idx[start_i:end_i]

                x_batch = x_train_t[batch_idx, ]
                y_batch = y_train_t[batch_idx]

                opt$zero_grad()
                y_pred = model(x_batch)
                batch_loss = loss_fn(y_pred, y_batch)
                reg_loss = regularizer(model, penalty, mixture)
                total_loss = batch_loss + reg_loss
                total_loss$backward()
                opt$step()

                epoch_loss = epoch_loss + total_loss$item()
            }

            loss_history[epoch] = epoch_loss / n_batches

            # ---- Validation ----
            if (!is.null(x_val)) {
                model$eval()
                torch::with_no_grad({
                    y_val_pred = model(x_val_t)
                    val_loss = loss_fn(y_val_pred, y_val_t)
                    val_loss_history[epoch] = val_loss$item()
                })
            }

            # ---- Verbose ----
            if (verbose && (epoch %% max(1L, epochs %/% 10L) == 0L || epoch == epochs)) {
                msg = glue::glue(
                    "Epoch {epoch}/{epochs} - Loss: {round(loss_history[epoch], 4)}"
                )
                if (!is.null(val_loss_history))
                    msg = glue::glue("{msg} - Val Loss: {round(val_loss_history[epoch], 4)}")

                message(msg)
            }

            # ---- Early stopping check ----
            if (!is.null(es_state)) {
                monitored = if (es_state$spec$monitor == "val_loss") {
                    val_loss_history[epoch]
                } else {
                    loss_history[epoch]
                }

                if (monitored < es_state$best_loss - es_state$spec$min_delta) {
                    es_state$best_loss = monitored
                    es_state$wait = 0L
                    if (es_state$spec$restore_best_weights) {
                        es_state$best_weights = lapply(
                            model$parameters,
                            function(p) p$detach()$clone()
                        )
                    }
                } else {
                    es_state$wait = es_state$wait + 1L
                    if (es_state$wait >= es_state$spec$patience) {
                        es_state$stopped_epoch = epoch
                        if (verbose) cli::cli_alert_warning(
                            "Early stopping at epoch {epoch}. Best {es_state$spec$monitor}: {round(es_state$best_loss, 4)}."
                        )
                        break
                    }
                }
            }
        }

        # ---- Restore best weights ----
        if (!is.null(es_state) && es_state$spec$restore_best_weights && !is.null(es_state$best_weights)) {
            params = model$parameters
            torch::with_no_grad({
                for (i in seq_along(params)) {
                    params[[i]]$set_data(es_state$best_weights[[i]])
                }
            })
            if (verbose) {
                best_ep = if (!is.na(es_state$stopped_epoch)) {
                    es_state$stopped_epoch - es_state$spec$patience
                } else {
                    epoch
                }
                cli::cli_alert_info("Restored best weights from epoch {best_ep}.")
            }
        }

        # ---- Trim histories to actual epochs run ----
        actual_epochs = if (!is.null(es_state) && !is.na(es_state$stopped_epoch)) {
            es_state$stopped_epoch
        } else {
            epochs
        }
        loss_history = loss_history[seq_len(actual_epochs)]
        val_loss_history = if (!is.null(val_loss_history)) val_loss_history[seq_len(actual_epochs)] else NULL
        stopped_epoch = if (!is.null(es_state) && !is.na(es_state$stopped_epoch)) {
            es_state$stopped_epoch
        } else {
            NA_integer_
        }

        # ---- Fitted values ----
        model$eval()
        x_full_t = .make_tensor(x)
        fitted_tensor = torch::with_no_grad(model(x_full_t))

        fitted_values = if (is_classification) {
            fitted_probs = torch::nnf_softmax(fitted_tensor, dim = 2L)
            fitted_classes = torch::torch_argmax(fitted_probs, dim = 2L)
            factor(
                as.integer(fitted_classes$cpu()),
                levels = seq_along(y_levels),
                labels = y_levels
            )
        } else {
            convert = as.matrix(fitted_tensor$cpu())
            if (no_y == 1L) {
                as.vector(convert)
            } else {
                convert
            }
        }

        # ---- Weight caching ----
        # For visualizing weight distribution
        # Allow `cache_weights` to generate
        cached_weights = if (cache_weights) {
            tryCatch(
                lapply(model$parameters, function(p) as.matrix(p$cpu())),
                error = function(e) {
                    cli::cli_warn("Weight caching failed: {conditionMessage(e)}")
                    NULL
                }
            )
        } else {
            NULL
        }

        structure(
            list(
                model = model,
                fitted = fitted_values,
                loss_history = loss_history,
                val_loss_history = val_loss_history,
                n_epochs = actual_epochs,
                stopped_epoch = stopped_epoch,
                hidden_neurons = hidden_neurons,
                activations = activations,
                output_activation = output_activation,
                penalty = penalty,
                mixture = mixture,
                feature_names = feature_names,
                response_name = response_name,
                no_x = no_x,
                no_y = no_y,
                is_classification = is_classification,
                y_levels = y_levels,
                n_classes = n_classes,
                device = device,
                cached_weights = cached_weights,
                arch = arch
            ),
            class = unique(c(fit_class, "nn_fit"))
        )
    }

#' Predict from a trained neural network
#'
#' @description
#' Generate predictions from an `"nn_fit"` object produced by [train_nn()].
#'
#' Three S3 methods are registered:
#'
#' - `predict.nn_fit()` — base method for `matrix`-trained models.
#' - `predict.nn_fit_tab()` — extends the base method for tabular fits; runs new
#'   data through `hardhat::forge()` before predicting.
#' - `predict.nn_fit_ds()` — extends the base method for torch `dataset` fits.
#'
#' @param object A fitted model object returned by [train_nn()].
#' @param newdata New predictor data. Accepted forms depend on the method:
#'   - `predict.nn_fit()`: a numeric `matrix` or coercible object.
#'   - `predict.nn_fit_tab()`: a `data.frame` with the same columns used during
#'     training; preprocessing is applied automatically via `hardhat::forge()`.
#'   - `predict.nn_fit_ds()`: a `torch` `dataset`, numeric `array`, `matrix`, or
#'     `data.frame`.
#'   If `NULL`, the cached fitted values from training are returned (not
#'   available for `type = "prob"`).
#' @param new_data Legacy alias for `newdata`. Retained for compatibility.
#' @param type Character. Output type:
#'   - `"response"` (default): predicted class labels (factor) for
#'     classification, or a numeric vector / matrix for regression.
#'   - `"prob"`: a numeric matrix of class probabilities (classification only).
#' @param ... Currently unused; reserved for future extensions.
#'
#' @return
#' - **Regression**: a numeric vector (single output) or matrix (multiple outputs).
#' - **Classification**, `type = "response"`: a factor with levels matching those
#'   seen during training.
#' - **Classification**, `type = "prob"`: a numeric matrix with one column per
#'   class, columns named by class label.
#'
#' @seealso [train_nn()]
#' @name gen-nn-predict
#' @export
predict.nn_fit =
    function(
        object,
        newdata = NULL,
        new_data = NULL,
        type = "response",
        ...
    ) {
        if (!requireNamespace("torch", quietly = TRUE)) {
            cli::cli_abort("Package {.pkg torch} is required but not installed.")
        }

        if (!type %in% c("response", "prob")) {
            cli::cli_abort("{.arg type} must be one of {.val response} or {.val prob}.")
        }

        newdata = .resolve_predict_newdata(newdata = newdata, new_data = new_data)

        device = object$device
        input_fn = if (!is.null(object$arch) && !is.null(object$arch$input_transform)) {
            rlang::as_function(object$arch$input_transform)
        } else {
            identity
        }

        # Early return is a must, when `newdata` / `new_data` is not used
        # Do not attempt to use `type = "prob"` when classification is not applied
        if (is.null(newdata)) {
            if (type == "prob" && object$is_classification) {
                cli::cli_abort("Cannot compute probabilities without {.arg newdata}. Use fitted values instead.")
            }
            if (type == "prob" && !object$is_classification) {
                cli::cli_abort("{.arg type = 'prob'} is only available for classification models.")
            }
            return(object$fitted)
        }

        if (!is.matrix(newdata)) newdata = as.matrix(newdata)

        x_new_t = input_fn(torch::torch_tensor(newdata, dtype = torch::torch_float32(), device = device))
        object$model$eval()
        pred_tensor = torch::with_no_grad(object$model(x_new_t))

        predictions = if (object$is_classification) {
            probs = torch::nnf_softmax(pred_tensor, dim = 2L)

            if (type == "prob") {
                prob_matrix = as.matrix(probs$cpu())
                colnames(prob_matrix) = object$y_levels
                prob_matrix
            } else {
                factor(
                    as.integer(torch::torch_argmax(probs, dim = 2L)$cpu()),
                    levels = seq_along(object$y_levels),
                    labels = object$y_levels
                )
            }
        } else {
            if (type == "prob") {
                cli::cli_abort("{.arg type = 'prob'} is only available for classification models.")
            }

            convert = as.matrix(pred_tensor$cpu())
            if (object$no_y == 1L) {
                as.vector(convert)
            } else {
                convert
            }
        }

        predictions
    }


#' @rdname gen-nn-predict
#' @export
predict.nn_fit_tab =
    function(
        object,
        newdata = NULL,
        new_data = NULL,
        type = "response",
        ...
    ) {
        newdata = .resolve_predict_newdata(newdata = newdata, new_data = new_data)

        if (!is.null(newdata) && !is.null(object$blueprint)) {
            processed = hardhat::forge(newdata, object$blueprint)
            newdata = as.matrix(processed$predictors)
        }

        NextMethod()
    }


#' @keywords internal
#' @export
`$.nn_fit` = function(x, name) {
    if (name %in% names(x)) return(x[[name]])
    attr(x, name, exact = TRUE)
}
