R/nn_model.R

In ggmlR: 'GGML' Tensor Operations for Machine Learning

# High-level sequential model API for ggmlR
# Provides Keras-like model building, compilation, training and evaluation

# ============================================================================
# Model Constructor
# ============================================================================

#' Create a Sequential Neural Network Model
#'
#' Creates an empty sequential model that layers can be added to using
#' pipe (\code{|>}) operators.
#'
#' @return A ggml_sequential_model object
#' @export
#' @examples
#' \dontrun{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_conv_2d(32, c(3,3), activation = "relu",
#'                      input_shape = c(28, 28, 1)) |>
#'   ggml_layer_max_pooling_2d(c(2,2)) |>
#'   ggml_layer_flatten() |>
#'   ggml_layer_dense(128, activation = "relu") |>
#'   ggml_layer_dense(10, activation = "softmax")
#' }
ggml_model_sequential <- function() {
  model <- list(
    layers = list(),
    input_shape = NULL,
    compiled = FALSE,
    compilation = list(
      ctx_weights = NULL,
      backend = NULL,
      sched = NULL,
      buffer = NULL,
      optimizer = NULL,
      loss = NULL,
      metrics = NULL
    )
  )
  class(model) <- c("ggml_sequential_model", "list")
  model
}

# ============================================================================
# Layer access and manipulation
# ============================================================================

#' Get a Layer from a Sequential Model
#'
#' Retrieves a layer by name or by integer index (1-based).
#'
#' @param model A ggml_sequential_model object
#' @param index Integer index of the layer (1-based), or NULL
#' @param name Character name of the layer, or NULL
#' @return The layer list object
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(64, activation = "relu", name = "hidden") |>
#'   ggml_layer_dense(10, activation = "softmax", name = "output")
#'
#' ggml_get_layer(model, index = 1)
#' ggml_get_layer(model, name = "output")
#' }
ggml_get_layer <- function(model, index = NULL, name = NULL) {
  if (is.null(index) && is.null(name)) {
    stop("Provide either index or name.")
  }
  if (!is.null(index) && !is.null(name)) {
    stop("Provide either index or name, not both.")
  }
  if (!is.null(index)) {
    n <- length(model$layers)
    if (index < 1L || index > n) {
      stop("Layer index ", index, " out of range (model has ", n, " layers).")
    }
    return(model$layers[[index]])
  }
  # by name
  for (layer in model$layers) {
    if (!is.null(layer$name) && layer$name == name) return(layer)
  }
  stop("No layer with name '", name, "'.")
}

#' Remove the Last Layer from a Sequential Model
#'
#' Removes the last layer from the model. The model must not be compiled.
#'
#' @param model A ggml_sequential_model object
#' @return The model with the last layer removed.
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(64, activation = "relu") |>
#'   ggml_layer_dense(10, activation = "softmax")
#'
#' model <- ggml_pop_layer(model)
#' length(model$layers)  # 1
#' }
ggml_pop_layer <- function(model) {
  if (model$compiled) {
    stop("Cannot pop layers from a compiled model.")
  }
  n <- length(model$layers)
  if (n == 0L) {
    stop("Model has no layers to remove.")
  }
  model$layers <- model$layers[-n]
  if (n == 1L) model$input_shape <- NULL
  model
}

# ============================================================================
# Freeze / unfreeze weights
# ============================================================================

#' Freeze Layer Weights
#'
#' Sets \code{trainable = FALSE} on layers, preventing their weights from being
#' updated during training. Accepts optional \code{from} / \code{to} to freeze
#' a range of layers by index, or \code{layer_names} to freeze by name.
#' If none are provided, all layers are frozen.
#'
#' @param model A model object (ggml_sequential_model or ggml_functional_model)
#' @param from Integer index of the first layer to freeze (default: 1)
#' @param to Integer index of the last layer to freeze (default: last layer)
#' @param layer_names Character vector of layer names to freeze (overrides from/to)
#' @param ... Additional arguments passed to methods
#' @return The model with selected layers frozen.
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(64, activation = "relu") |>
#'   ggml_layer_dense(10, activation = "softmax")
#'
#' # Freeze all layers
#' model <- ggml_freeze_weights(model)
#'
#' # Freeze only the first layer
#' model <- ggml_freeze_weights(model, from = 1, to = 1)
#' }
ggml_freeze_weights <- function(model, from = 1L, to = length(model$layers),
                                 layer_names = NULL, ...) {
  UseMethod("ggml_freeze_weights")
}

#' @export
ggml_freeze_weights.ggml_sequential_model <- function(model,
                                                        from = 1L,
                                                        to = length(model$layers),
                                                        layer_names = NULL, ...) {
  n <- length(model$layers)
  if (n == 0L) stop("Model has no layers.")
  if (!is.null(layer_names)) {
    names_all <- vapply(model$layers, function(l) if (!is.null(l$name)) l$name else "", character(1))
    idx <- which(names_all %in% layer_names)
    if (length(idx) == 0L) stop("No layers found with names: ", paste(layer_names, collapse = ", "))
    for (i in idx) model$layers[[i]]$trainable <- FALSE
  } else {
    from <- as.integer(from); to <- as.integer(to)
    if (from < 1L || to > n || from > to)
      stop("Invalid range: from=", from, " to=", to, " (model has ", n, " layers).")
    for (i in from:to) model$layers[[i]]$trainable <- FALSE
  }
  model
}

#' @export
ggml_freeze_weights.ggml_functional_model <- function(model,
                                                        from = NULL, to = NULL,
                                                        layer_names = NULL, ...) {
  nodes <- nn_topo_sort(model$outputs)
  frozen <- if (is.null(model$frozen_nodes)) list() else model$frozen_nodes
  if (!is.null(layer_names)) {
    for (node in nodes) {
      nm <- if (!is.null(node$config$name)) node$config$name else ""
      if (nm %in% layer_names) frozen[[node$id]] <- FALSE
    }
  } else {
    for (node in nodes) {
      if (!is.null(node$node_type) && node$node_type != "input")
        frozen[[node$id]] <- FALSE
    }
  }
  model$frozen_nodes <- frozen
  model
}

#' Unfreeze Layer Weights
#'
#' Sets \code{trainable = TRUE} on layers. Accepts optional \code{from} / \code{to}
#' to unfreeze a range of layers, or \code{layer_names} to unfreeze by name.
#' If none are provided, all layers are unfrozen.
#'
#' @param model A model object (ggml_sequential_model or ggml_functional_model)
#' @param from Integer index of the first layer to unfreeze (default: 1)
#' @param to Integer index of the last layer to unfreeze (default: last layer)
#' @param layer_names Character vector of layer names to unfreeze (overrides from/to)
#' @param ... Additional arguments passed to methods
#' @return The model with selected layers unfrozen.
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(64, activation = "relu") |>
#'   ggml_layer_dense(10, activation = "softmax")
#'
#' model <- ggml_freeze_weights(model)
#' model <- ggml_unfreeze_weights(model, from = 2)  # unfreeze last layer only
#' }
ggml_unfreeze_weights <- function(model, from = 1L, to = length(model$layers),
                                   layer_names = NULL, ...) {
  UseMethod("ggml_unfreeze_weights")
}

#' @export
ggml_unfreeze_weights.ggml_sequential_model <- function(model,
                                                          from = 1L,
                                                          to = length(model$layers),
                                                          layer_names = NULL, ...) {
  n <- length(model$layers)
  if (n == 0L) stop("Model has no layers.")
  if (!is.null(layer_names)) {
    names_all <- vapply(model$layers, function(l) if (!is.null(l$name)) l$name else "", character(1))
    idx <- which(names_all %in% layer_names)
    if (length(idx) == 0L) stop("No layers found with names: ", paste(layer_names, collapse = ", "))
    for (i in idx) model$layers[[i]]$trainable <- TRUE
  } else {
    from <- as.integer(from); to <- as.integer(to)
    if (from < 1L || to > n || from > to)
      stop("Invalid range: from=", from, " to=", to, " (model has ", n, " layers).")
    for (i in from:to) model$layers[[i]]$trainable <- TRUE
  }
  model
}

#' @export
ggml_unfreeze_weights.ggml_functional_model <- function(model,
                                                          from = NULL, to = NULL,
                                                          layer_names = NULL, ...) {
  frozen <- if (is.null(model$frozen_nodes)) list() else model$frozen_nodes
  if (!is.null(layer_names)) {
    nodes <- nn_topo_sort(model$outputs)
    ids_to_unfreeze <- vapply(nodes, function(n) {
      nm <- if (!is.null(n$config$name)) n$config$name else ""
      if (nm %in% layer_names) n$id else ""
    }, character(1))
    ids_to_unfreeze <- ids_to_unfreeze[nzchar(ids_to_unfreeze)]
    for (id in ids_to_unfreeze) frozen[[id]] <- NULL
  } else {
    frozen <- list()
  }
  model$frozen_nodes <- frozen
  model
}

