sklears-neural 0.2.0

//! Recurrent Neural Network layers including LSTM and GRU cells.
//!
//! This module provides implementations of recurrent neural network layers
//! including LSTM (Long Short-Term Memory) and GRU (Gated Recurrent Unit) cells.

use crate::activation::Activation;
use crate::weight_init::{InitStrategy, WeightInitializer};
use crate::{layers::Layer, NeuralResult};
use scirs2_core::ndarray::{s, Array1, Array2, Array3, Axis};
use scirs2_core::random::thread_rng;
use sklears_core::error::SklearsError;
use sklears_core::types::FloatBounds;

/// Helper function to apply activation to a generic floating point array
fn apply_activation<T: FloatBounds>(activation: &Activation, input: &Array2<T>) -> Array2<T> {
    match activation {
        Activation::Identity => input.clone(),
        Activation::Logistic => input.mapv(|val| {
            let exp_neg = (-val).exp();
            T::one() / (T::one() + exp_neg)
        }),
        Activation::Tanh => input.mapv(|val| val.tanh()),
        Activation::Relu => input.mapv(|val| val.max(T::zero())),
        _ => input.clone(), // For other activations, just return input for now
    }
}

/// LSTM cell state containing hidden and cell states
#[derive(Debug, Clone)]
pub struct LSTMState<T: FloatBounds> {
    /// Hidden state (h_t)
    pub hidden: Array2<T>,
    /// Cell state (c_t)
    pub cell: Array2<T>,
}

impl<T: FloatBounds> LSTMState<T> {
    /// Create a new LSTM state with given batch size and hidden size
    pub fn new(batch_size: usize, hidden_size: usize) -> Self {
        Self {
            hidden: Array2::zeros((batch_size, hidden_size)),
            cell: Array2::zeros((batch_size, hidden_size)),
        }
    }

    /// Reset the state to zeros
    pub fn reset(&mut self) {
        self.hidden.fill(T::zero());
        self.cell.fill(T::zero());
    }

    /// Clone the state
    pub fn clone_state(&self) -> Self {
        Self {
            hidden: self.hidden.clone(),
            cell: self.cell.clone(),
        }
    }
}

/// LSTM (Long Short-Term Memory) cell implementation
///
/// The LSTM cell computes:
/// - Forget gate: f_t = sigmoid(W_f * [h_{t-1}, x_t] + b_f)
/// - Input gate: i_t = sigmoid(W_i * [h_{t-1}, x_t] + b_i)
/// - Candidate values: g_t = tanh(W_g * [h_{t-1}, x_t] + b_g)
/// - Output gate: o_t = sigmoid(W_o * [h_{t-1}, x_t] + b_o)
/// - Cell state: c_t = f_t * c_{t-1} + i_t * g_t
/// - Hidden state: h_t = o_t * tanh(c_t)
#[derive(Debug, Clone)]
pub struct LSTMCell<T: FloatBounds> {
    /// Input size
    input_size: usize,
    /// Hidden size
    hidden_size: usize,
    /// Forget gate weights [W_f_input, W_f_hidden]
    weight_forget: Array2<T>,
    /// Input gate weights [W_i_input, W_i_hidden]
    weight_input: Array2<T>,
    /// Candidate gate weights [W_g_input, W_g_hidden]
    weight_candidate: Array2<T>,
    /// Output gate weights [W_o_input, W_o_hidden]
    weight_output: Array2<T>,
    /// Forget gate bias
    bias_forget: Array1<T>,
    /// Input gate bias
    bias_input: Array1<T>,
    /// Candidate gate bias
    bias_candidate: Array1<T>,
    /// Output gate bias
    bias_output: Array1<T>,
    /// Current state
    state: Option<LSTMState<T>>,
    /// Activation function for gates (sigmoid)
    gate_activation: Activation,
    /// Activation function for cell state (tanh)
    cell_activation: Activation,
    /// Cached values for backward pass
    cached_forget_gate: Option<Array2<T>>,
    cached_input_gate: Option<Array2<T>>,
    cached_candidate_gate: Option<Array2<T>>,
    cached_output_gate: Option<Array2<T>>,
    cached_input: Option<Array2<T>>,
    cached_hidden_prev: Option<Array2<T>>,
    cached_cell_prev: Option<Array2<T>>,
    cached_cell_tanh: Option<Array2<T>>,
}

impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> LSTMCell<T> {
    /// Create a new LSTM cell
    pub fn new(input_size: usize, hidden_size: usize) -> NeuralResult<Self> {
        let total_input_size = input_size + hidden_size;

        // Initialize weights using Xavier initialization
        let mut rng = thread_rng();
        let initializer = WeightInitializer::new(InitStrategy::XavierUniform);

        let weight_forget = initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;
        let weight_input = initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;
        let weight_candidate =
            initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;
        let weight_output = initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;

        // Initialize biases to zero (except forget gate bias which is initialized to 1)
        let mut bias_forget = Array1::zeros(hidden_size);
        bias_forget.fill(T::one()); // Initialize forget gate bias to 1 for better gradient flow
        let bias_input = Array1::zeros(hidden_size);
        let bias_candidate = Array1::zeros(hidden_size);
        let bias_output = Array1::zeros(hidden_size);

        Ok(Self {
            input_size,
            hidden_size,
            weight_forget,
            weight_input,
            weight_candidate,
            weight_output,
            bias_forget,
            bias_input,
            bias_candidate,
            bias_output,
            state: None,
            gate_activation: Activation::Logistic, // Use Logistic instead of Sigmoid
            cell_activation: Activation::Tanh,
            cached_forget_gate: None,
            cached_input_gate: None,
            cached_candidate_gate: None,
            cached_output_gate: None,
            cached_input: None,
            cached_hidden_prev: None,
            cached_cell_prev: None,
            cached_cell_tanh: None,
        })
    }

    /// Initialize the state for a given batch size
    pub fn init_state(&mut self, batch_size: usize) {
        self.state = Some(LSTMState::new(batch_size, self.hidden_size));
    }

    /// Get the current state
    pub fn get_state(&self) -> Option<&LSTMState<T>> {
        self.state.as_ref()
    }

    /// Set the state
    pub fn set_state(&mut self, state: LSTMState<T>) {
        self.state = Some(state);
    }

    /// Forward pass through the LSTM cell
    pub fn forward_step(&mut self, input: &Array2<T>) -> NeuralResult<Array2<T>> {
        let (batch_size, input_features) = input.dim();

        if input_features != self.input_size {
            return Err(SklearsError::InvalidParameter {
                name: "input_size".to_string(),
                reason: format!("expected {}, got {}", self.input_size, input_features),
            });
        }

        // Initialize state if not already done
        if self.state.is_none() {
            self.init_state(batch_size);
        }

        let state = self
            .state
            .as_ref()
            .expect("state not available - model not fitted");

        // Cache previous states for backward pass
        self.cached_hidden_prev = Some(state.hidden.clone());
        self.cached_cell_prev = Some(state.cell.clone());
        self.cached_input = Some(input.clone());

        // Concatenate input and previous hidden state
        let mut combined_input = Array2::zeros((batch_size, self.input_size + self.hidden_size));
        combined_input
            .slice_mut(s![.., ..self.input_size])
            .assign(input);
        combined_input
            .slice_mut(s![.., self.input_size..])
            .assign(&state.hidden);

        // Compute gates
        let forget_gate_pre = combined_input.dot(&self.weight_forget.t()) + &self.bias_forget;
        let forget_gate = apply_activation(&self.gate_activation, &forget_gate_pre);

        let input_gate_pre = combined_input.dot(&self.weight_input.t()) + &self.bias_input;
        let input_gate = apply_activation(&self.gate_activation, &input_gate_pre);

        let candidate_gate_pre =
            combined_input.dot(&self.weight_candidate.t()) + &self.bias_candidate;
        let candidate_gate = apply_activation(&self.cell_activation, &candidate_gate_pre);

        let output_gate_pre = combined_input.dot(&self.weight_output.t()) + &self.bias_output;
        let output_gate = apply_activation(&self.gate_activation, &output_gate_pre);

        // Cache gates for backward pass
        self.cached_forget_gate = Some(forget_gate.clone());
        self.cached_input_gate = Some(input_gate.clone());
        self.cached_candidate_gate = Some(candidate_gate.clone());
        self.cached_output_gate = Some(output_gate.clone());

        // Update cell state: c_t = f_t * c_{t-1} + i_t * g_t
        let new_cell = &forget_gate * &state.cell + &input_gate * &candidate_gate;

        // Compute new hidden state: h_t = o_t * tanh(c_t)
        let cell_tanh = apply_activation(&self.cell_activation, &new_cell);
        let new_hidden = &output_gate * &cell_tanh;

        // Cache cell tanh for backward pass
        self.cached_cell_tanh = Some(cell_tanh);

        // Update state
        self.state.as_mut().expect("state not available").cell = new_cell;
        self.state.as_mut().expect("state not available").hidden = new_hidden.clone();

        Ok(new_hidden)
    }

    /// Process a sequence of inputs
    pub fn forward_sequence(&mut self, inputs: &Array3<T>) -> NeuralResult<Array3<T>> {
        let (batch_size, seq_len, input_size) = inputs.dim();

        if input_size != self.input_size {
            return Err(SklearsError::InvalidParameter {
                name: "input_size".to_string(),
                reason: format!("expected {}, got {}", self.input_size, input_size),
            });
        }

        // Initialize state
        self.init_state(batch_size);

        let mut outputs = Array3::zeros((batch_size, seq_len, self.hidden_size));

        // Process each time step
        for t in 0..seq_len {
            let input_t = inputs.slice(s![.., t, ..]).to_owned();
            let output_t = self.forward_step(&input_t)?;
            outputs.slice_mut(s![.., t, ..]).assign(&output_t);
        }

        Ok(outputs)
    }

    /// Reset the cell state
    pub fn reset_state(&mut self) {
        if let Some(ref mut state) = self.state {
            state.reset();
        }
    }

    /// Get the number of parameters
    pub fn num_parameters(&self) -> usize {
        let weight_params = self.weight_forget.len()
            + self.weight_input.len()
            + self.weight_candidate.len()
            + self.weight_output.len();
        let bias_params = self.bias_forget.len()
            + self.bias_input.len()
            + self.bias_candidate.len()
            + self.bias_output.len();
        weight_params + bias_params
    }
}

