#' EncoderRNN Neural Network Module
#'
#' An RNN-based encoder module using Long Short-Term Memory (LSTM) architecture, implemented with the \code{torch} package.
#' This module is designed to process sequential data and is commonly used in tasks like sequence-to-sequence modeling, language modeling, and time series forecasting.
#'
#' @section Usage:
#' \preformatted{
#' EncoderRNN <- nn_module(
#'   "EncoderRNN",
#'
#'   initialize = function(input_size, nhidden, nlayers, dropout) {
#'     ...
#'   },
#'
#'   forward = function(x) {
#'     ...
#'   }
#' )
#' }
#'
#' @section Methods:
#' \describe{
#'   \item{\code{initialize(input_size, nhidden, nlayers, dropout)}}{
#'     Initializes the EncoderRNN module.
#'
#'     \itemize{
#'       \item{\code{input_size}: A numeric value representing the number of features in the input data.}
#'       \item{\code{nhidden}: A numeric value representing the number of hidden units in each LSTM layer.}
#'       \item{\code{nlayers}: A numeric value indicating the number of LSTM layers.}
#'       \item{\code{dropout}: A numeric value between 0 and 1 indicating the dropout probability to be applied between LSTM layers.}
#'     }
#'   }
#'
#'   \item{\code{forward(x)}}{
#'     Executes the forward pass of the EncoderRNN module.
#'
#'     \itemize{
#'       \item{\code{x}: A tensor of shape (batch_size, sequence_length, input_size) representing the input data.}
#'     }
#'
#'     \value{
#'     A list containing:
#'       \itemize{
#'         \item{\code{output}: The output tensor from the LSTM, processed and flipped.}
#'         \item{\code{list(hn, cn)}: A list containing the hidden state (\code{hn}) and cell state (\code{cn}) from the LSTM.}
#'         \item{\code{newinput}: A tensor representing the input data with an additional zero tensor at the beginning.}
#'       }
#'     }
#'   }
#' }
#'
#' @param input_size A numeric value representing the number of features in the input data.
#' @param nhidden A numeric value representing the number of hidden units in each LSTM layer.
#' @param nlayers A numeric value indicating the number of LSTM layers.
#' @param dropout A numeric value between 0 and 1 indicating the dropout probability to be applied between LSTM layers.
# @param x A tensor of shape \code{(batch_size, sequence_length, input_size)} representing the input data.
# @param h A list containing the hidden state (\code{hn}) and cell state (\code{cn}) from the LSTM, typically passed from the encoder or previous decoder step.

#' @return An \code{nn_module} object (torch LSTM-based encoder). When called
#'   with a tensor of shape \code{(batch_size, seq_len, input_size)}, returns a
#'   list of: (1) processed output tensor, (2) list of hidden and cell states
#'   \code{(hn, cn)}, and (3) the input with a leading zero tensor prepended.
#' @examples
#' \donttest{
#' enc <- EncoderRNN(input_size = 4L, nhidden = 8L, nlayers = 1L, dropout = 0.0)
#' x   <- torch_randn(2L, 1L, 4L)   # batch=2, seq_len=1, features=4
#' out <- enc(x)
#' dim(out[[1]])  # encoded output
#' }
#' @import torch
#' @export




EncoderRNN <- torch::nn_module(
  "EncoderRNN",

  initialize = function(input_size, nhidden, nlayers, dropout) {
    self$nhidden <- nhidden
    self$feasize <- input_size
    self$nlayers <- nlayers
    self$dropout <- dropout

    # Define the LSTM layer
    self$lstm <- nn_lstm(
      input_size = input_size, #self$feasize,
      hidden_size = nhidden, #self$nhidden,
      num_layers = nlayers, #self$nlayers,
      dropout = dropout, #self$dropout,
      batch_first = TRUE
    )

    # Initialize weights directly
    #params <- self$lstm$parameters()
    #for (p in params) {
    #  p$data$uniform_(-0.1, 0.1)
    #}
  },

  forward = function(x) {
    batch_size <- x$size(1)
    # Forward pass through LSTM
    out_state <- self$lstm(x)
    output <- out_state[[1]]
    state <- out_state[[2]]
    hn <- state[[1]]
    cn <- state[[2]]

    # Process the output
    output <- torch_flip(output, dims = c(1))
    newinput <- torch_flip(x, dims = c(1))

    # Add a zero tensor
    zeros <- torch_zeros(c(batch_size, 1, x$size(3)))
    newinput <- torch_cat(list(zeros, newinput), 2)
    #newinput <- newinput$narrow(2, 0, newinput$size(2) - 1)
    #newinput <- torch::torch_reshape(newinput, c(batch_size, sequence_length, input_size))
    #newinput <- torch::torch_reshape(newinput, c(batch_size, sequence_length, input_size))

    # Return the processed output and new input
    list(output, list(hn, cn), newinput)
  }

)