# ============================================================================
# Compile
# ============================================================================

#' Compile a Sequential Model
#'
#' Configures the model for training: infers shapes, creates backend.
#' Weight tensors are created at training time when batch_size is known.
#'
#' @param model A ggml_sequential_model object
#' @param optimizer Optimizer name: "adam" or "sgd"
#' @param loss Loss function name: "categorical_crossentropy" or "mse"
#' @param metrics Character vector of metrics (currently "accuracy")
#' @param backend Backend to use: "auto" (GPU if available, else CPU), "cpu", or "vulkan"
#' @return The compiled model (invisibly).
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_conv_2d(32, c(3,3), activation = "relu",
#'                      input_shape = c(28, 28, 1)) |>
#'   ggml_layer_max_pooling_2d(c(2, 2)) |>
#'   ggml_layer_flatten() |>
#'   ggml_layer_dense(10, activation = "softmax")
#' model <- ggml_compile(model, optimizer = "adam",
#'                       loss = "categorical_crossentropy")
#' }
ggml_compile <- function(model, optimizer = "adam",
                          loss = "categorical_crossentropy",
                          metrics = c("accuracy"),
                          backend = "auto") {
  UseMethod("ggml_compile")
}

#' @rdname ggml_compile
#' @export
ggml_compile.ggml_sequential_model <- function(model, optimizer = "adam",
                                                loss = "categorical_crossentropy",
                                                metrics = c("accuracy"),
                                                backend = "auto") {
  if (length(model$layers) == 0) {
    stop("Model has no layers. Add layers before compiling.")
  }

  # 1. Shape inference
  model <- nn_infer_shapes(model)

  # 2. Create backend and scheduler
  use_vulkan <- FALSE
  if (backend == "auto") {
    if (ggml_vulkan_available() && ggml_vulkan_device_count() > 0) {
      use_vulkan <- TRUE
    }
  } else if (backend == "vulkan") {
    if (!ggml_vulkan_available() || ggml_vulkan_device_count() == 0) {
      stop("Vulkan backend requested but not available. ",
           "Install libvulkan-dev and glslc, then reinstall ggmlR")
    }
    use_vulkan <- TRUE
  } else if (backend != "cpu") {
    stop("Unknown backend: '", backend, "'. Use 'auto', 'cpu', or 'vulkan'.")
  }

  if (use_vulkan) {
    gpu_backend <- ggml_vulkan_init(0L)
    # sched_new auto-adds CPU as fallback for unsupported ops
    sched <- ggml_backend_sched_new(list(gpu_backend), parallel = FALSE)
    # Use separate CPU backend for weight allocation
    cpu_backend <- ggml_backend_cpu_init()
    if (!isTRUE(.ggmlr_state$backend_msg_shown)) {
      message("Using Vulkan GPU backend: ", ggml_vulkan_device_description(0L))
      .ggmlr_state$backend_msg_shown <- TRUE
    }
  } else {
    cpu_backend <- ggml_backend_cpu_init()
    sched <- ggml_backend_sched_new(list(cpu_backend), parallel = FALSE)
    if (!isTRUE(.ggmlr_state$backend_msg_shown)) {
      message("Using CPU backend")
      .ggmlr_state$backend_msg_shown <- TRUE
    }
  }

  # 3. Store compilation config (weights created later in fit/evaluate)
  # Weights go on GPU for performance (avoids per-iteration copies)
  if (use_vulkan) {
    model$compilation$backend <- gpu_backend
    model$compilation$cpu_backend <- cpu_backend
  } else {
    model$compilation$backend <- cpu_backend
  }
  model$compilation$sched <- sched
  model$compilation$optimizer <- optimizer
  model$compilation$loss <- loss
  model$compilation$metrics <- metrics
  # Record requested vs actually-used backend so a silent "auto" -> CPU
  # fallback is inspectable later (see ggml_model_backend()).
  model$compilation$backend_requested <- backend
  model$compilation$backend_used      <- if (use_vulkan) "vulkan" else "cpu"
  model$compilation$device <- if (use_vulkan) {
    ggml_vulkan_device_description(0L)
  } else "cpu"
  model$compiled <- TRUE

  invisible(model)
}

# ============================================================================
# Internal: Build graph with weights for a given batch_size
# ============================================================================

#' Build computation graph with allocated weights and inputs
#' @param model Compiled model
#' @param batch_size Batch size
#' @return List with ctx_weights, ctx_compute, inputs, outputs, buffer
#' @keywords internal
nn_build_graph <- function(model, batch_size, training = TRUE) {
  input_shape <- model$input_shape
  ne_datapoint <- prod(input_shape)
  backend <- model$compilation$backend

  # Count total parameters + input tensor size for memory estimation
  total_elements <- 0
  for (layer in model$layers) {
    total_elements <- total_elements + nn_count_layer_params(layer)
  }
  # Add input tensor
  total_elements <- total_elements + ne_datapoint * batch_size

  mem_size <- max((total_elements + 1000) * 4 + length(model$layers) * 2048,
                  2 * 1024 * 1024)
  ctx_weights <- ggml_init(mem_size, no_alloc = TRUE)

  # Create input tensor in ctx_weights (will be allocated with backend)
  if (length(input_shape) == 3) {
    # Image: R [H, W, C] -> ggml [W, H, C, N]
    inputs <- ggml_new_tensor_4d(ctx_weights, GGML_TYPE_F32,
                                  input_shape[2], input_shape[1],
                                  input_shape[3], batch_size)
  } else if (length(input_shape) == 2) {
    # Sequence: R [seq_len, input_size] -> ggml [input_size, seq_len, N]
    # (input_size is dim0 so each time step's features are contiguous)
    inputs <- ggml_new_tensor_3d(ctx_weights, GGML_TYPE_F32,
                                  input_shape[2], input_shape[1], batch_size)
  } else {
    # Flat vector: R [features] -> ggml [features, N]
    inputs <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                  ne_datapoint, batch_size)
  }
  ggml_set_name(inputs, "inputs")
  ggml_set_input(inputs)

  # Create weight tensors in ctx_weights
  layers_built <- model$layers
  for (i in seq_along(layers_built)) {
    layer <- layers_built[[i]]

    if (layer$type == "conv_1d") {
      k <- layer$config$kernel_size
      ic <- layer$input_shape[2]
      oc <- layer$config$filters

      # ggml conv_1d kernel: [K, IC, OC]
      layer$weights$kernel <- ggml_new_tensor_3d(ctx_weights, GGML_TYPE_F32,
                                                  k, ic, oc)
      layer$weights$bias <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, oc)
      ggml_set_name(layer$weights$kernel, paste0("conv1d_", i, "_kernel"))
      ggml_set_name(layer$weights$bias, paste0("conv1d_", i, "_bias"))

    } else if (layer$type == "conv_2d") {
      kh <- layer$config$kernel_size[1]
      kw <- layer$config$kernel_size[2]
      ic <- layer$input_shape[3]
      oc <- layer$config$filters

      layer$weights$kernel <- ggml_new_tensor_4d(ctx_weights, GGML_TYPE_F32,
                                                  kw, kh, ic, oc)
      layer$weights$bias <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, oc)
      ggml_set_name(layer$weights$kernel, paste0("conv_", i, "_kernel"))
      ggml_set_name(layer$weights$bias, paste0("conv_", i, "_bias"))

    } else if (layer$type == "dense") {
      fan_in <- if (length(layer$input_shape) == 1) layer$input_shape else prod(layer$input_shape)
      units <- layer$config$units

      layer$weights$weight <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                  fan_in, units)
      layer$weights$bias <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, units)
      ggml_set_name(layer$weights$weight, paste0("dense_", i, "_weight"))
      ggml_set_name(layer$weights$bias, paste0("dense_", i, "_bias"))

    } else if (layer$type == "batch_norm") {
      # Determine number of features for gamma/beta
      n_features <- if (length(layer$input_shape) == 1) layer$input_shape
                    else if (length(layer$input_shape) == 2) layer$input_shape[2]
                    else layer$input_shape[3]

      layer$weights$gamma <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_features)
      layer$weights$beta <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_features)
      ggml_set_name(layer$weights$gamma, paste0("bn_", i, "_gamma"))
      ggml_set_name(layer$weights$beta, paste0("bn_", i, "_beta"))

    } else if (layer$type == "lstm") {
      # input_shape: c(seq_len, input_size)
      input_sz <- layer$input_shape[2]
      units    <- layer$config$units
      nm       <- layer$name

      layer$weights$W_gates <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                    input_sz, 4L * units)
      layer$weights$U_gates <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                    units, 4L * units)
      layer$weights$b_gates <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32,
                                                    4L * units)
      layer$weights$h0      <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32,
                                                    units)
      layer$weights$c0      <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32,
                                                    units)
      ggml_set_name(layer$weights$W_gates, paste0(nm, "_W_gates"))
      ggml_set_name(layer$weights$U_gates, paste0(nm, "_U_gates"))
      ggml_set_name(layer$weights$b_gates, paste0(nm, "_b_gates"))
      ggml_set_name(layer$weights$h0,      paste0(nm, "_h0"))
      ggml_set_name(layer$weights$c0,      paste0(nm, "_c0"))

    } else if (layer$type == "gru") {
      input_sz <- layer$input_shape[2]
      units    <- layer$config$units
      nm       <- layer$name

      layer$weights$W_zh <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                 input_sz, 2L * units)
      layer$weights$U_zh <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                 units, 2L * units)
      layer$weights$b_zh <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32,
                                                 2L * units)
      layer$weights$W_n  <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                 input_sz, units)
      layer$weights$U_n  <- ggml_new_tensor_2d(ctx_weights, GGML_TYPE_F32,
                                                 units, units)
      layer$weights$b_n  <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, units)
      layer$weights$h0   <- ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, units)
      ggml_set_name(layer$weights$W_zh, paste0(nm, "_W_zh"))
      ggml_set_name(layer$weights$U_zh, paste0(nm, "_U_zh"))
      ggml_set_name(layer$weights$b_zh, paste0(nm, "_b_zh"))
      ggml_set_name(layer$weights$W_n,  paste0(nm, "_W_n"))
      ggml_set_name(layer$weights$U_n,  paste0(nm, "_U_n"))
      ggml_set_name(layer$weights$b_n,  paste0(nm, "_b_n"))
      ggml_set_name(layer$weights$h0,   paste0(nm, "_h0"))
    }

    layers_built[[i]] <- layer
  }

  # Allocate all tensors in ctx_weights (inputs + weights) via backend
  buffer <- ggml_backend_alloc_ctx_tensors(ctx_weights, backend)

  # Initialize weights: prefer weights_data (R vectors from load), then trained
  # tensor weights, then random init
  for (i in seq_along(layers_built)) {
    layer <- layers_built[[i]]
    old_layer <- model$layers[[i]]
    has_weights_data <- !is.null(old_layer$weights_data)
    has_trained_weights <- !is.null(old_layer$weights) &&
      (!is.null(old_layer$weights$kernel) || !is.null(old_layer$weights$weight) ||
       !is.null(old_layer$weights$gamma))

    if (layer$type == "conv_1d") {
      if (has_weights_data && !is.null(old_layer$weights_data$kernel)) {
        ggml_backend_tensor_set_data(layer$weights$kernel, old_layer$weights_data$kernel)
        ggml_backend_tensor_set_data(layer$weights$bias, old_layer$weights_data$bias)
      } else if (has_trained_weights && !is.null(old_layer$weights$kernel)) {
        kernel_data <- ggml_backend_tensor_get_data(old_layer$weights$kernel)
        ggml_backend_tensor_set_data(layer$weights$kernel, kernel_data)
        bias_data <- ggml_backend_tensor_get_data(old_layer$weights$bias)
        ggml_backend_tensor_set_data(layer$weights$bias, bias_data)
      } else {
        k <- layer$config$kernel_size
        fan_in <- k * layer$input_shape[2]
        nn_init_he_uniform(layer$weights$kernel, fan_in)
        nn_init_zeros(layer$weights$bias)
      }
      if (isTRUE(layer$trainable)) {
        ggml_set_param(layer$weights$kernel)
        ggml_set_param(layer$weights$bias)
      }

    } else if (layer$type == "conv_2d") {
      if (has_weights_data && !is.null(old_layer$weights_data$kernel)) {
        # Load from R vectors (save/load path)
        ggml_backend_tensor_set_data(layer$weights$kernel, old_layer$weights_data$kernel)
        ggml_backend_tensor_set_data(layer$weights$bias, old_layer$weights_data$bias)
      } else if (has_trained_weights && !is.null(old_layer$weights$kernel)) {
        # Copy trained weights from tensors
        kernel_data <- ggml_backend_tensor_get_data(old_layer$weights$kernel)
        ggml_backend_tensor_set_data(layer$weights$kernel, kernel_data)
        bias_data <- ggml_backend_tensor_get_data(old_layer$weights$bias)
        ggml_backend_tensor_set_data(layer$weights$bias, bias_data)
      } else {
        # Random initialization
        kh <- layer$config$kernel_size[1]
        kw <- layer$config$kernel_size[2]
        fan_in <- kw * kh * layer$input_shape[3]
        nn_init_he_uniform(layer$weights$kernel, fan_in)
        nn_init_zeros(layer$weights$bias)
      }
      if (isTRUE(layer$trainable)) {
        ggml_set_param(layer$weights$kernel)
        ggml_set_param(layer$weights$bias)
      }

    } else if (layer$type == "dense") {
      if (has_weights_data && !is.null(old_layer$weights_data$weight)) {
        # Load from R vectors (save/load path)
        ggml_backend_tensor_set_data(layer$weights$weight, old_layer$weights_data$weight)
        ggml_backend_tensor_set_data(layer$weights$bias, old_layer$weights_data$bias)
      } else if (has_trained_weights && !is.null(old_layer$weights$weight)) {
        # Copy trained weights from tensors
        weight_data <- ggml_backend_tensor_get_data(old_layer$weights$weight)
        ggml_backend_tensor_set_data(layer$weights$weight, weight_data)
        bias_data <- ggml_backend_tensor_get_data(old_layer$weights$bias)
        ggml_backend_tensor_set_data(layer$weights$bias, bias_data)
      } else {
        # Random initialization
        fan_in <- if (length(layer$input_shape) == 1) layer$input_shape else prod(layer$input_shape)
        fan_out <- layer$config$units
        nn_init_glorot_uniform(layer$weights$weight, fan_in, fan_out)
        nn_init_zeros(layer$weights$bias)
      }
      if (isTRUE(layer$trainable)) {
        ggml_set_param(layer$weights$weight)
        ggml_set_param(layer$weights$bias)
      }

    } else if (layer$type == "batch_norm") {
      if (has_weights_data && !is.null(old_layer$weights_data$gamma)) {
        ggml_backend_tensor_set_data(layer$weights$gamma, old_layer$weights_data$gamma)
        ggml_backend_tensor_set_data(layer$weights$beta, old_layer$weights_data$beta)
      } else if (has_trained_weights && !is.null(old_layer$weights$gamma)) {
        gamma_data <- ggml_backend_tensor_get_data(old_layer$weights$gamma)
        ggml_backend_tensor_set_data(layer$weights$gamma, gamma_data)
        beta_data <- ggml_backend_tensor_get_data(old_layer$weights$beta)
        ggml_backend_tensor_set_data(layer$weights$beta, beta_data)
      } else {
        # gamma=1, beta=0
        n <- ggml_nelements(layer$weights$gamma)
        ggml_backend_tensor_set_data(layer$weights$gamma, rep(1.0, n))
        nn_init_zeros(layer$weights$beta)
      }
      if (isTRUE(layer$trainable)) {
        ggml_set_param(layer$weights$gamma)
        ggml_set_param(layer$weights$beta)
      }

    } else if (layer$type == "lstm") {
      units    <- layer$config$units
      input_sz <- layer$input_shape[2]
      # Restore or init W_gates
      if (!is.null(old_layer$weights$W_gates)) {
        ggml_backend_tensor_set_data(layer$weights$W_gates,
          ggml_backend_tensor_get_data(old_layer$weights$W_gates))
        ggml_backend_tensor_set_data(layer$weights$U_gates,
          ggml_backend_tensor_get_data(old_layer$weights$U_gates))
        ggml_backend_tensor_set_data(layer$weights$b_gates,
          ggml_backend_tensor_get_data(old_layer$weights$b_gates))
      } else if (!is.null(old_layer$weights_data$W_gates)) {
        ggml_backend_tensor_set_data(layer$weights$W_gates, old_layer$weights_data$W_gates)
        ggml_backend_tensor_set_data(layer$weights$U_gates, old_layer$weights_data$U_gates)
        ggml_backend_tensor_set_data(layer$weights$b_gates, old_layer$weights_data$b_gates)
      } else {
        nn_init_recurrent_uniform(layer$weights$W_gates)
        nn_init_recurrent_uniform(layer$weights$U_gates)
        nn_init_zeros(layer$weights$b_gates)
      }
      nn_init_zeros(layer$weights$h0)
      nn_init_zeros(layer$weights$c0)
      if (isTRUE(layer$trainable)) {
        ggml_set_param(layer$weights$W_gates)
        ggml_set_param(layer$weights$U_gates)
        ggml_set_param(layer$weights$b_gates)
      }

    } else if (layer$type == "gru") {
      units    <- layer$config$units
      input_sz <- layer$input_shape[2]
      if (!is.null(old_layer$weights$W_zh)) {
        ggml_backend_tensor_set_data(layer$weights$W_zh,
          ggml_backend_tensor_get_data(old_layer$weights$W_zh))
        ggml_backend_tensor_set_data(layer$weights$U_zh,
          ggml_backend_tensor_get_data(old_layer$weights$U_zh))
        ggml_backend_tensor_set_data(layer$weights$b_zh,
          ggml_backend_tensor_get_data(old_layer$weights$b_zh))
        ggml_backend_tensor_set_data(layer$weights$W_n,
          ggml_backend_tensor_get_data(old_layer$weights$W_n))
        ggml_backend_tensor_set_data(layer$weights$U_n,
          ggml_backend_tensor_get_data(old_layer$weights$U_n))
        ggml_backend_tensor_set_data(layer$weights$b_n,
          ggml_backend_tensor_get_data(old_layer$weights$b_n))
      } else if (!is.null(old_layer$weights_data$W_zh)) {
        ggml_backend_tensor_set_data(layer$weights$W_zh, old_layer$weights_data$W_zh)
        ggml_backend_tensor_set_data(layer$weights$U_zh, old_layer$weights_data$U_zh)
        ggml_backend_tensor_set_data(layer$weights$b_zh, old_layer$weights_data$b_zh)
        ggml_backend_tensor_set_data(layer$weights$W_n,  old_layer$weights_data$W_n)
        ggml_backend_tensor_set_data(layer$weights$U_n,  old_layer$weights_data$U_n)
        ggml_backend_tensor_set_data(layer$weights$b_n,  old_layer$weights_data$b_n)
      } else {
        nn_init_recurrent_uniform(layer$weights$W_zh)
        nn_init_recurrent_uniform(layer$weights$U_zh)
        nn_init_zeros(layer$weights$b_zh)
        nn_init_recurrent_uniform(layer$weights$W_n)
        nn_init_recurrent_uniform(layer$weights$U_n)
        nn_init_zeros(layer$weights$b_n)
      }
      nn_init_zeros(layer$weights$h0)
      if (isTRUE(layer$trainable)) {
        ggml_set_param(layer$weights$W_zh)
        ggml_set_param(layer$weights$U_zh)
        ggml_set_param(layer$weights$b_zh)
        ggml_set_param(layer$weights$W_n)
        ggml_set_param(layer$weights$U_n)
        ggml_set_param(layer$weights$b_n)
      }
    }
  }

  # Create compute context for intermediate tensors (no_alloc, ggml_opt manages)
  compute_mem <- max(64 * 1024 * 1024,
                     ne_datapoint * batch_size * 4 * 20)
  ctx_compute <- ggml_init(compute_mem, no_alloc = TRUE)

  # Build forward graph
  current <- inputs
  for (i in seq_along(layers_built)) {
    current <- nn_build_layer(ctx_compute, current, layers_built[[i]],
                              training = training)
  }
  outputs <- current
  ggml_set_output(outputs)

  list(
    ctx_weights = ctx_weights,
    ctx_compute = ctx_compute,
    inputs = inputs,
    outputs = outputs,
    buffer = buffer,
    layers_built = layers_built
  )
}

