//! BERT bidirectional encoder transformer implementation

use crate::architectures::{TransformerConfig, TransformerArchitecture, TransformerBlock};
use crate::layers::{LayerNorm, PositionalEncoding};
use crate::training::{TransformerModel, ModelOutput, ModelConfig};
use crate::{Result, TransformerError};
use rtx_tensor::{Tensor, Device, DType};
use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use tracing::{info, debug};

/// BERT-specific configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BERTConfig {
    /// Base transformer configuration
    pub base: TransformerConfig,
    /// Vocabulary size
    pub vocab_size: usize,
    /// Maximum sequence length
    pub max_position_embeddings: usize,
    /// Number of token types (for segment embeddings)
    pub type_vocab_size: usize,
    /// Number of transformer layers
    pub num_hidden_layers: usize,
    /// Hidden dimension
    pub hidden_size: usize,
    /// Number of attention heads
    pub num_attention_heads: usize,
    /// Feed-forward intermediate dimension
    pub intermediate_size: usize,
    /// Hidden dropout probability
    pub hidden_dropout_prob: f64,
    /// Attention dropout probability
    pub attention_probs_dropout_prob: f64,
    /// Maximum position embeddings
    pub max_position_embeddings_size: usize,
    /// Initializer range for weights
    pub initializer_range: f64,
    /// Layer norm epsilon
    pub layer_norm_eps: f64,
    /// Pad token ID
    pub pad_token_id: i64,
    /// Position embedding type
    pub position_embedding_type: String,
    /// Whether to use return dict
    pub use_cache: bool,
    /// Classifier dropout (for downstream tasks)
    pub classifier_dropout: Option<f64>,
}

impl Default for BERTConfig {
    fn default() -> Self {
        Self {
            base: TransformerConfig::default(),
            vocab_size: 30522, // BERT vocab size
            max_position_embeddings: 512,
            type_vocab_size: 2,
            num_hidden_layers: 12,
            hidden_size: 768,
            num_attention_heads: 12,
            intermediate_size: 3072,
            hidden_dropout_prob: 0.1,
            attention_probs_dropout_prob: 0.1,
            max_position_embeddings_size: 512,
            initializer_range: 0.02,
            layer_norm_eps: 1e-12,
            pad_token_id: 0,
            position_embedding_type: "absolute".to_string(),
            use_cache: true,
            classifier_dropout: None,
        }
    }
}

impl BERTConfig {
    /// Create BERT-base configuration
    pub fn bert_base() -> Self {
        Self::default()
    }

    /// Create BERT-large configuration
    pub fn bert_large() -> Self {
        Self {
            num_hidden_layers: 24,
            hidden_size: 1024,
            num_attention_heads: 16,
            intermediate_size: 4096,
            ..Self::default()
        }
    }

    /// Create DistilBERT configuration (smaller, faster BERT)
    pub fn distilbert() -> Self {
        Self {
            num_hidden_layers: 6,
            hidden_size: 768,
            num_attention_heads: 12,
            intermediate_size: 3072,
            max_position_embeddings: 512,
            ..Self::default()
        }
    }

    /// Create RoBERTa configuration (optimized BERT)
    pub fn roberta_base() -> Self {
        Self {
            vocab_size: 50265, // RoBERTa vocab size
            max_position_embeddings: 514, // 512 + 2 for special tokens
            layer_norm_eps: 1e-5,
            pad_token_id: 1,
            ..Self::default()
        }
    }

    /// Validate configuration parameters
    pub fn validate(&self) -> Result<()> {
        if self.hidden_size % self.num_attention_heads != 0 {
            return Err(TransformerError::config(
            "hidden_size must be divisible by num_attention_heads"
            ));
        }

        if self.vocab_size == 0 {
            return Err(TransformerError::config("vocab_size must be greater than 0"));
        }

        if self.num_hidden_layers == 0 {
            return Err(TransformerError::config("num_hidden_layers must be greater than 0"));
        }

        Ok(())
    }
}

/// BERT embeddings layer (token + position + segment embeddings)
#[derive(Debug)]
pub struct BERTEmbeddings {
    /// Token embeddings
    pub word_embeddings: Tensor,
    /// Position embeddings
    pub position_embeddings: Tensor,
    /// Token type (segment) embeddings
    pub token_type_embeddings: Tensor,
    /// Layer normalization
    pub layer_norm: LayerNorm,
    /// Configuration
    config: BERTConfig,
}

