//! LSTM-augmented networks for Recurrent PPO
//!
//! Wraps `cuda_nn::CudaLSTM` and `cuda_nn::CudaLinear` to provide LSTM-based
//! policy and value networks. All computation is GPU-native via cuBLAS.

use ml_core::MLError;

use crate::cuda_nn::{GpuContext, CudaLinear, CudaLSTM, cuda_relu, cuda_from_slice};

// Re-export cuda_nn LSTM for direct access.
pub use crate::cuda_nn::CudaLSTM as CudaLSTMExport;

/// LSTM-augmented policy network for temporal action selection
///
/// Architecture: Linear -> `ReLU` -> LSTM layers -> Linear (logits)
pub struct LSTMPolicyNetwork {
    /// Input projection layer
    input_layer: CudaLinear,
    /// Stacked LSTM layers
    lstm_layers: Vec<CudaLSTM>,
    /// Output layer (logits)
    output_layer: CudaLinear,
    /// GPU context
    ctx: GpuContext,
    /// Hidden dim (for state initialization)
    hidden_dim: usize,
    /// Number of LSTM layers
    num_layers: usize,
}

impl std::fmt::Debug for LSTMPolicyNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LSTMPolicyNetwork")
            .field("num_layers", &self.num_layers)
            .field("hidden_dim", &self.hidden_dim)
            .finish()
    }
}

impl LSTMPolicyNetwork {
    /// Create new LSTM policy network
    pub fn new(
        input_dim: usize,
        hidden_dim: usize,
        num_layers: usize,
        output_dim: usize,
    ) -> Result<Self, MLError> {
        let ctx = GpuContext::new()?;

        let input_layer = CudaLinear::new(ctx.clone(), input_dim, hidden_dim)?;

        let mut lstm_layers = Vec::new();
        for _ in 0..num_layers {
            let lstm = CudaLSTM::new(ctx.clone(), hidden_dim, hidden_dim)?;
            lstm_layers.push(lstm);
        }

        let output_layer = CudaLinear::new(ctx.clone(), hidden_dim, output_dim)?;

        Ok(Self {
            input_layer,
            lstm_layers,
            output_layer,
            ctx,
            hidden_dim,
            num_layers,
        })
    }

    /// Forward pass with LSTM hidden state propagation.
    ///
    /// # Arguments
    /// * `state_flat` - Flat state data `[batch_size * input_dim]`
    /// * `h_flat` - Hidden state `[num_layers * batch_size * hidden_dim]`
    /// * `c_flat` - Cell state `[num_layers * batch_size * hidden_dim]`
    /// * `batch_size` - Number of samples
    ///
    /// # Returns
    /// (`logits_host`, `new_h_flat`, `new_c_flat`)
    #[allow(clippy::type_complexity)]
    pub fn forward(
        &self,
        state_flat: &[f32],
        h_flat: &[f32],
        c_flat: &[f32],
        batch_size: usize,
    ) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>), MLError> {
        // Project input
        let input_gpu = cuda_from_slice(&self.ctx.stream, state_flat)?;
        let projected = self.input_layer.forward(&input_gpu.data, batch_size)?;
        let mut x = cuda_relu(&self.ctx.stream, &projected.data, projected.len)?;

        // Process through LSTM layers
        let hc_size = batch_size * self.hidden_dim;
        let mut all_new_h = Vec::with_capacity(self.num_layers * hc_size);
        let mut all_new_c = Vec::with_capacity(self.num_layers * hc_size);

        for layer_idx in 0..self.num_layers {
            let offset = layer_idx * hc_size;
            let h_slice = h_flat.get(offset..offset + hc_size).ok_or_else(|| {
                MLError::InvalidInput(format!(
                    "h_flat too short for layer {}: need {} but got {}",
                    layer_idx, offset + hc_size, h_flat.len()
                ))
            })?;
            let c_slice = c_flat.get(offset..offset + hc_size).ok_or_else(|| {
                MLError::InvalidInput(format!(
                    "c_flat too short for layer {}: need {} but got {}",
                    layer_idx, offset + hc_size, c_flat.len()
                ))
            })?;

            let h_gpu = cuda_from_slice(&self.ctx.stream, h_slice)?;
            let c_gpu = cuda_from_slice(&self.ctx.stream, c_slice)?;

            let lstm = self.lstm_layers.get(layer_idx).ok_or_else(|| {
                MLError::ConfigError(format!("LSTM layer {} not found", layer_idx))
            })?;

            let (h_new, c_new) = lstm.forward(&x.data, &h_gpu.data, &c_gpu.data, batch_size)?;

            let h_host = h_new.to_vec(&self.ctx.stream)?;
            let c_host = c_new.to_vec(&self.ctx.stream)?;
            all_new_h.extend_from_slice(&h_host);
            all_new_c.extend_from_slice(&c_host);

            x = h_new;
        }

        // Output projection
        let logits = self.output_layer.forward(&x.data, batch_size)?;
        let logits_host = logits.to_vec(&self.ctx.stream)?;

        Ok((logits_host, all_new_h, all_new_c))
    }

    /// Get hidden dimension
    pub const fn hidden_dim(&self) -> usize {
        self.hidden_dim
    }

    /// Get number of layers
    pub const fn num_layers(&self) -> usize {
        self.num_layers
    }
}