# ============================================================================
# Fit (Training)
# ============================================================================

#' Train a Model (dispatcher)
#'
#' Dispatcher: if the first argument is a \code{ggml_sequential_model}, delegates
#' to the Keras-style high-level API (\code{ggml_fit_sequential}); otherwise
#' delegates to the low-level optimizer loop (\code{ggml_fit_opt}).
#'
#' \strong{Keras-style (Sequential model):}
#' \describe{
#'   \item{model}{A compiled \code{ggml_sequential_model}}
#'   \item{x}{Training data (matrix or array)}
#'   \item{y}{Training labels (matrix, one-hot encoded for classification)}
#'   \item{epochs}{Number of training epochs (default: 1)}
#'   \item{batch_size}{Batch size (default: 32)}
#'   \item{validation_split}{Fraction of data for validation (default: 0)}
#'   \item{validation_data}{Optional list(x_val, y_val) for validation. Overrides validation_split.}
#'   \item{class_weight}{Named vector of weights per class, e.g. c("0"=1, "1"=10). Cannot be used with sample_weight.}
#'   \item{sample_weight}{Numeric vector of per-sample weights (length = nrow(x)). Cannot be used with class_weight.}
#'   \item{verbose}{0 = silent, 1 = progress (default: 1)}
#' }
#'
#' \strong{Low-level (optimizer loop):}
#' \describe{
#'   \item{sched}{Backend scheduler}
#'   \item{ctx_compute}{Compute context}
#'   \item{inputs}{Input tensor}
#'   \item{outputs}{Output tensor}
#'   \item{dataset}{Dataset from \code{ggml_opt_dataset_init()}}
#'   \item{loss_type}{Loss type (default: MSE)}
#'   \item{optimizer}{Optimizer type (default: AdamW)}
#'   \item{nepoch}{Number of epochs (default: 10)}
#'   \item{nbatch_logical}{Logical batch size (default: 32)}
#'   \item{val_split}{Validation fraction (default: 0)}
#'   \item{callbacks}{List of callback objects}
#'   \item{silent}{Suppress output (default: FALSE)}
#' }
#'
#' @param ... Arguments passed to the appropriate implementation.
#' @return For Sequential models: the trained model (invisibly).
#'   For the low-level API: a data frame with columns
#'   \code{epoch}, \code{train_loss}, \code{train_accuracy},
#'   \code{val_loss}, \code{val_accuracy}.
#' @seealso \code{\link{ggml_fit_opt}}, \code{\link{ggml_compile}}
#' @examples
#' \donttest{
#' ggml_set_n_threads(1L)  # deterministic, single OpenMP pool
#' n <- 128
#' x <- matrix(runif(n * 4), nrow = n, ncol = 4)
#' y <- matrix(0, nrow = n, ncol = 2)
#' for (i in seq_len(n)) { y[i, if (sum(x[i,]) > 2) 1L else 2L] <- 1 }
#'
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(8, activation = "relu") |>
#'   ggml_layer_dense(2, activation = "softmax")
#' model$input_shape <- 4L
#' model <- ggml_compile(model, optimizer = "adam",
#'                       loss = "categorical_crossentropy")
#'
#' # Basic training
#' model <- ggml_fit(model, x, y, epochs = 5, batch_size = 32, verbose = 0)
#'
#' # With validation_data
#' x_val <- matrix(runif(32 * 4), nrow = 32, ncol = 4)
#' y_val <- matrix(0, nrow = 32, ncol = 2)
#' for (i in seq_len(32)) { y_val[i, if (sum(x_val[i,]) > 2) 1L else 2L] <- 1 }
#' model <- ggml_fit(model, x, y, epochs = 3, batch_size = 32,
#'                   validation_data = list(x_val, y_val), verbose = 0)
#'
#' # With class_weight (useful for imbalanced classes)
#' model <- ggml_fit(model, x, y, epochs = 3, batch_size = 32,
#'                   class_weight = c("0" = 1, "1" = 2), verbose = 0)
#'
#' # With sample_weight
#' sw <- runif(n, 0.5, 1.5)
#' model <- ggml_fit(model, x, y, epochs = 3, batch_size = 32,
#'                   sample_weight = sw, verbose = 0)
#' }
#' @export
ggml_fit <- function(model, ...) {
  # Detect low-level call: all keyword args, no positional 'model'
  # (sched = ..., ctx_compute = ..., etc.)
  if (missing(model)) {
    return(do.call(ggml_fit_opt, list(...)))
  }
  UseMethod("ggml_fit")
}

