//! xLSTM network -- stack of xLSTM blocks with output projection.
//!
//! Processes sequential input `(batch, seq_len, features)` through blocks,
//! takes the final hidden state, and projects to `output_dim`.
//!
//! Block allocation: blocks 0..N*`slstm_ratio` are sLSTM, rest are mLSTM.
//! The first block has `input_dim` -> `hidden_dim`, subsequent blocks are `hidden_dim` -> `hidden_dim`.

use std::sync::Arc;

use cudarc::driver::CudaStream;
use ml_core::MLError;

use crate::gpu_tensor::{gpu_select_dim1, GpuLinear, GpuTensor};
use super::block::XLSTMBlock;
use super::config::XLSTMConfig;

/// xLSTM network.
#[derive(Debug)]
pub struct XLSTMNetwork {
    /// Input projection: features -> `hidden_dim` (only if features != `hidden_dim`)
    input_proj: Option<GpuLinear>,
    blocks: Vec<XLSTMBlock>,
    output_head: GpuLinear,
    hidden_dim: usize,
}

impl XLSTMNetwork {
    /// Build an xLSTM network from config.
    pub fn new(config: &XLSTMConfig, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
        if config.num_blocks == 0 {
            return Err(MLError::ConfigError("num_blocks must be >= 1".to_owned()));
        }
        if config.input_dim == 0 || config.hidden_dim == 0 {
            return Err(MLError::ConfigError(format!(
                "xLSTM requires input_dim > 0 and hidden_dim > 0 (got {}x{})",
                config.input_dim, config.hidden_dim
            )));
        }

        let num_slstm = ((config.num_blocks as f64) * config.slstm_ratio).round() as usize;

        let input_proj = (config.input_dim != config.hidden_dim)
            .then(|| GpuLinear::new(config.input_dim, config.hidden_dim, stream))
            .transpose()?;

        let mut blocks = Vec::with_capacity(config.num_blocks);
        for i in 0..config.num_blocks {
            let block = if i < num_slstm {
                XLSTMBlock::new_slstm(config.hidden_dim, config.hidden_dim, stream)?
            } else {
                XLSTMBlock::new_mlstm(
                    config.hidden_dim,
                    config.hidden_dim,
                    config.num_heads,
                    stream,
                )?
            };
            blocks.push(block);
        }

        let output_head = GpuLinear::new(config.hidden_dim, config.output_dim, stream)?;

        Ok(Self {
            input_proj,
            blocks,
            output_head,
            hidden_dim: config.hidden_dim,
        })
    }

    /// Forward pass.
    ///
    /// * `input` - `(batch, seq_len, features)` or `(batch, features)` for single-step.
    ///
    /// Returns `(batch, output_dim)`.
    pub fn forward(&self, input: &GpuTensor) -> Result<GpuTensor, MLError> {
        let (batch, seq_len, _features) = match input.shape.len() {
            2 => (input.dim(0)?, 1_usize, input.dim(1)?),
            3 => (input.dim(0)?, input.dim(1)?, input.dim(2)?),
            _ => {
                return Err(MLError::ModelError(format!(
                    "xLSTM expects 2D or 3D input, got {}D",
                    input.shape.len()
                )))
            }
        };

        // Process each timestep through all blocks
        let mut states: Vec<Option<(GpuTensor, GpuTensor)>> = vec![None; self.blocks.len()];
        let mut final_h = GpuTensor::zeros(&[batch, self.hidden_dim], &input.stream)?;

        for t in 0..seq_len {
            // Extract timestep on GPU: (batch, features)
            let x_t = if input.shape.len() == 3 {
                gpu_select_dim1(input, t)?
            } else {
                input.clone()
            };

            // Input projection
            let mut h = match &self.input_proj {
                Some(proj) => proj.forward(&x_t)?,
                None => x_t,
            };

            // Process through blocks
            for (i, block) in self.blocks.iter().enumerate() {
                let state_ref = states.get(i)
                    .and_then(|s| s.as_ref())
                    .map(|(h_s, c_s)| (h_s, c_s));

                let (h_new, c_new) = block.forward(&h, state_ref)?;

                if let Some(slot) = states.get_mut(i) {
                    *slot = Some((h_new.clone(), c_new));
                }
                h = h_new;
            }

            final_h = h;
        }

        // Output projection
        self.output_head.forward(&final_h)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use cudarc::driver::CudaContext;

    fn test_stream() -> Arc<CudaStream> {
        let ctx = CudaContext::new(0).expect("CUDA context required");
        ctx.new_stream().expect("Failed to create CUDA stream")
    }

    fn small_config() -> XLSTMConfig {
        XLSTMConfig {
            input_dim: 8,
            hidden_dim: 16,
            num_blocks: 2,
            num_heads: 2,
            slstm_ratio: 0.5,
            output_dim: 1,
            dropout: 0.0,
            learning_rate: 1e-3,
            weight_decay: 1e-4,
            grad_clip: 1.0,
        }
    }

    #[test]
    fn test_xlstm_network_forward_3d() {
        let stream = test_stream();
        let config = small_config();
        let net = XLSTMNetwork::new(&config, &stream).unwrap();

        let input = GpuTensor::randn(&[2, 4, 8], 0.5, &stream).unwrap();
        let out = net.forward(&input).unwrap();
        assert_eq!(out.shape, vec![2, 1]);
    }

    #[test]
    fn test_xlstm_network_forward_2d() {
        let stream = test_stream();
        let config = small_config();
        let net = XLSTMNetwork::new(&config, &stream).unwrap();

        let input = GpuTensor::randn(&[2, 8], 0.5, &stream).unwrap();
        let out = net.forward(&input).unwrap();
        assert_eq!(out.shape, vec![2, 1]);
    }

    #[test]
    fn test_xlstm_block_allocation() {
        let stream = test_stream();
        let config = XLSTMConfig {
            num_blocks: 4,
            slstm_ratio: 0.5,
            ..small_config()
        };
        let net = XLSTMNetwork::new(&config, &stream).unwrap();
        assert_eq!(net.blocks.len(), 4);
    }

    #[test]
    fn test_xlstm_zero_blocks_errors() {
        let stream = test_stream();
        let config = XLSTMConfig {
            num_blocks: 0,
            ..small_config()
        };
        let result = XLSTMNetwork::new(&config, &stream);
        assert!(result.is_err());
    }
}