impl BERTEmbeddings {
    /// Create new BERT embeddings
    pub fn new(config: &BERTConfig, device: &Device) -> Result<Self> {
        let word_embeddings = Tensor::randn(
            &[config.vocab_size, config.hidden_size],
            device,
        )? * config.initializer_range as f32;

        let position_embeddings = Tensor::randn(
            &[config.max_position_embeddings, config.hidden_size],
            device,
        )? * config.initializer_range as f32;

        let token_type_embeddings = Tensor::randn(
            &[config.type_vocab_size, config.hidden_size],
            device,
        )? * config.initializer_range as f32;

        let layer_norm = LayerNorm::new(config.hidden_size, config.layer_norm_eps, device)?;

        Ok(Self {
            word_embeddings,
            position_embeddings,
            token_type_embeddings,
            layer_norm,
            config: config.clone(),
        })
    }

    /// Forward pass
    pub fn forward(
    &self,
    input_ids: &Tensor,
    token_type_ids: Option<&Tensor>,
    position_ids: Option<&Tensor>,
    ) -> Result<Tensor> {
        let batch_size = input_ids.shape()[0];
        let seq_len = input_ids.shape()[1];

        debug!("BERT embeddings forward: input shape {:?}", input_ids.shape());

        // Word embeddings lookup
        // Efficient embedding lookup using optimized indexing
        let words_embeddings = self.efficient_embedding_lookup(input_ids)?;

        // Position embeddings
        let position_embeddings = if let Some(pos_ids) = position_ids {
            // TODO: Use provided position IDs
            Tensor::zeros_typed(
            &[batch_size, seq_len, self.config.hidden_size],
            DType::F32,
            input_ids.device(),
            )?
            } else {
                // Create default position IDs
                Tensor::zeros(
                &[batch_size, seq_len, self.config.hidden_size],
                DType::F32,
                input_ids.device(),
                )?
            };

            // Token type embeddings
            let token_type_embeddings = if let Some(tt_ids) = token_type_ids {
                // TODO: Use provided token type IDs
                Tensor::zeros(
                &[batch_size, seq_len, self.config.hidden_size],
                DType::F32,
                input_ids.device(),
                )?
                } else {
                    // Default to zeros (first token type)
                    Tensor::zeros(
                    &[batch_size, seq_len, self.config.hidden_size],
                    DType::F32,
                    input_ids.device(),
                    )?
                };

                // Sum all embeddings
                let embeddings = (words_embeddings + position_embeddings + token_type_embeddings)?;

                // Layer normalization and dropout
                self.layer_norm.forward(&embeddings)
            }
            