#' @rdname ggml_fit
#' @export
ggml_fit.ggml_sequential_model <- function(model, ...) {
  ggml_fit_sequential(model, ...)
}

#' @rdname ggml_fit
#' @export
ggml_fit.default <- function(model, ...) {
  # Positional first arg — forward to low-level optimizer loop
  ggml_fit_opt(model, ...)
}

ggml_fit_sequential <- function(model, x, y, epochs = 1, batch_size = 32,
                                validation_split = 0.0, validation_data = NULL,
                                class_weight = NULL, sample_weight = NULL,
                                verbose = 1, callbacks = list()) {
  if (!model$compiled) {
    stop("Model must be compiled before training. Call ggml_compile() first.")
  }

  # Apply class_weight / sample_weight by scaling labels (one-hot)
  # For cross-entropy: CE(p, w*y_onehot) = w * CE(p, y_onehot)
  if (!is.null(class_weight) && !is.null(sample_weight)) {
    stop("Specify either class_weight or sample_weight, not both.")
  }
  if (!is.null(class_weight)) {
    # class_weight: named vector or list, e.g. c("0"=1, "1"=10) or list("0"=1, "1"=10)
    # y must be one-hot; true class = argmax per row
    true_classes <- max.col(y) - 1L  # 0-based
    keys <- as.character(true_classes)
    w <- unlist(class_weight)[keys]
    if (anyNA(w)) stop("class_weight must cover all classes present in y.")
    y <- y * w  # scale each row
  }
  # When TRUE, sample weights are applied through the weighted-MSE loss node
  # (sum(w*(pred-y)^2)) instead of by scaling the labels, which would be wrong
  # for squared error. CE is unaffected: CE(p, w*y) == w*CE(p, y).
  use_weighted_mse <- FALSE
  if (!is.null(sample_weight)) {
    if (length(sample_weight) != nrow(y)) {
      stop("sample_weight length must match number of training samples.")
    }
    is_mse <- model$compilation$loss %in% c("mse", "mean_squared_error")
    if (is_mse) {
      use_weighted_mse <- TRUE   # weights go into the loss node below, not y
    } else {
      y <- y * sample_weight     # cross-entropy: scaling the label is correct
    }
  }

  if (!is.null(validation_data)) {
    if (!is.list(validation_data) || length(validation_data) < 2) {
      stop("validation_data must be a list: list(x_val, y_val)")
    }
    x_val <- validation_data[[1]]
    y_val <- validation_data[[2]]
    n_val <- if (is.matrix(x_val)) nrow(x_val) else dim(x_val)[1]
    n_train <- if (is.matrix(x)) nrow(x) else dim(x)[1]
    # Combine train + val; val_split = fraction of combined that is val
    if (is.matrix(x)) {
      x <- rbind(x, x_val)
    } else {
      x <- abind_first(x, x_val)
    }
    y <- rbind(y, y_val)
    # Validation rows carry no user weight; default to 1.0 so weighted MSE
    # reduces to plain MSE on the validation split.
    if (use_weighted_mse) {
      sample_weight <- c(sample_weight, rep(1.0, n_val))
    }
    validation_split <- n_val / (n_train + n_val)
  }

  input_shape <- model$input_shape
  n_samples <- if (is.matrix(x)) nrow(x) else dim(x)[1]
  ne_datapoint <- prod(input_shape)
  ne_label <- ncol(y)

  # Ensure batch_size divides data evenly
  usable_samples <- (n_samples %/% batch_size) * batch_size
  if (usable_samples < n_samples) {
    dropped <- n_samples - usable_samples
    if (verbose > 0) {
      message("Note: dropping last ", dropped, " sample(s) (", n_samples,
              " -> ", usable_samples, ") because batch_size=", batch_size,
              " must divide evenly. Training metrics are computed on ",
              usable_samples, " samples only.")
    }
    x <- slice_first_dim(x, seq_len(usable_samples))
    y <- y[seq_len(usable_samples), , drop = FALSE]
    if (use_weighted_mse) {
      sample_weight <- sample_weight[seq_len(usable_samples)]
    }
    n_samples <- usable_samples
  }

  # Prepare dataset
  dataset <- ggml_opt_dataset_init(
    type_data = GGML_TYPE_F32,
    type_label = GGML_TYPE_F32,
    ne_datapoint = ne_datapoint,
    ne_label = ne_label,
    ndata = n_samples,
    ndata_shard = 1
  )

  # Convert data to ggml format
  if (length(input_shape) == 3) {
    # Image data: R [N, H, W, C] -> ggml [W, H, C, N]
    x_ggml <- as.vector(aperm(x, c(3, 2, 4, 1)))
  } else if (length(input_shape) == 2) {
    # Sequence: R [N, seq_len, input_size] -> ggml [input_size, seq_len, N]
    x_ggml <- as.vector(aperm(x, c(3, 2, 1)))
  } else if (length(input_shape) == 1) {
    # Vector data: R [N, features] -> ggml [features, N]
    x_ggml <- as.vector(t(x))
  } else {
    stop("Unsupported input_shape length: ", length(input_shape))
  }

  # Labels: R [N, classes] -> ggml [classes, N]
  y_ggml <- as.vector(t(y))

  # Fill dataset
  data_tensor <- ggml_opt_dataset_data(dataset)
  labels_tensor <- ggml_opt_dataset_labels(dataset)
  ggml_backend_tensor_set_data(data_tensor, x_ggml)
  ggml_backend_tensor_set_data(labels_tensor, y_ggml)

  # Per-datapoint weights for weighted MSE (loss = sum(w*(pred-y)^2)/nelem)
  if (use_weighted_mse) {
    weights_tensor <- ggml_opt_dataset_weights(dataset)
    ggml_backend_tensor_set_data(weights_tensor, as.numeric(sample_weight))
  }

  # Build graph (creates contexts, weights, inputs, outputs)
  graph_info <- nn_build_graph(model, batch_size)

  # Map optimizer and loss
  optimizer_type <- switch(model$compilation$optimizer,
    "adam" = , "adamw" = ggml_opt_optimizer_type_adamw(),
    "sgd" = ggml_opt_optimizer_type_sgd(),
    ggml_opt_optimizer_type_adamw()
  )

  loss_type <- if (use_weighted_mse) {
    ggml_opt_loss_type_weighted_mse()
  } else {
    switch(model$compilation$loss,
      "categorical_crossentropy" = , "crossentropy" = ggml_opt_loss_type_cross_entropy(),
      "mse" = , "mean_squared_error" = ggml_opt_loss_type_mse(),
      ggml_opt_loss_type_cross_entropy()
    )
  }

  # Train (returns history list from C)
  history_raw <- ggml_opt_fit(
    sched = model$compilation$sched,
    ctx_compute = graph_info$ctx_compute,
    inputs = graph_info$inputs,
    outputs = graph_info$outputs,
    dataset = dataset,
    loss_type = loss_type,
    optimizer = optimizer_type,
    nepoch = epochs,
    nbatch_logical = batch_size,
    val_split = validation_split,
    silent = (verbose == 0)
  )

  # Store built layers (with trained weights)
  model$layers <- graph_info$layers_built
  model$compilation$ctx_weights <- graph_info$ctx_weights
  model$compilation$buffer <- graph_info$buffer

  # Build history object
  model$history <- structure(
    list(
      train_loss     = history_raw$train_loss,
      train_accuracy = history_raw$train_accuracy,
      val_loss       = history_raw$val_loss,
      val_accuracy   = history_raw$val_accuracy,
      epochs         = seq_len(epochs)
    ),
    class = "ggml_history"
  )

  # Cleanup
  ggml_free(graph_info$ctx_compute)
  ggml_opt_dataset_free(dataset)

  invisible(model)
}