/// LSTM-augmented value network for temporal state value estimation
///
/// Architecture: Linear -> `ReLU` -> LSTM layers -> Linear (1)
pub struct LSTMValueNetwork {
    /// Input projection layer
    input_layer: CudaLinear,
    /// Stacked LSTM layers
    lstm_layers: Vec<CudaLSTM>,
    /// Output layer (value)
    output_layer: CudaLinear,
    /// GPU context
    ctx: GpuContext,
    /// Hidden dim
    hidden_dim: usize,
    /// Number of LSTM layers
    num_layers: usize,
}

impl std::fmt::Debug for LSTMValueNetwork {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("LSTMValueNetwork")
            .field("num_layers", &self.num_layers)
            .field("hidden_dim", &self.hidden_dim)
            .finish()
    }
}

impl LSTMValueNetwork {
    /// Create new LSTM value network
    pub fn new(
        input_dim: usize,
        hidden_dim: usize,
        num_layers: usize,
    ) -> Result<Self, MLError> {
        let ctx = GpuContext::new()?;

        let input_layer = CudaLinear::new(ctx.clone(), input_dim, hidden_dim)?;

        let mut lstm_layers = Vec::new();
        for _ in 0..num_layers {
            let lstm = CudaLSTM::new(ctx.clone(), hidden_dim, hidden_dim)?;
            lstm_layers.push(lstm);
        }

        let output_layer = CudaLinear::new(ctx.clone(), hidden_dim, 1)?;

        Ok(Self {
            input_layer,
            lstm_layers,
            output_layer,
            ctx,
            hidden_dim,
            num_layers,
        })
    }

    /// Forward pass returning (`values_host`, `new_h_flat`, `new_c_flat`)
    #[allow(clippy::type_complexity)]
    pub fn forward(
        &self,
        state_flat: &[f32],
        h_flat: &[f32],
        c_flat: &[f32],
        batch_size: usize,
    ) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>), MLError> {
        let input_gpu = cuda_from_slice(&self.ctx.stream, state_flat)?;
        let projected = self.input_layer.forward(&input_gpu.data, batch_size)?;
        let mut x = cuda_relu(&self.ctx.stream, &projected.data, projected.len)?;

        let hc_size = batch_size * self.hidden_dim;
        let mut all_new_h = Vec::with_capacity(self.num_layers * hc_size);
        let mut all_new_c = Vec::with_capacity(self.num_layers * hc_size);

        for layer_idx in 0..self.num_layers {
            let offset = layer_idx * hc_size;
            let h_slice = h_flat.get(offset..offset + hc_size).ok_or_else(|| {
                MLError::InvalidInput(format!("h_flat too short for layer {}", layer_idx))
            })?;
            let c_slice = c_flat.get(offset..offset + hc_size).ok_or_else(|| {
                MLError::InvalidInput(format!("c_flat too short for layer {}", layer_idx))
            })?;

            let h_gpu = cuda_from_slice(&self.ctx.stream, h_slice)?;
            let c_gpu = cuda_from_slice(&self.ctx.stream, c_slice)?;

            let lstm = self.lstm_layers.get(layer_idx).ok_or_else(|| {
                MLError::ConfigError(format!("LSTM layer {} not found", layer_idx))
            })?;

            let (h_new, c_new) = lstm.forward(&x.data, &h_gpu.data, &c_gpu.data, batch_size)?;

            let h_host = h_new.to_vec(&self.ctx.stream)?;
            let c_host = c_new.to_vec(&self.ctx.stream)?;
            all_new_h.extend_from_slice(&h_host);
            all_new_c.extend_from_slice(&c_host);

            x = h_new;
        }

        let value_out = self.output_layer.forward(&x.data, batch_size)?;
        let values_host = value_out.to_vec(&self.ctx.stream)?;

        Ok((values_host, all_new_h, all_new_c))
    }

    /// Get hidden dimension
    pub const fn hidden_dim(&self) -> usize {
        self.hidden_dim
    }

    /// Get number of layers
    pub const fn num_layers(&self) -> usize {
        self.num_layers
    }
}
