#' Unpack MLP parameter vector
#'
#' Converts a flat parameter vector into weight matrices and bias vectors
#' for a single-hidden-layer multilayer perceptron.
#'
#' @param params Numeric parameter vector.
#' @param input_dim Number of input variables.
#' @param hidden_dim Number of hidden neurons.
#' @param output_dim Number of output neurons.
#'
#' @return A list containing \code{W1}, \code{b1}, \code{W2}, and \code{b2}.
unpack_mlp_params <- function(params, input_dim, hidden_dim, output_dim) {
  idx1 <- input_dim * hidden_dim
  idx2 <- idx1 + hidden_dim
  idx3 <- idx2 + hidden_dim * output_dim

  W1 <- matrix(params[1:idx1], nrow = input_dim, ncol = hidden_dim)
  b1 <- params[(idx1 + 1):idx2]
  W2 <- matrix(params[(idx2 + 1):idx3], nrow = hidden_dim, ncol = output_dim)
  b2 <- params[(idx3 + 1):(idx3 + output_dim)]

  list(
    W1 = W1,
    b1 = b1,
    W2 = W2,
    b2 = b2
  )
}

#' Forward pass for a single-hidden-layer MLP
#'
#' Computes hidden-layer activations and output predictions.
#'
#' @param params Numeric parameter vector.
#' @param X Input matrix.
#' @param input_dim Number of input variables.
#' @param hidden_dim Number of hidden neurons.
#' @param output_dim Number of output neurons.
#'
#' @return A list containing hidden activations and predictions.
forward_mlp <- function(params, X, input_dim, hidden_dim, output_dim) {
  net <- unpack_mlp_params(params, input_dim, hidden_dim, output_dim)

  Z1 <- X %*% net$W1 + matrix(net$b1, nrow = nrow(X), ncol = length(net$b1), byrow = TRUE)
  A1 <- sigmoid(Z1)

  Z2 <- A1 %*% net$W2 + matrix(net$b2, nrow = nrow(A1), ncol = length(net$b2), byrow = TRUE)
  Y_hat <- sigmoid(Z2)

  list(
    A1 = A1,
    Y_hat = Y_hat
  )
}

#' MSE fitness function for MLP training
#'
#' Computes mean squared error for a single-hidden-layer MLP.
#'
#' @param params Numeric parameter vector.
#' @param X Input matrix.
#' @param Y Output matrix.
#' @param input_dim Number of input variables.
#' @param hidden_dim Number of hidden neurons.
#' @param output_dim Number of output neurons.
#'
#' @return Mean squared error.
mlp_mse_fitness <- function(params, X, Y, input_dim, hidden_dim, output_dim) {
  pred <- forward_mlp(params, X, input_dim, hidden_dim, output_dim)$Y_hat
  mean((Y - pred)^2)
}