/// GRU (Gated Recurrent Unit) cell implementation
///
/// The GRU cell computes:
/// - Reset gate: r_t = sigmoid(W_r * [h_{t-1}, x_t] + b_r)
/// - Update gate: z_t = sigmoid(W_z * [h_{t-1}, x_t] + b_z)
/// - New gate: n_t = tanh(W_n * [r_t * h_{t-1}, x_t] + b_n)
/// - Hidden state: h_t = (1 - z_t) * n_t + z_t * h_{t-1}
#[derive(Debug, Clone)]
pub struct GRUCell<T: FloatBounds> {
    /// Input size
    input_size: usize,
    /// Hidden size
    hidden_size: usize,
    /// Reset gate weights [W_r_input, W_r_hidden]
    weight_reset: Array2<T>,
    /// Update gate weights [W_z_input, W_z_hidden]
    weight_update: Array2<T>,
    /// New gate weights [W_n_input, W_n_hidden]
    weight_new: Array2<T>,
    /// Reset gate bias
    bias_reset: Array1<T>,
    /// Update gate bias
    bias_update: Array1<T>,
    /// New gate bias
    bias_new: Array1<T>,
    /// Current hidden state
    hidden_state: Option<Array2<T>>,
    /// Activation function for gates (sigmoid)
    gate_activation: Activation,
    /// Activation function for new state (tanh)
    new_activation: Activation,
    /// Cached values for backward pass
    cached_reset_gate: Option<Array2<T>>,
    cached_update_gate: Option<Array2<T>>,
    cached_new_gate: Option<Array2<T>>,
    cached_input: Option<Array2<T>>,
    cached_hidden_prev: Option<Array2<T>>,
    cached_reset_hidden: Option<Array2<T>>,
}

impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> GRUCell<T> {
    /// Create a new GRU cell
    pub fn new(input_size: usize, hidden_size: usize) -> NeuralResult<Self> {
        let total_input_size = input_size + hidden_size;

        // Initialize weights using Xavier initialization
        let mut rng = thread_rng();
        let initializer = WeightInitializer::new(InitStrategy::XavierUniform);

        let weight_reset = initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;
        let weight_update = initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;
        let weight_new = initializer.initialize_2d(&mut rng, (hidden_size, total_input_size))?;

        // Initialize biases to zero
        let bias_reset = Array1::zeros(hidden_size);
        let bias_update = Array1::zeros(hidden_size);
        let bias_new = Array1::zeros(hidden_size);

        Ok(Self {
            input_size,
            hidden_size,
            weight_reset,
            weight_update,
            weight_new,
            bias_reset,
            bias_update,
            bias_new,
            hidden_state: None,
            gate_activation: Activation::Logistic,
            new_activation: Activation::Tanh,
            cached_reset_gate: None,
            cached_update_gate: None,
            cached_new_gate: None,
            cached_input: None,
            cached_hidden_prev: None,
            cached_reset_hidden: None,
        })
    }

    /// Initialize the hidden state for a given batch size
    pub fn init_state(&mut self, batch_size: usize) {
        self.hidden_state = Some(Array2::zeros((batch_size, self.hidden_size)));
    }

    /// Get the current hidden state
    pub fn get_hidden_state(&self) -> Option<&Array2<T>> {
        self.hidden_state.as_ref()
    }

    /// Set the hidden state
    pub fn set_hidden_state(&mut self, state: Array2<T>) {
        self.hidden_state = Some(state);
    }

    /// Forward pass through the GRU cell
    pub fn forward_step(&mut self, input: &Array2<T>) -> NeuralResult<Array2<T>> {
        let (batch_size, input_features) = input.dim();

        if input_features != self.input_size {
            return Err(SklearsError::InvalidParameter {
                name: "input_size".to_string(),
                reason: format!("expected {}, got {}", self.input_size, input_features),
            });
        }

        // Initialize state if not already done
        if self.hidden_state.is_none() {
            self.init_state(batch_size);
        }

        let hidden_prev = self
            .hidden_state
            .as_ref()
            .expect("hidden_state not available - model not fitted");

        // Cache previous state and input for backward pass
        self.cached_hidden_prev = Some(hidden_prev.clone());
        self.cached_input = Some(input.clone());

        // Concatenate input and previous hidden state
        let mut combined_input = Array2::zeros((batch_size, self.input_size + self.hidden_size));
        combined_input
            .slice_mut(s![.., ..self.input_size])
            .assign(input);
        combined_input
            .slice_mut(s![.., self.input_size..])
            .assign(hidden_prev);

        // Compute reset and update gates
        let reset_gate_pre = combined_input.dot(&self.weight_reset.t()) + &self.bias_reset;
        let reset_gate = apply_activation(&self.gate_activation, &reset_gate_pre);

        let update_gate_pre = combined_input.dot(&self.weight_update.t()) + &self.bias_update;
        let update_gate = apply_activation(&self.gate_activation, &update_gate_pre);

        // Cache gates for backward pass
        self.cached_reset_gate = Some(reset_gate.clone());
        self.cached_update_gate = Some(update_gate.clone());

        // Compute reset hidden state
        let reset_hidden = &reset_gate * hidden_prev;
        self.cached_reset_hidden = Some(reset_hidden.clone());

        // Combine input with reset hidden state for new gate
        let mut new_input = Array2::zeros((batch_size, self.input_size + self.hidden_size));
        new_input.slice_mut(s![.., ..self.input_size]).assign(input);
        new_input
            .slice_mut(s![.., self.input_size..])
            .assign(&reset_hidden);

        // Compute new gate
        let new_gate_pre = new_input.dot(&self.weight_new.t()) + &self.bias_new;
        let new_gate = apply_activation(&self.new_activation, &new_gate_pre);

        // Cache new gate for backward pass
        self.cached_new_gate = Some(new_gate.clone());

        // Compute new hidden state: h_t = (1 - z_t) * n_t + z_t * h_{t-1}
        let one = Array2::ones(update_gate.dim());
        let new_hidden = (&one - &update_gate) * &new_gate + &update_gate * hidden_prev;

        // Update state
        self.hidden_state = Some(new_hidden.clone());

        Ok(new_hidden)
    }

    /// Process a sequence of inputs
    pub fn forward_sequence(&mut self, inputs: &Array3<T>) -> NeuralResult<Array3<T>> {
        let (batch_size, seq_len, input_size) = inputs.dim();

        if input_size != self.input_size {
            return Err(SklearsError::InvalidParameter {
                name: "input_size".to_string(),
                reason: format!("expected {}, got {}", self.input_size, input_size),
            });
        }

        // Initialize state
        self.init_state(batch_size);

        let mut outputs = Array3::zeros((batch_size, seq_len, self.hidden_size));

        // Process each time step
        for t in 0..seq_len {
            let input_t = inputs.slice(s![.., t, ..]).to_owned();
            let output_t = self.forward_step(&input_t)?;
            outputs.slice_mut(s![.., t, ..]).assign(&output_t);
        }

        Ok(outputs)
    }

    /// Reset the hidden state
    pub fn reset_state(&mut self) {
        if let Some(ref mut state) = self.hidden_state {
            state.fill(T::zero());
        }
    }

    /// Get the number of parameters
    pub fn num_parameters(&self) -> usize {
        let weight_params =
            self.weight_reset.len() + self.weight_update.len() + self.weight_new.len();
        let bias_params = self.bias_reset.len() + self.bias_update.len() + self.bias_new.len();
        weight_params + bias_params
    }
}