# ============================================================================
# Evaluate
# ============================================================================

#' Evaluate a Trained Model
#'
#' @param model A trained ggml_sequential_model
#' @param x Test data
#' @param y Test labels (one-hot encoded)
#' @param batch_size Batch size for evaluation
#' @param sample_weight Numeric vector of per-sample weights (length = nrow(x)).
#' @param class_weight Named vector of weights per class, e.g. c("0"=1, "1"=10). Cannot be used with sample_weight.
#' @return Named list with \code{loss} and \code{accuracy}.
#' @examples
#' \donttest{
#' ggml_set_n_threads(1L)  # deterministic, single OpenMP pool
#' n <- 128
#' x <- matrix(runif(n * 4), nrow = n, ncol = 4)
#' y <- matrix(0, nrow = n, ncol = 2)
#' for (i in seq_len(n)) { y[i, if (sum(x[i,]) > 2) 1L else 2L] <- 1 }
#'
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(8, activation = "relu") |>
#'   ggml_layer_dense(2, activation = "softmax")
#' model$input_shape <- 4L
#' model <- ggml_compile(model, optimizer = "adam",
#'                       loss = "categorical_crossentropy")
#' model <- ggml_fit(model, x, y, epochs = 5, batch_size = 32, verbose = 0)
#'
#' # Basic evaluation
#' result <- ggml_evaluate(model, x, y, batch_size = 32)
#'
#' # With sample_weight
#' sw <- runif(n, 0.5, 1.5)
#' result <- ggml_evaluate(model, x, y, batch_size = 32, sample_weight = sw)
#'
#' # With class_weight
#' result <- ggml_evaluate(model, x, y, batch_size = 32,
#'                         class_weight = c("0" = 1, "1" = 2))
#' }
#' @export
ggml_evaluate <- function(model, ...) {
  UseMethod("ggml_evaluate")
}