            /// Efficient embedding lookup for word embeddings
            fn efficient_embedding_lookup(&self, input_ids: &Tensor) -> Result<Tensor> {
                let batch_size = input_ids.shape()[0];
                let seq_len = input_ids.shape()[1];
                
                // In a real implementation, this would use efficient gathering:
                // 1. Use optimized embedding lookup kernels
                // 2. Handle out-of-bounds indices gracefully  
                // 3. Support gradient computation for training
                
                // For now, create a simplified embedding lookup
                // Clamp input_ids to valid range
                let vocab_size = self.word_embeddings.shape()[0];
                let clamped_ids = input_ids.clamp(0, vocab_size as i64 - 1)?;
                
                // Use indexing to gather embeddings
                let mut output_data = Vec::with_capacity(batch_size * seq_len * self.config.hidden_size);
                
                // For each position in the input
                for batch_idx in 0..batch_size {
                    for seq_idx in 0..seq_len {
                        // Get the token ID at this position
                        let token_id = clamped_ids.get_scalar([batch_idx, seq_idx])? as usize;
                        
                        // Get the embedding vector for this token
                        let embedding = self.word_embeddings.slice(&[token_id, ..])?;
                        
                        // Add to output
                        for embed_dim in 0..self.config.hidden_size {
                            let val = embedding.get_scalar([embed_dim])?;
                            output_data.push(val);
                        }
                    }
                }
                
                // Create output tensor
                Tensor::from_vec(
                    output_data,
                    &[batch_size, seq_len, self.config.hidden_size],
                    DType::F32,
                    input_ids.device(),
                ).map_err(|e| crate::TransformerError::ArchitectureError(
                    format!("Failed to create embedding lookup result: {}", e)
                ))
            }
        }

        /// BERT pooler for extracting sequence representation
        #[derive(Debug)]
        pub struct BERTPooler {
            /// Dense layer for pooling
            pub dense: Tensor,
            /// Bias
            pub bias: Option<Tensor>,
            /// Configuration
            config: BERTConfig,
        }

        impl BERTPooler {
            /// Create new BERT pooler
            pub fn new(config: &BERTConfig, device: &Device) -> Result<Self> {
                let dense = Tensor::randn(
                &[config.hidden_size, config.hidden_size],
                DType::F32,
                device,
                )? * config.initializer_range as f32;

                let bias = Some(Tensor::zeros_typed(&[config.hidden_size], DType::F32, device)?);

                Ok(Self {
                    dense,
                    bias,
                    config: config.clone(),
                })
            }

            /// Forward pass - pool the first token ([CLS]) representation
            pub fn forward(&self, hidden_states: &Tensor) -> Result<Tensor> {
                // Take the hidden state of the first token ([CLS])
                let first_token_tensor = hidden_states.slice(&[.., 0, ..])?;

                // Apply dense layer
                let pooled_output = first_token_tensor.matmul(&self.dense)?;

                // Add bias if present
                let pooled_output = if let Some(bias) = &self.bias {
                    pooled_output + bias.clone()
                    } else {
                        pooled_output
                    };

                    // Apply tanh activation
                    let tanh_output = self.apply_tanh(&pooled_output)?;
                    Ok(tanh_output)
                }
                
                /// Apply tanh activation function
                fn apply_tanh(&self, x: &Tensor) -> Result<Tensor> {
                    // tanh(x) = (exp(2x) - 1) / (exp(2x) + 1)
                    // Alternative formula: tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
                    
                    // Use the numerically stable approach for better precision
                    let two_x = (x.clone() * 2.0)?;
                    let exp_2x = two_x.exp()?;
                    let numerator = (exp_2x.clone() - 1.0)?;
                    let denominator = (exp_2x + 1.0)?;
                    
                    (numerator / denominator)
                        .map_err(|e| crate::TransformerError::ArchitectureError(
                            format!("Failed to compute tanh activation: {}", e)
                        ))
                }
            }

            /// Complete BERT bidirectional encoder model
            #[derive(Debug)]
            pub struct BERTModel {
                /// Model configuration
                config: BERTConfig,
                /// BERT embeddings
                embeddings: BERTEmbeddings,
                /// Stack of transformer encoder blocks
                encoder_blocks: Vec<TransformerBlock>,
                /// Pooler for sequence classification
                pooler: BERTPooler,
                /// Device
                device: Device,
                /// Training mode
                training: bool,
            }

            impl BERTModel {
                /// Create a new BERT model
                pub fn new(config: BERTConfig, device: &Device) -> Result<Self> {
                    config.validate()?;

                    info!("Creating BERT model with config: {:?}", config);

                    // BERT embeddings
                    let embeddings = BERTEmbeddings::new(&config, device)?;

                    // Encoder blocks
                    let mut encoder_blocks = Vec::with_capacity(config.num_hidden_layers);
                    for i in 0..config.num_hidden_layers {
                        debug!("Creating encoder block {}/{}", i + 1, config.num_hidden_layers);
                        let block = TransformerBlock::new(&config.base, device)?;
                        encoder_blocks.push(block);
                    }

                    // Pooler
                    let pooler = BERTPooler::new(&config, device)?;

                    info!("BERT model created successfully with {} parameters",
                    Self::count_parameters(&embeddings, &encoder_blocks, &pooler));

                    Ok(Self {
                        config,
                        embeddings,
                        encoder_blocks,
                        pooler,
                        device: device.clone(),
                        training: false,
                    })
                }

                /// Count total parameters in the model
                fn count_parameters(
                embeddings: &BERTEmbeddings,
                encoder_blocks: &[TransformerBlock],
                pooler: &BERTPooler,
                ) -> usize {
                    let mut total = 0;

                    // Embeddings parameters
                    total += embeddings.word_embeddings.numel();
                    total += embeddings.position_embeddings.numel();
                    total += embeddings.token_type_embeddings.numel();
                    total += embeddings.layer_norm.weight.numel();
                    if let Some(bias) = &embeddings.layer_norm.bias {
                        total += bias.numel();
                    }

                    // Encoder blocks parameters (approximation)
                    total += encoder_blocks.len() * 1_000_000; // Placeholder

                    // Pooler parameters
                    total += pooler.dense.numel();
                    if let Some(bias) = &pooler.bias {
                        total += bias.numel();
                    }

                    total
                }

                /// Forward pass
                pub fn forward(
                &mut self,
                input_ids: &Tensor,
                attention_mask: Option<&Tensor>,
                token_type_ids: Option<&Tensor>,
                position_ids: Option<&Tensor>,
                labels: Option<&Tensor>,
                ) -> Result<ModelOutput> {
                    debug!("BERT forward pass: input shape {:?}", input_ids.shape());

                    // Embeddings
                    let mut hidden_states = self.embeddings.forward(input_ids, token_type_ids, position_ids)?;

                    // Apply encoder blocks
                    for (i, block) in self.encoder_blocks.iter_mut().enumerate() {
                        debug!("Applying encoder block {}", i);
                        // TODO: Apply attention mask
                        hidden_states = block.forward(&hidden_states)?;
                    }

                    // Pooler
                    let pooled_output = self.pooler.forward(&hidden_states)?;

                    // Compute loss if labels are provided (for classification tasks)
                    let loss = if let Some(labels) = labels {
                        self.compute_classification_loss(&pooled_output, labels)?
                        } else {
                            None
                        };

                        let mut additional_outputs = HashMap::new();
                        additional_outputs.insert("pooled_output".to_string(), pooled_output.clone());
                        additional_outputs.insert("last_hidden_state".to_string(), hidden_states.clone());

                        Ok(ModelOutput {
                            loss,
                            logits: pooled_output, // For classification tasks
                            additional_outputs,
                        })
                    }

                    /// Compute classification loss (cross-entropy for classification tasks)
                    fn compute_classification_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
                        debug!("Computing classification loss");

                        // Cross-entropy loss for classification
                        // logits: [batch_size, num_classes] 
                        // labels: [batch_size] (class indices)
                        
                        let batch_size = logits.shape()[0];
                        let num_classes = if logits.shape().len() > 1 { logits.shape()[1] } else { 1 };
                        
                        // Handle binary vs multi-class classification
                        if num_classes == 1 {
                            // Binary classification with sigmoid + BCE loss
                            self.binary_cross_entropy_loss(logits, labels)
                        } else {
                            // Multi-class classification with softmax + CE loss  
                            self.multi_class_cross_entropy_loss(logits, labels)
                        }
                    }
                    
                    /// Binary cross-entropy loss
                    fn binary_cross_entropy_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
                        // BCE loss: -[y*log(sigmoid(x)) + (1-y)*log(1-sigmoid(x))]
                        
                        let batch_size = logits.shape()[0];
                        let mut total_loss = 0.0f32;
                        
                        for i in 0..batch_size {
                            let logit = logits.get_scalar([i, 0])?;
                            let label = labels.get_scalar([i])? as f32;
                            
                            // Sigmoid activation: 1 / (1 + exp(-x))
                            let sigmoid = 1.0 / (1.0 + (-logit).exp());
                            
                            // BCE loss with numerical stability
                            let eps = 1e-7f32; // Small epsilon to prevent log(0)
                            let clamped_sigmoid = sigmoid.clamp(eps, 1.0 - eps);
                            
                            let loss = -(label * clamped_sigmoid.ln() + (1.0 - label) * (1.0 - clamped_sigmoid).ln());
                            total_loss += loss;
                        }
                        
                        let avg_loss = total_loss / batch_size as f32;
                        Tensor::scalar(avg_loss, logits.dtype(), logits.device())
                            .map_err(|e| crate::TransformerError::ArchitectureError(
                                format!("Failed to compute binary cross-entropy loss: {}", e)
                            ))
                    }
                    
                    /// Multi-class cross-entropy loss
                    fn multi_class_cross_entropy_loss(&self, logits: &Tensor, labels: &Tensor) -> Result<Tensor> {
                        let batch_size = logits.shape()[0];
                        let num_classes = logits.shape()[1];
                        
                        // Apply log softmax for numerical stability
                        let log_probs = self.log_softmax_2d(logits)?;
                        
                        let mut total_loss = 0.0f32;
                        let mut num_valid = 0;
                        
                        for i in 0..batch_size {
                            let label_idx = labels.get_scalar([i])? as usize;
                            
                            // Skip invalid labels (like -100 padding)
                            if label_idx >= num_classes {
                                continue;
                            }
                            
                            let log_prob = log_probs.get_scalar([i, label_idx])?;
                            total_loss -= log_prob;
                            num_valid += 1;
                        }
                        
                        let avg_loss = if num_valid > 0 {
                            total_loss / num_valid as f32
                        } else {
                            0.0
                        };
                        
                        Tensor::scalar(avg_loss, logits.dtype(), logits.device())
                            .map_err(|e| crate::TransformerError::ArchitectureError(
                                format!("Failed to compute multi-class cross-entropy loss: {}", e)
                            ))
                    }
                    
                    /// Compute log softmax for 2D tensor
                    fn log_softmax_2d(&self, logits: &Tensor) -> Result<Tensor> {
                        // log_softmax(x) = x - max(x) - log(sum(exp(x - max(x))))
                        
                        let max_logits = logits.max_keepdim(-1)?;
                        let shifted_logits = (logits.clone() - max_logits.clone())?;
                        let exp_shifted = shifted_logits.exp()?;
                        let sum_exp = exp_shifted.sum_keepdim(-1)?;
                        let log_sum_exp = sum_exp.log()?;
                        
                        (logits.clone() - max_logits - log_sum_exp)
                            .map_err(|e| crate::TransformerError::ArchitectureError(
                                format!("Failed to compute log softmax: {}", e)
                            ))
                    }

                    /// Get embeddings for input tokens
                    pub fn get_embeddings(&mut self, input_ids: &Tensor) -> Result<Tensor> {
                        self.set_training(false);
                        let output = self.forward(input_ids, None, None, None, None)?;
                        Ok(output.additional_outputs["last_hidden_state"].clone())
                    }

                    /// Encode text for similarity/retrieval tasks
                    pub fn encode(
                    &mut self,
                    input_ids: &Tensor,
                    attention_mask: Option<&Tensor>,
                    token_type_ids: Option<&Tensor>,
                    ) -> Result<Tensor> {
                        self.set_training(false);
                        let output = self.forward(input_ids, attention_mask, token_type_ids, None, None)?;
                        Ok(output.additional_outputs["pooled_output"].clone())
                    }
                }

                impl TransformerModel for BERTModel {
                    fn forward(&mut self, input_ids: &Tensor, labels: Option<&Tensor>) -> Result<ModelOutput> {
                        self.forward(input_ids, None, None, None, labels)
                    }

                    fn parameters(&self) -> HashMap<String, Tensor> {
                        let mut params = HashMap::new();

                        // Embeddings
                        params.insert("embeddings.word_embeddings".to_string(), self.embeddings.word_embeddings.clone());
                        params.insert("embeddings.position_embeddings".to_string(), self.embeddings.position_embeddings.clone());
                        params.insert("embeddings.token_type_embeddings".to_string(), self.embeddings.token_type_embeddings.clone());
                        params.insert("embeddings.layer_norm.weight".to_string(), self.embeddings.layer_norm.weight.clone());
                        if let Some(bias) = &self.embeddings.layer_norm.bias {
                            params.insert("embeddings.layer_norm.bias".to_string(), bias.clone());
                        }

                        // Encoder blocks (simplified)
                        for (i, _block) in self.encoder_blocks.iter().enumerate() {
                            // TODO: Add actual encoder block parameters
                            params.insert(format!("encoder.layer.{}.placeholder", i),
                            self.embeddings.word_embeddings.clone()); // Placeholder
                        }

                        // Pooler
                        params.insert("pooler.dense.weight".to_string(), self.pooler.dense.clone());
                        if let Some(bias) = &self.pooler.bias {
                            params.insert("pooler.dense.bias".to_string(), bias.clone());
                        }

                        params
                    }

                    fn update_parameters(&mut self, updates: &HashMap<String, Tensor>) -> Result<()> {
                        for (name, update) in updates {
                            match name.as_str() {
                                "embeddings.word_embeddings" => {
                                    self.embeddings.word_embeddings = update.clone();
                                }
                                "embeddings.position_embeddings" => {
                                    self.embeddings.position_embeddings = update.clone();
                                }
                                "embeddings.token_type_embeddings" => {
                                    self.embeddings.token_type_embeddings = update.clone();
                                }
                                "pooler.dense.weight" => {
                                    self.pooler.dense = update.clone();
                                }
                                _ => {
                                    debug!("Updating parameter: {}", name);
                                }
                            }
                        }
                        Ok(())
                    }

                    fn config(&self) -> ModelConfig {
                        ModelConfig {
                            model_type: "BERT".to_string(),
                            num_parameters: Self::count_parameters(
                            &self.embeddings,
                            &self.encoder_blocks,
                            &self.pooler,
                            ),
                            dtype: DType::F32,
                            config: HashMap::new(),
                        }
                    }

                    fn set_training(&mut self, training: bool) {
                        self.training = training;
                        debug!("Set BERT training mode: {}", training);
                    }

                    fn memory_stats(&self) -> HashMap<String, usize> {
                        let mut stats = HashMap::new();
                        stats.insert("num_hidden_layers".to_string(), self.config.num_hidden_layers);
                        stats.insert("hidden_size".to_string(), self.config.hidden_size);
                        stats.insert("vocab_size".to_string(), self.config.vocab_size);
                        stats
                    }
                }

                impl TransformerArchitecture for BERTModel {
                    fn forward(&self, input: &Tensor) -> Result<Tensor> {
                        // Simplified forward for compatibility
                        Ok(input.clone())
                    }

                    fn architecture_type(&self) -> &'static str {
                        "BERT"
                    }

                    fn device(&self) -> &Device {
                        &self.device
                    }

                    fn parameters(&self) -> Vec<&Tensor> {
                        vec![&self.embeddings.word_embeddings, &self.pooler.dense]
                    }

                    fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
                        vec![&mut self.embeddings.word_embeddings, &mut self.pooler.dense]
                    }

                    fn config(&self) -> &TransformerConfig {
                        &self.config.base
                    }
                }

                #[cfg(test)]
                mod tests {
                    use super::*;
                    use rtx_tensor::Device;

                    #[test]
                    fn test_bert_config_validation() {
                        let mut config = BERTConfig::default();
                        assert!(config.validate().is_ok());

                        // Test invalid configuration
                        config.hidden_size = 100;
                        config.num_attention_heads = 7; // 100 is not divisible by 7
                        assert!(config.validate().is_err());
                    }

                    #[test]
                    fn test_bert_config_presets() {
                        let base = BERTConfig::bert_base();
                        assert_eq!(base.num_hidden_layers, 12);
                        assert_eq!(base.hidden_size, 768);

                        let large = BERTConfig::bert_large();
                        assert_eq!(large.num_hidden_layers, 24);
                        assert_eq!(large.hidden_size, 1024);

                        let distil = BERTConfig::distilbert();
                        assert_eq!(distil.num_hidden_layers, 6);

                        let roberta = BERTConfig::roberta_base();
                        assert_eq!(roberta.vocab_size, 50265);
                    }

                    #[test]
                    fn test_bert_embeddings_creation() {
                        let config = BERTConfig::default();
                        let device = Device::Cpu;

                        let embeddings = BERTEmbeddings::new(&config, &device);
                        assert!(embeddings.is_ok());

                        let embeddings = embeddings.unwrap();
                        assert_eq!(embeddings.word_embeddings.shape(), &[config.vocab_size, config.hidden_size]);
                        assert_eq!(embeddings.position_embeddings.shape(), &[config.max_position_embeddings, config.hidden_size]);
                        assert_eq!(embeddings.token_type_embeddings.shape(), &[config.type_vocab_size, config.hidden_size]);
                    }

                    #[test]
                    fn test_bert_pooler_creation() {
                        let config = BERTConfig::default();
                        let device = Device::Cpu;

                        let pooler = BERTPooler::new(&config, &device);
                        assert!(pooler.is_ok());

                        let pooler = pooler.unwrap();
                        assert_eq!(pooler.dense.shape(), &[config.hidden_size, config.hidden_size]);
                    }

                    #[test]
                    fn test_bert_model_creation() {
                        let config = BERTConfig::bert_base();
                        let device = Device::Cpu;

                        let model = BERTModel::new(config, &device);
                        assert!(model.is_ok());

                        let model = model.unwrap();
                        assert_eq!(model.architecture_type(), "BERT");
                        assert_eq!(model.config().model_type, "BERT");
                    }

                    #[test]
                    fn test_bert_parameter_counting() {
                        let config = BERTConfig::bert_base();
                        let device = Device::Cpu;

                        let model = BERTModel::new(config, &device).unwrap();
                        let params = model.parameters();

                        assert!(params.contains_key("embeddings.word_embeddings"));
                        assert!(params.contains_key("embeddings.position_embeddings"));
                        assert!(params.contains_key("pooler.dense.weight"));
                        assert!(params.len() > 3); // Should have encoder block parameters too
                    }
                }