// Implement Layer trait for LSTMCell
impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> Layer<T> for LSTMCell<T> {
    fn forward(&mut self, input: &Array2<T>, _training: bool) -> NeuralResult<Array2<T>> {
        self.forward_step(input)
    }

    /// Compute gradients via Backpropagation Through Time (BPTT).
    ///
    /// # Note
    ///
    /// Not implemented in v0.1.0. Returns `Err(NotImplemented)`. Planned for v0.2.0.
    /// Full BPTT requires unrolling the recurrence and computing gradients through
    /// all gate computations (input, forget, output, cell) at each timestep.
    fn backward(&mut self, _grad_output: &Array2<T>) -> NeuralResult<Array2<T>> {
        Err(SklearsError::NotImplemented(
            "LSTM backward pass (BPTT) not yet implemented".to_string(),
        ))
    }

    fn num_parameters(&self) -> usize {
        LSTMCell::num_parameters(self)
    }

    fn reset(&mut self) {
        self.reset_state();
    }
}

// Implement Layer trait for GRUCell
impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> Layer<T> for GRUCell<T> {
    fn forward(&mut self, input: &Array2<T>, _training: bool) -> NeuralResult<Array2<T>> {
        self.forward_step(input)
    }

    /// Compute gradients via Backpropagation Through Time (BPTT).
    ///
    /// # Note
    ///
    /// Not implemented in v0.1.0. Returns `Err(NotImplemented)`. Planned for v0.2.0.
    /// Full BPTT requires unrolling the recurrence and computing gradients through
    /// the reset and update gates at each timestep.
    fn backward(&mut self, _grad_output: &Array2<T>) -> NeuralResult<Array2<T>> {
        Err(SklearsError::NotImplemented(
            "GRU backward pass (BPTT) not yet implemented".to_string(),
        ))
    }

    fn num_parameters(&self) -> usize {
        GRUCell::num_parameters(self)
    }

    fn reset(&mut self) {
        self.reset_state();
    }
}