#' @rdname ggml_evaluate
#' @export
ggml_evaluate.ggml_sequential_model <- function(model, x, y, batch_size = 32,
                                                  sample_weight = NULL,
                                                  class_weight = NULL, ...) {
  if (!model$compiled) {
    stop("Model must be compiled before evaluation.")
  }

  if (!is.null(class_weight) && !is.null(sample_weight)) {
    stop("Specify either class_weight or sample_weight, not both.")
  }
  if (!is.null(class_weight)) {
    true_classes <- max.col(y) - 1L
    keys <- as.character(true_classes)
    w <- unlist(class_weight)[keys]
    if (anyNA(w)) stop("class_weight must cover all classes present in y.")
    y <- y * w
  }
  if (!is.null(sample_weight)) {
    if (length(sample_weight) != nrow(y)) {
      stop("sample_weight length must match number of samples.")
    }
    y <- y * sample_weight
  }

  n_samples <- nrow(y)
  ne_label <- ncol(y)

  # Get predictions for ALL samples (no truncation)
  preds <- ggml_predict(model, x, batch_size = batch_size)

  # Compute loss
  loss_name <- model$compilation$loss
  if (loss_name %in% c("categorical_crossentropy", "crossentropy")) {
    # Cross-entropy: -sum(y * log(p)) / n
    eps <- 1e-7
    preds_clipped <- pmax(pmin(preds, 1 - eps), eps)
    loss_val <- -mean(rowSums(y * log(preds_clipped)))
  } else if (loss_name %in% c("mse", "mean_squared_error")) {
    loss_val <- mean(rowSums((y - preds)^2) / ne_label)
  } else {
    loss_val <- NA_real_
  }

  # Compute accuracy (classification: argmax match)
  if (ne_label > 1L) {
    pred_classes <- max.col(preds)
    true_classes <- max.col(y)
    acc_val <- mean(pred_classes == true_classes)
  } else {
    acc_val <- NA_real_
  }

  list(loss = loss_val, accuracy = acc_val, n_samples = n_samples)
}

# ============================================================================
# Predict
# ============================================================================

#' Get Predictions from a Trained Model
#'
#' Runs forward pass on input data and returns prediction probabilities
#' (or raw output values for regression). Unlike \code{ggml_evaluate()}, this
#' does not require labels.
#'
#' @param model A trained ggml_sequential_model
#' @param x Input data (matrix or array)
#' @param batch_size Batch size for inference
#' @return Matrix of predictions with shape \code{[N, output_units]}
#' @export
ggml_predict <- function(model, ...) {
  UseMethod("ggml_predict")
}

#' @rdname ggml_predict
#' @export
ggml_predict.ggml_sequential_model <- function(model, x, batch_size = 32L, ...) {
  if (!model$compiled) {
    stop("Model must be compiled before prediction.")
  }

  input_shape <- model$input_shape
  n_samples <- if (is.matrix(x)) nrow(x) else dim(x)[1]
  ne_datapoint <- prod(input_shape)

  # Get output size from last layer
  last_layer <- model$layers[[length(model$layers)]]
  ne_output <- if (length(last_layer$output_shape) == 1) {
    last_layer$output_shape
  } else {
    prod(last_layer$output_shape)
  }

  if (n_samples == 0L) {
    stop("No samples provided.")
  }

  # Convert all data to ggml format upfront (before any slicing)
  nn_convert_to_ggml <- function(x, input_shape) {
    if (length(input_shape) == 3) {
      as.vector(aperm(x, c(3, 2, 4, 1)))
    } else if (length(input_shape) == 2) {
      as.vector(aperm(x, c(3, 2, 1)))
    } else {
      as.vector(t(x))
    }
  }

  # Helper: run forward pass for a batch of given size, return prediction matrix
  nn_predict_batch_run <- function(model, x_ggml, bs, ne_datapoint, ne_output,
                                   n_batches, input_shape) {
    graph_info <- nn_build_graph(model, bs, training = FALSE)
    graph <- ggml_build_forward_expand(graph_info$ctx_compute, graph_info$outputs)
    sched <- model$compilation$sched
    preds <- matrix(0, nrow = n_batches * bs, ncol = ne_output)

    for (ib in seq_len(n_batches)) {
      data_start <- (ib - 1L) * bs * ne_datapoint + 1L
      data_end <- ib * bs * ne_datapoint
      batch_data <- x_ggml[data_start:data_end]
      ggml_backend_tensor_set_data(graph_info$inputs, batch_data)

      ggml_backend_sched_reset(sched)
      ggml_backend_sched_alloc_graph(sched, graph)
      ggml_backend_sched_graph_compute(sched, graph)

      batch_output <- ggml_backend_tensor_get_data(graph_info$outputs)
      batch_matrix <- matrix(batch_output, nrow = ne_output, ncol = bs)
      row_start <- (ib - 1L) * bs + 1L
      row_end <- ib * bs
      preds[row_start:row_end, ] <- t(batch_matrix)
    }

    ggml_free(graph_info$ctx_compute)
    ggml_backend_buffer_free(graph_info$buffer)
    ggml_free(graph_info$ctx_weights)

    preds
  }

  n_full <- (n_samples %/% batch_size) * batch_size
  remainder <- n_samples - n_full

  all_preds <- matrix(0, nrow = n_samples, ncol = ne_output)

  # Main batches
  if (n_full > 0L) {
    x_main <- slice_first_dim(x, seq_len(n_full))
    x_ggml_main <- nn_convert_to_ggml(x_main, input_shape)
    n_batches <- n_full %/% batch_size
    all_preds[seq_len(n_full), ] <- nn_predict_batch_run(
      model, x_ggml_main, batch_size, ne_datapoint, ne_output,
      n_batches, input_shape
    )
  }

  # Remainder batch with rebuilt graph
  if (remainder > 0L) {
    idx_rem <- (n_full + 1L):n_samples
    x_rem <- slice_first_dim(x, idx_rem)
    x_ggml_rem <- nn_convert_to_ggml(x_rem, input_shape)
    all_preds[idx_rem, ] <- nn_predict_batch_run(
      model, x_ggml_rem, remainder, ne_datapoint, ne_output,
      1L, input_shape
    )
  }

  all_preds
}

#' Predict Classes from a Trained Model
#'
#' Returns predicted class indices (1-based) by applying argmax
#' to the output of \code{ggml_predict()}.
#'
#' @param model A trained ggml_sequential_model
#' @param x Input data (matrix or array)
#' @param batch_size Batch size for inference
#' @return Integer vector of predicted class indices (1-based)
#' @export
ggml_predict_classes <- function(model, x, batch_size = 32L) {
  probs <- ggml_predict(model, x, batch_size)
  apply(probs, 1, which.max)
}

# ============================================================================
# Save / Load Weights
# ============================================================================

#' Save Model Weights to File
#'
#' Saves the trained weights of a sequential model to an RDS file.
#' The file includes both weights and architecture metadata for validation
#' when loading.
#'
#' @param model A trained ggml_sequential_model
#' @param path File path to save weights (typically with .rds extension)
#' @return The model (invisibly).
#' @export
ggml_save_weights <- function(model, path) {
  if (!model$compiled) {
    stop("Model must be compiled before saving weights.")
  }

  # Extract weights as R vectors from each layer
  weights_list <- list()
  for (i in seq_along(model$layers)) {
    layer <- model$layers[[i]]
    layer_weights <- list()

    if (layer$type %in% c("conv_1d", "conv_2d")) {
      if (!is.null(layer$weights$kernel)) {
        layer_weights$kernel <- ggml_backend_tensor_get_data(layer$weights$kernel)
        layer_weights$bias <- ggml_backend_tensor_get_data(layer$weights$bias)
      }
    } else if (layer$type == "dense") {
      if (!is.null(layer$weights$weight)) {
        layer_weights$weight <- ggml_backend_tensor_get_data(layer$weights$weight)
        layer_weights$bias <- ggml_backend_tensor_get_data(layer$weights$bias)
      }
    } else if (layer$type == "batch_norm") {
      if (!is.null(layer$weights$gamma)) {
        layer_weights$gamma <- ggml_backend_tensor_get_data(layer$weights$gamma)
        layer_weights$beta <- ggml_backend_tensor_get_data(layer$weights$beta)
      }
    }

    weights_list[[i]] <- layer_weights
  }

  # Build architecture description for validation on load
  architecture <- list(
    input_shape = model$input_shape,
    n_layers = length(model$layers),
    layer_configs = lapply(model$layers, function(l) {
      list(type = l$type, config = l$config,
           input_shape = l$input_shape, output_shape = l$output_shape)
    })
  )

  data <- list(
    weights = weights_list,
    architecture = architecture,
    version = 1L
  )

  saveRDS(data, path)
  invisible(model)
}