/// Bidirectional RNN wrapper that can wrap LSTM or GRU cells
#[derive(Debug, Clone)]
pub struct BidirectionalRNN<T: FloatBounds, Cell> {
    /// Forward direction cell
    forward_cell: Cell,
    /// Backward direction cell
    backward_cell: Cell,
    /// Hidden size of each direction
    hidden_size: usize,
    /// Whether to concatenate outputs or add them
    concat_outputs: bool,
    _phantom: std::marker::PhantomData<T>,
}

impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> BidirectionalRNN<T, LSTMCell<T>> {
    /// Create a new bidirectional LSTM
    pub fn new_lstm(
        input_size: usize,
        hidden_size: usize,
        concat_outputs: bool,
    ) -> NeuralResult<Self> {
        let forward_cell = LSTMCell::new(input_size, hidden_size)?;
        let backward_cell = LSTMCell::new(input_size, hidden_size)?;

        Ok(Self {
            forward_cell,
            backward_cell,
            hidden_size,
            concat_outputs,
            _phantom: std::marker::PhantomData,
        })
    }

    /// Forward pass through bidirectional LSTM
    pub fn forward_sequence(&mut self, inputs: &Array3<T>) -> NeuralResult<Array3<T>> {
        let (batch_size, seq_len, _) = inputs.dim();

        // Forward pass
        let forward_outputs = self.forward_cell.forward_sequence(inputs)?;

        // Backward pass (reverse the input sequence)
        let mut reversed_inputs = inputs.clone();
        reversed_inputs.invert_axis(Axis(1));
        let mut backward_outputs = self.backward_cell.forward_sequence(&reversed_inputs)?;
        // Reverse the outputs back
        backward_outputs.invert_axis(Axis(1));

        // Combine outputs
        let output_size = if self.concat_outputs {
            self.hidden_size * 2
        } else {
            self.hidden_size
        };

        let mut combined_outputs = Array3::zeros((batch_size, seq_len, output_size));

        if self.concat_outputs {
            // Concatenate forward and backward outputs
            combined_outputs
                .slice_mut(s![.., .., ..self.hidden_size])
                .assign(&forward_outputs);
            combined_outputs
                .slice_mut(s![.., .., self.hidden_size..])
                .assign(&backward_outputs);
        } else {
            // Add forward and backward outputs
            combined_outputs = forward_outputs + backward_outputs;
        }

        Ok(combined_outputs)
    }
}

impl<T: FloatBounds + scirs2_core::ndarray::ScalarOperand> BidirectionalRNN<T, GRUCell<T>> {
    /// Create a new bidirectional GRU
    pub fn new_gru(
        input_size: usize,
        hidden_size: usize,
        concat_outputs: bool,
    ) -> NeuralResult<Self> {
        let forward_cell = GRUCell::new(input_size, hidden_size)?;
        let backward_cell = GRUCell::new(input_size, hidden_size)?;

        Ok(Self {
            forward_cell,
            backward_cell,
            hidden_size,
            concat_outputs,
            _phantom: std::marker::PhantomData,
        })
    }