#' Load Model Weights from File
#'
#' Loads previously saved weights into a compiled model. The model architecture
#' must match the saved weights (same layer types, sizes, and shapes).
#'
#' @param model A compiled ggml_sequential_model (same architecture as saved)
#' @param path File path to load weights from
#' @return The model with loaded weights.
#' @export
ggml_load_weights <- function(model, path) {
  if (!model$compiled) {
    stop("Model must be compiled before loading weights.")
  }

  data <- readRDS(path)

  if (is.null(data$weights) || is.null(data$architecture)) {
    stop("Invalid weights file format.")
  }

  # Validate architecture match
  arch <- data$architecture
  if (length(model$layers) != arch$n_layers) {
    stop("Architecture mismatch: model has ", length(model$layers),
         " layers, saved weights have ", arch$n_layers)
  }

  for (i in seq_along(model$layers)) {
    if (model$layers[[i]]$type != arch$layer_configs[[i]]$type) {
      stop("Architecture mismatch at layer ", i, ": model has '",
           model$layers[[i]]$type, "', saved weights have '",
           arch$layer_configs[[i]]$type, "'")
    }
  }

  # Store R weight vectors in layers for nn_build_graph to pick up
  for (i in seq_along(model$layers)) {
    if (length(data$weights[[i]]) > 0) {
      model$layers[[i]]$weights_data <- data$weights[[i]]
    }
  }

  invisible(model)
}

# ============================================================================
# Print / Summary
# ============================================================================

#' Print method for ggml_sequential_model
#'
#' Prints a summary of the model architecture including layer types,
#' output shapes, and parameter counts.
#'
#' @param x A ggml_sequential_model object
#' @param ... Additional arguments (ignored)
#' @return The model object (invisibly).
#' @export
print.ggml_sequential_model <- function(x, ...) {
  model <- x
  cat("ggmlR Sequential Model\n")
  cat(paste(rep("=", 60), collapse = ""), "\n")

  if (length(model$layers) == 0) {
    cat("  (no layers)\n")
    return(invisible(model))
  }

  # Infer shapes if not done yet
  if (is.null(model$layers[[1]]$output_shape) && !is.null(model$input_shape)) {
    model <- nn_infer_shapes(model)
  }

  total_params <- 0
  cat(sprintf("%-20s %-20s %-10s\n", "Layer", "Output Shape", "Params"))
  cat(paste(rep("-", 60), collapse = ""), "\n")

  for (i in seq_along(model$layers)) {
    layer <- model$layers[[i]]

    n_params <- nn_count_layer_params(layer)
    total_params <- total_params + n_params

    shape_str <- if (!is.null(layer$output_shape)) {
      paste0("(", paste(layer$output_shape, collapse = ", "), ")")
    } else {
      "?"
    }

    cat(sprintf("%-20s %-20s %-10d\n", layer$type, shape_str, n_params))
  }

  cat(paste(rep("=", 60), collapse = ""), "\n")
  cat(sprintf("Total parameters: %d\n", total_params))
  cat(sprintf("Compiled: %s\n", if (model$compiled) "yes" else "no"))

  invisible(x)
}

#' Summary method for ggml_sequential_model
#'
#' Prints a detailed summary including input shape, layer details,
#' trainable/non-trainable parameter counts, and memory estimate.
#'
#' @param object A ggml_sequential_model object
#' @param ... Additional arguments (ignored)
#' @return The model object (invisibly).
#' @export
summary.ggml_sequential_model <- function(object, ...) {
  model <- object

  # Infer shapes if needed
  if (length(model$layers) > 0 &&
      is.null(model$layers[[1]]$output_shape) &&
      !is.null(model$input_shape)) {
    model <- nn_infer_shapes(model)
  }

  cat("ggmlR Sequential Model Summary\n")
  cat(paste(rep("=", 70), collapse = ""), "\n")

  if (!is.null(model$input_shape)) {
    cat(sprintf("Input shape: (%s)\n", paste(model$input_shape, collapse = ", ")))
  }
  cat("\n")

  if (length(model$layers) == 0) {
    cat("  (no layers)\n")
    return(invisible(object))
  }

  trainable <- 0
  non_trainable <- 0

  cat(sprintf("%-4s %-20s %-20s %-12s %-12s\n",
              "#", "Layer", "Output Shape", "Trainable", "Non-train."))
  cat(paste(rep("-", 70), collapse = ""), "\n")

  for (i in seq_along(model$layers)) {
    layer <- model$layers[[i]]
    n_params <- nn_count_layer_params(layer)
    trainable <- trainable + n_params

    shape_str <- if (!is.null(layer$output_shape)) {
      paste0("(", paste(layer$output_shape, collapse = ", "), ")")
    } else {
      "?"
    }

    cat(sprintf("%-4d %-20s %-20s %-12d %-12d\n",
                i, layer$type, shape_str, n_params, 0L))
  }

  cat(paste(rep("=", 70), collapse = ""), "\n")
  total <- trainable + non_trainable
  cat(sprintf("Total parameters:         %s\n", format(total, big.mark = ",")))
  cat(sprintf("  Trainable:              %s\n", format(trainable, big.mark = ",")))
  cat(sprintf("  Non-trainable:          %s\n", format(non_trainable, big.mark = ",")))

  # Memory estimate (F32 = 4 bytes per param)
  mem_bytes <- total * 4
  if (mem_bytes >= 1024 * 1024) {
    cat(sprintf("Estimated weight memory:  %.1f MB\n", mem_bytes / (1024 * 1024)))
  } else {
    cat(sprintf("Estimated weight memory:  %.1f KB\n", mem_bytes / 1024))
  }

  cat(sprintf("Compiled:                 %s\n", if (model$compiled) "yes" else "no"))

  invisible(object)
}

#' Count parameters for a single layer
#' @param layer A layer list
#' @return Number of parameters
#' @keywords internal
nn_count_layer_params <- function(layer) {
  if (layer$type == "conv_2d") {
    if (!is.null(layer$input_shape)) {
      ksize <- layer$config$kernel_size
      ksize[1] * ksize[2] * layer$input_shape[3] * layer$config$filters +
        layer$config$filters
    } else 0
  } else if (layer$type == "conv_1d") {
    if (!is.null(layer$input_shape)) {
      layer$config$kernel_size * layer$input_shape[2] * layer$config$filters +
        layer$config$filters
    } else 0
  } else if (layer$type == "dense") {
    if (!is.null(layer$input_shape)) {
      fan_in <- if (length(layer$input_shape) == 1) layer$input_shape else prod(layer$input_shape)
      fan_in * layer$config$units + layer$config$units
    } else 0
  } else if (layer$type == "batch_norm") {
    if (!is.null(layer$input_shape)) {
      n <- if (length(layer$input_shape) == 1) layer$input_shape
           else if (length(layer$input_shape) == 2) layer$input_shape[2]
           else layer$input_shape[3]
      n * 2L  # gamma + beta
    } else 0
  } else if (layer$type == "lstm") {
    if (!is.null(layer$input_shape)) {
      input_sz <- layer$input_shape[2]
      units    <- layer$config$units
      # W_gates + U_gates + b_gates
      input_sz * 4L * units + units * 4L * units + 4L * units
    } else 0
  } else if (layer$type == "gru") {
    if (!is.null(layer$input_shape)) {
      input_sz <- layer$input_shape[2]
      units    <- layer$config$units
      # W_zh + U_zh + b_zh + W_n + U_n + b_n
      input_sz * 2L * units + units * 2L * units + 2L * units +
        input_sz * units + units * units + units
    } else 0
  } else {
    0
  }
}

# ============================================================================
# History Class
# ============================================================================

#' Print method for ggml_history
#'
#' @param x A ggml_history object
#' @param ... Additional arguments (ignored)
#' @return The history object (invisibly).
#' @export
print.ggml_history <- function(x, ...) {
  n <- length(x$epochs)
  cat("Training History (", n, " epoch", if (n != 1) "s", ")\n", sep = "")
  cat(sprintf("  Final train loss:     %.4f\n", x$train_loss[n]))
  cat(sprintf("  Final train accuracy: %.4f\n", x$train_accuracy[n]))
  if (!is.na(x$val_loss[n])) {
    cat(sprintf("  Final val loss:       %.4f\n", x$val_loss[n]))
    cat(sprintf("  Final val accuracy:   %.4f\n", x$val_accuracy[n]))
  }
  invisible(x)
}

#' Plot training history
#'
#' Plots loss and accuracy curves over epochs.
#'
#' @param x A ggml_history object
#' @param ... Additional arguments (ignored)
#' @return The history object (invisibly).
#' @importFrom graphics plot lines legend par
#' @export
plot.ggml_history <- function(x, ...) {
  has_val <- !is.na(x$val_loss[1])
  old_par <- par(mfrow = c(1, 2))
  on.exit(par(old_par))

  # Loss plot
  ylim_loss <- range(c(x$train_loss, if (has_val) x$val_loss), na.rm = TRUE)
  plot(x$epochs, x$train_loss, type = "l", col = "blue",
       xlab = "Epoch", ylab = "Loss", main = "Loss", ylim = ylim_loss)
  if (has_val) {
    lines(x$epochs, x$val_loss, col = "red")
    legend("topright", legend = c("Train", "Val"),
           col = c("blue", "red"), lty = 1, cex = 0.8)
  }

  # Accuracy plot
  ylim_acc <- range(c(x$train_accuracy, if (has_val) x$val_accuracy), na.rm = TRUE)
  plot(x$epochs, x$train_accuracy, type = "l", col = "blue",
       xlab = "Epoch", ylab = "Accuracy", main = "Accuracy", ylim = ylim_acc)
  if (has_val) {
    lines(x$epochs, x$val_accuracy, col = "red")
    legend("bottomright", legend = c("Train", "Val"),
           col = c("blue", "red"), lty = 1, cex = 0.8)
  }

  invisible(x)
}

# ============================================================================
# Save / load full model (architecture + weights)
# ============================================================================

#' Save a Full Model (Architecture + Weights)
#'
#' Saves both the architecture and trained weights of a model to an RDS file.
#' Unlike \code{ggml_save_weights()}, which requires the model to be manually
#' reconstructed before loading, \code{ggml_save_model()} saves everything
#' needed to restore the model with a single call to \code{ggml_load_model()}.
#'
#' @section Supported model types:
#' \itemize{
#'   \item \code{ggml_sequential_model} — input shape, layer configs, trained
#'     weights, and compilation settings are all saved.
#'   \item \code{ggml_functional_model} — input/output node graphs (pure R
#'     lists, no ggml pointers) and trained \code{node_weights} are saved.
#' }
#'
#' @param model A trained \code{ggml_sequential_model} or
#'   \code{ggml_functional_model}.
#' @param path File path (typically \code{.rds}).
#' @return The model (invisibly).
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(16L, activation = "relu", input_shape = 4L) |>
#'   ggml_layer_dense(2L,  activation = "softmax")
#' model <- ggml_compile(model, optimizer = "adam",
#'                        loss = "categorical_crossentropy")
#' x <- matrix(runif(64 * 4), 64, 4)
#' y <- matrix(c(rep(c(1,0), 32), rep(c(0,1), 32)), 64, 2)
#' model <- ggml_fit(model, x, y, epochs = 1L, batch_size = 32L, verbose = 0L)
#' tmp <- tempfile(fileext = ".rds")
#' ggml_save_model(model, tmp)
#' model2 <- ggml_load_model(tmp)
#' }
ggml_save_model <- function(model, path) {
  UseMethod("ggml_save_model")
}

#' @export
ggml_save_model.ggml_sequential_model <- function(model, path) {
  if (!model$compiled) stop("Model must be compiled before saving.")

  # Strip live ggml tensors; keep only serialisable config + weight values
  layers_clean <- lapply(model$layers, function(l) {
    wdata <- list()
    if (l$type %in% c("conv_1d", "conv_2d") && !is.null(l$weights$kernel)) {
      wdata$kernel <- ggml_backend_tensor_get_data(l$weights$kernel)
      wdata$bias   <- ggml_backend_tensor_get_data(l$weights$bias)
    } else if (l$type == "dense" && !is.null(l$weights$weight)) {
      wdata$weight <- ggml_backend_tensor_get_data(l$weights$weight)
      wdata$bias   <- ggml_backend_tensor_get_data(l$weights$bias)
    } else if (l$type == "batch_norm" && !is.null(l$weights$gamma)) {
      wdata$gamma <- ggml_backend_tensor_get_data(l$weights$gamma)
      wdata$beta  <- ggml_backend_tensor_get_data(l$weights$beta)
    } else if (l$type == "lstm" && !is.null(l$weights$W_gates)) {
      wdata$W_gates <- ggml_backend_tensor_get_data(l$weights$W_gates)
      wdata$U_gates <- ggml_backend_tensor_get_data(l$weights$U_gates)
      wdata$b_gates <- ggml_backend_tensor_get_data(l$weights$b_gates)
    } else if (l$type == "gru" && !is.null(l$weights$W_zh)) {
      wdata$W_zh <- ggml_backend_tensor_get_data(l$weights$W_zh)
      wdata$U_zh <- ggml_backend_tensor_get_data(l$weights$U_zh)
      wdata$b_zh <- ggml_backend_tensor_get_data(l$weights$b_zh)
      wdata$W_n  <- ggml_backend_tensor_get_data(l$weights$W_n)
      wdata$U_n  <- ggml_backend_tensor_get_data(l$weights$U_n)
      wdata$b_n  <- ggml_backend_tensor_get_data(l$weights$b_n)
    } else if (l$type == "embedding" && !is.null(l$weights$weight)) {
      wdata$weight <- ggml_backend_tensor_get_data(l$weights$weight)
    }
    list(type        = l$type,
         name        = l$name,
         trainable   = l$trainable,
         config      = l$config,
         input_shape = l$input_shape,
         output_shape= l$output_shape,
         weights_data= wdata)
  })

  saveRDS(list(
    model_class  = "ggml_sequential_model",
    input_shape  = model$input_shape,
    layers       = layers_clean,
    compilation  = list(
      optimizer = model$compilation$optimizer,
      loss      = model$compilation$loss,
      metrics   = model$compilation$metrics
    ),
    version = 2L
  ), path)
  invisible(model)
}

#' @export
ggml_save_model.ggml_functional_model <- function(model, path) {
  if (!model$compiled) stop("Model must be compiled before saving.")

  # node_weights: list of node_id -> list of ggml tensors
  # Extract numeric data from each tensor
  nw_data <- NULL
  if (!is.null(model$node_weights)) {
    nw_data <- lapply(model$node_weights, function(wlist) {
      lapply(wlist, function(t) {
        if (is.null(t)) NULL else ggml_backend_tensor_get_data(t)
      })
    })
  }

  saveRDS(list(
    model_class  = "ggml_functional_model",
    inputs       = model$inputs,     # pure R ggml_tensor_node lists
    outputs      = model$outputs,
    compilation  = list(
      optimizer = model$compilation$optimizer,
      loss      = model$compilation$loss,
      metrics   = model$compilation$metrics
    ),
    node_weights_data = nw_data,
    version = 2L
  ), path)
  invisible(model)
}

#' Load a Full Model (Architecture + Weights)
#'
#' Restores a model previously saved with \code{ggml_save_model()}.  The
#' returned model is compiled and ready for \code{ggml_predict()} /
#' \code{ggml_evaluate()}.  Call \code{ggml_fit()} again to continue training.
#'
#' @param path File path to an RDS file written by \code{ggml_save_model()}.
#' @param backend Backend selection: \code{"auto"}, \code{"cpu"}, or
#'   \code{"vulkan"}.
#' @return A compiled model object.
#' @export
#' @examples
#' \donttest{
#' model <- ggml_model_sequential() |>
#'   ggml_layer_dense(16L, activation = "relu", input_shape = 4L) |>
#'   ggml_layer_dense(2L,  activation = "softmax")
#' model <- ggml_compile(model, optimizer = "adam",
#'                        loss = "categorical_crossentropy")
#' x <- matrix(runif(64 * 4), 64, 4)
#' y <- matrix(c(rep(c(1,0), 32), rep(c(0,1), 32)), 64, 2)
#' model <- ggml_fit(model, x, y, epochs = 1L, batch_size = 32L, verbose = 0L)
#' tmp <- tempfile(fileext = ".rds")
#' ggml_save_model(model, tmp)
#' model2 <- ggml_load_model(tmp)
#' }
ggml_load_model <- function(path, backend = "auto") {
  data <- readRDS(path)
  if (is.null(data$version) || data$version < 2L) {
    stop("This file was saved with ggml_save_weights(). ",
         "Use ggml_load_weights() instead, or re-save with ggml_save_model().")
  }

  if (data$model_class == "ggml_sequential_model") {
    # Reconstruct sequential model
    model <- ggml_model_sequential()
    model$input_shape <- data$input_shape

    for (ldata in data$layers) {
      layer <- list(
        type         = ldata$type,
        name         = ldata$name,
        trainable    = ldata$trainable,
        config       = ldata$config,
        input_shape  = ldata$input_shape,
        output_shape = ldata$output_shape,
        weights      = list(),
        weights_data = ldata$weights_data
      )
      model$layers <- c(model$layers, list(layer))
    }

    model <- ggml_compile(model,
      optimizer = data$compilation$optimizer,
      loss      = data$compilation$loss,
      metrics   = data$compilation$metrics,
      backend   = backend
    )
    return(invisible(model))

  } else if (data$model_class == "ggml_functional_model") {
    model <- ggml_model(inputs = data$inputs, outputs = data$outputs)
    model <- ggml_compile(model,
      optimizer = data$compilation$optimizer,
      loss      = data$compilation$loss,
      metrics   = data$compilation$metrics,
      backend   = backend
    )

    # Ensure no stale ggml tensor pointers from the freshly-created model.
    model$node_weights <- NULL

    # Restore node_weights as R-vector lists so nn_build_functional_graph
    # picks them up via the swd (saved_weights_data) path.
    if (!is.null(data$node_weights_data)) {
      model$node_weights_data <- data$node_weights_data
    }
    return(invisible(model))
  } else {
    stop("Unknown model class in saved file: ", data$model_class)
  }
}