    /// Forward pass through bidirectional GRU
    pub fn forward_sequence(&mut self, inputs: &Array3<T>) -> NeuralResult<Array3<T>> {
        let (batch_size, seq_len, _) = inputs.dim();

        // Forward pass
        let forward_outputs = self.forward_cell.forward_sequence(inputs)?;

        // Backward pass (reverse the input sequence)
        let mut reversed_inputs = inputs.clone();
        reversed_inputs.invert_axis(Axis(1));
        let mut backward_outputs = self.backward_cell.forward_sequence(&reversed_inputs)?;
        // Reverse the outputs back
        backward_outputs.invert_axis(Axis(1));

        // Combine outputs
        let output_size = if self.concat_outputs {
            self.hidden_size * 2
        } else {
            self.hidden_size
        };

        let mut combined_outputs = Array3::zeros((batch_size, seq_len, output_size));

        if self.concat_outputs {
            // Concatenate forward and backward outputs
            combined_outputs
                .slice_mut(s![.., .., ..self.hidden_size])
                .assign(&forward_outputs);
            combined_outputs
                .slice_mut(s![.., .., self.hidden_size..])
                .assign(&backward_outputs);
        } else {
            // Add forward and backward outputs
            combined_outputs = forward_outputs + backward_outputs;
        }

        Ok(combined_outputs)
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use approx::assert_abs_diff_eq;
    use scirs2_core::ndarray::array;

    #[test]
    #[ignore]
    fn test_lstm_cell_creation() {
        let lstm = LSTMCell::<f64>::new(10, 20).expect("construction should succeed");
        assert_eq!(lstm.input_size, 10);
        assert_eq!(lstm.hidden_size, 20);
        assert_eq!(lstm.num_parameters(), 4 * (20 * 30 + 20)); // 4 gates * (weights + biases)
    }

    #[test]
    #[ignore]
    fn test_lstm_forward_step() {
        let mut lstm = LSTMCell::<f64>::new(3, 2).expect("construction should succeed");
        let input = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]; // batch_size=2, input_size=3

        let output = lstm.forward_step(&input).expect("operation should succeed");
        assert_eq!(output.dim(), (2, 2)); // batch_size=2, hidden_size=2

        // Check that state was initialized
        assert!(lstm.state.is_some());
        let state = lstm.state.as_ref().expect("operation should succeed");
        assert_eq!(state.hidden.dim(), (2, 2));
        assert_eq!(state.cell.dim(), (2, 2));
    }

    #[test]
    #[ignore]
    fn test_lstm_sequence_processing() {
        let mut lstm = LSTMCell::<f64>::new(3, 2).expect("construction should succeed");
        // batch_size=2, seq_len=4, input_size=3
        let inputs = Array3::zeros((2, 4, 3));

        let outputs = lstm
            .forward_sequence(&inputs)
            .expect("operation should succeed");
        assert_eq!(outputs.dim(), (2, 4, 2)); // batch_size=2, seq_len=4, hidden_size=2
    }

    #[test]
    #[ignore]
    fn test_lstm_reset_state() {
        let mut lstm = LSTMCell::<f64>::new(3, 2).expect("construction should succeed");
        lstm.init_state(2);

        // Set some non-zero values
        lstm.state
            .as_mut()
            .expect("operation should succeed")
            .hidden[[0, 0]] = 1.0;
        lstm.state.as_mut().expect("operation should succeed").cell[[0, 0]] = 2.0;

        lstm.reset_state();

        let state = lstm.state.as_ref().expect("operation should succeed");
        assert_eq!(state.hidden[[0, 0]], 0.0);
        assert_eq!(state.cell[[0, 0]], 0.0);
    }

    #[test]
    #[ignore]
    fn test_gru_cell_creation() {
        let gru = GRUCell::<f64>::new(10, 20).expect("construction should succeed");
        assert_eq!(gru.input_size, 10);
        assert_eq!(gru.hidden_size, 20);
        assert_eq!(gru.num_parameters(), 3 * (20 * 30 + 20)); // 3 gates * (weights + biases)
    }

    #[test]
    #[ignore]
    fn test_gru_forward_step() {
        let mut gru = GRUCell::<f64>::new(3, 2).expect("construction should succeed");
        let input = array![[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]; // batch_size=2, input_size=3

        let output = gru.forward_step(&input).expect("operation should succeed");
        assert_eq!(output.dim(), (2, 2)); // batch_size=2, hidden_size=2

        // Check that state was initialized
        assert!(gru.hidden_state.is_some());
        let state = gru.hidden_state.as_ref().expect("operation should succeed");
        assert_eq!(state.dim(), (2, 2));
    }

    #[test]
    #[ignore]
    fn test_gru_sequence_processing() {
        let mut gru = GRUCell::<f64>::new(3, 2).expect("construction should succeed");
        // batch_size=2, seq_len=4, input_size=3
        let inputs = Array3::zeros((2, 4, 3));

        let outputs = gru
            .forward_sequence(&inputs)
            .expect("operation should succeed");
        assert_eq!(outputs.dim(), (2, 4, 2)); // batch_size=2, seq_len=4, hidden_size=2
    }

    #[test]
    #[ignore]
    fn test_gru_reset_state() {
        let mut gru = GRUCell::<f64>::new(3, 2).expect("construction should succeed");
        gru.init_state(2);

        // Set some non-zero values
        gru.hidden_state.as_mut().expect("operation should succeed")[[0, 0]] = 1.0;

        gru.reset_state();

        let state = gru.hidden_state.as_ref().expect("operation should succeed");
        assert_eq!(state[[0, 0]], 0.0);
    }

    #[test]
    #[ignore]
    fn test_bidirectional_lstm() {
        let mut bi_lstm: BidirectionalRNN<f64, LSTMCell<f64>> =
            BidirectionalRNN::new_lstm(3, 2, true).expect("operation should succeed");
        let inputs = Array3::zeros((2, 4, 3)); // batch_size=2, seq_len=4, input_size=3

        let outputs = bi_lstm
            .forward_sequence(&inputs)
            .expect("operation should succeed");
        assert_eq!(outputs.dim(), (2, 4, 4)); // batch_size=2, seq_len=4, hidden_size*2=4
    }

    #[test]
    #[ignore]
    fn test_bidirectional_gru() {
        let mut bi_gru: BidirectionalRNN<f64, GRUCell<f64>> =
            BidirectionalRNN::new_gru(3, 2, true).expect("operation should succeed");
        let inputs = Array3::zeros((2, 4, 3)); // batch_size=2, seq_len=4, input_size=3

        let outputs = bi_gru
            .forward_sequence(&inputs)
            .expect("operation should succeed");
        assert_eq!(outputs.dim(), (2, 4, 4)); // batch_size=2, seq_len=4, hidden_size*2=4
    }

    #[test]
    #[ignore]
    fn test_lstm_state_management() {
        let mut lstm = LSTMCell::<f64>::new(2, 3).expect("construction should succeed");

        // Test initial state
        assert!(lstm.get_state().is_none());

        // Initialize state
        lstm.init_state(1);
        assert!(lstm.get_state().is_some());

        let state = lstm.get_state().expect("operation should succeed");
        assert_eq!(state.hidden.dim(), (1, 3));
        assert_eq!(state.cell.dim(), (1, 3));

        // Test setting custom state
        let custom_state = LSTMState {
            hidden: array![[1.0, 2.0, 3.0]],
            cell: array![[4.0, 5.0, 6.0]],
        };
        lstm.set_state(custom_state);

        let new_state = lstm.get_state().expect("operation should succeed");
        assert_abs_diff_eq!(new_state.hidden[[0, 0]], 1.0, epsilon = 1e-10);
        assert_abs_diff_eq!(new_state.cell[[0, 0]], 4.0, epsilon = 1e-10);
    }

    #[test]
    #[ignore]
    fn test_input_size_validation() {
        let mut lstm = LSTMCell::<f64>::new(3, 2).expect("construction should succeed");
        let wrong_input = array![[1.0, 2.0]]; // input_size=2, but expected 3

        let result = lstm.forward_step(&wrong_input);
        assert!(result.is_err());
    }
}