//! EmbeddingGemma 完整模型实现
//!
//! 实现了完整的 EmbeddingGemma Encoder 模型，包括：
//! - Token Embedding
//! - 24 层 Transformer Encoder
//! - 最终的 RMSNorm
//! - Embedding 投影层
//!
//! # 模型结构
//!
//! ```text
//! Input IDs
//!     ↓
//! Token Embedding (vocab_size -> hidden_size)
//!     ↓
//! 24x Transformer Encoder Layer
//!     ├── RMSNorm (pre-norm)
//!     ├── Bidirectional Self-Attention + RoPE
//!     ├── Residual Connection
//!     ├── RMSNorm (pre-norm)
//!     ├── GeGLU FFN
//!     └── Residual Connection
//!     ↓
//! Final RMSNorm
//!     ↓
//! Output Hidden States
//! ```

use candle_core::{DType, Device, Module, Result, Tensor};
use candle_nn::VarBuilder;

use super::attention::{create_attention_mask, BidirectionalAttention};
use super::config::EmbeddingGemmaConfig;
use super::layers::{EmbeddingProjection, GeGLUFeedForward, RMSNorm};
use super::rope::RotaryEmbedding;

/// Transformer Encoder 层 (Gemma3 架构)
///
/// 单个 Transformer 层，包含自注意力和 FFN，使用 Gemma3 的4层 norm 结构：
/// - input_layernorm: 自注意力前
/// - post_attention_layernorm: 自注意力后残差连接前
/// - pre_feedforward_layernorm: FFN 前
/// - post_feedforward_layernorm: FFN 后残差连接前
struct EncoderLayer {
    /// 自注意力
    self_attn: BidirectionalAttention,
    /// Feed-Forward Network
    mlp: GeGLUFeedForward,
    /// 注意力前的归一化
    input_layernorm: RMSNorm,
    /// 注意力后残差前的归一化
    post_attention_layernorm: RMSNorm,
    /// FFN 前的归一化
    pre_feedforward_layernorm: RMSNorm,
    /// FFN 后残差前的归一化
    post_feedforward_layernorm: RMSNorm,
}

impl EncoderLayer {
    /// 创建 Encoder 层
    fn new(config: &EmbeddingGemmaConfig, vb: VarBuilder) -> Result<Self> {
        Ok(Self {
            self_attn: BidirectionalAttention::new(config, vb.pp("self_attn"))?,
            mlp: GeGLUFeedForward::new(
                config.hidden_size,
                config.intermediate_size,
                vb.pp("mlp"),
            )?,
            input_layernorm: RMSNorm::new(
                config.hidden_size,
                config.rms_norm_eps,
                vb.pp("input_layernorm"),
            )?,
            post_attention_layernorm: RMSNorm::new(
                config.hidden_size,
                config.rms_norm_eps,
                vb.pp("post_attention_layernorm"),
            )?,
            pre_feedforward_layernorm: RMSNorm::new(
                config.hidden_size,
                config.rms_norm_eps,
                vb.pp("pre_feedforward_layernorm"),
            )?,
            post_feedforward_layernorm: RMSNorm::new(
                config.hidden_size,
                config.rms_norm_eps,
                vb.pp("post_feedforward_layernorm"),
            )?,
        })
    }

    /// 前向传播 (Gemma3 架构)
    ///
    /// Gemma3 使用特殊的残差结构:
    /// x = x + post_attention_norm(attn(input_norm(x)))
    /// x = x + post_feedforward_norm(ffn(pre_feedforward_norm(x)))
    fn forward(
        &self,
        hidden_states: &Tensor,
        attention_mask: Option<&Tensor>,
        rope: &RotaryEmbedding,
    ) -> Result<Tensor> {
        // Self-Attention with pre and post norm
        let residual = hidden_states;
        let hidden_states = self.input_layernorm.forward(hidden_states)?;
        let hidden_states = self.self_attn.forward(&hidden_states, attention_mask, rope)?;
        let hidden_states = self.post_attention_layernorm.forward(&hidden_states)?;
        let hidden_states = (residual + hidden_states)?;

        // FFN with pre and post norm
        let residual = &hidden_states;
        let hidden_states = self.pre_feedforward_layernorm.forward(&hidden_states)?;
        let hidden_states = self.mlp.forward(&hidden_states)?;
        let hidden_states = self.post_feedforward_layernorm.forward(&hidden_states)?;
        
        residual + hidden_states
    }
}

/// EmbeddingGemma 完整模型
///
/// 基于 Gemma3 架构的双向 Encoder 模型，专为生成文本 embedding 设计。
pub struct EmbeddingGemmaModel {
    /// Token Embedding 层
    embed_tokens: candle_nn::Embedding,
    /// Transformer Encoder 层
    layers: Vec<EncoderLayer>,
    /// 最终归一化
    norm: RMSNorm,
    /// Embedding 投影层 (768 -> 3072 -> 768)
    projection: Option<EmbeddingProjection>,
    /// 旋转位置编码
    rope: RotaryEmbedding,
    /// 模型配置
    config: EmbeddingGemmaConfig,
    /// 计算设备
    device: Device,
}

impl EmbeddingGemmaModel {
    /// 从变量构建器创建模型
    ///
    /// # Arguments
    /// * `config` - 模型配置
    /// * `vb` - 变量构建器（包含预训练权重）
    pub fn new(config: EmbeddingGemmaConfig, vb: VarBuilder) -> Result<Self> {
        let device = vb.device().clone();

        // Token Embedding
        // Note: EmbeddingGemma model weights don't have "model." prefix
        let embed_tokens = candle_nn::embedding(
            config.vocab_size,
            config.hidden_size,
            vb.pp("embed_tokens"),
        )?;

        // Transformer 层
        let mut layers = Vec::with_capacity(config.num_hidden_layers);
        for i in 0..config.num_hidden_layers {
            let layer = EncoderLayer::new(&config, vb.pp(&format!("layers.{}", i)))?;
            layers.push(layer);
        }

        // 最终归一化
        let norm = RMSNorm::new(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;

        // 投影层（可选）
        let projection = if config.projection_dim > 0 {
            Some(EmbeddingProjection::new(
                config.hidden_size,
                config.projection_dim,
                vb.pp("projection"),
            )?)
        } else {
            None
        };

        // RoPE
        let rope = RotaryEmbedding::new(
            config.head_dim,
            config.max_position_embeddings,
            config.rope_theta,
            DType::F32, // RoPE 使用 F32 以获得更好的精度
            &device,
        )?;

        Ok(Self {
            embed_tokens,
            layers,
            norm,
            projection,
            rope,
            config,
            device,
        })
    }

    /// 前向传播 - 获取隐藏状态
    ///
    /// # Arguments
    /// * `input_ids` - 输入 token IDs，形状: (batch, seq_len)
    /// * `attention_mask` - 可选的注意力掩码，形状: (batch, seq_len)
    ///
    /// # Returns
    /// 隐藏状态，形状: (batch, seq_len, hidden_size)
    pub fn forward(
        &self,
        input_ids: &Tensor,
        attention_mask: Option<&Tensor>,
    ) -> Result<Tensor> {
        // Token Embedding
        let mut hidden_states = self.embed_tokens.forward(input_ids)?;

        // 处理注意力掩码
        let attention_mask = match attention_mask {
            Some(mask) => Some(create_attention_mask(mask)?),
            None => None,
        };

        // 通过所有 Transformer 层
        for layer in &self.layers {
            hidden_states = layer.forward(&hidden_states, attention_mask.as_ref(), &self.rope)?;
        }

        // 最终归一化
        self.norm.forward(&hidden_states)
    }

    /// 对隐藏状态进行 Mean Pooling
    ///
    /// # Arguments
    /// * `hidden_states` - 隐藏状态，形状: (batch, seq_len, hidden_size)
    /// * `attention_mask` - 注意力掩码，形状: (batch, seq_len)
    ///
    /// # Returns
    /// 池化后的向量，形状: (batch, hidden_size)
    pub fn mean_pooling(
        &self,
        hidden_states: &Tensor,
        attention_mask: &Tensor,
    ) -> Result<Tensor> {
        // 扩展 mask: (batch, seq) -> (batch, seq, 1)
        let mask = attention_mask
            .unsqueeze(2)?
            .to_dtype(hidden_states.dtype())?;

        // 广播到隐藏状态的形状: (batch, seq, hidden_size)
        let mask = mask.broadcast_as(hidden_states.shape())?;

        // 加权求和
        let masked = hidden_states.mul(&mask)?;
        let sum = masked.sum(1)?; // (batch, hidden_size)

        // 计算有效 token 数
        let mask_sum = mask.sum(1)?; // (batch, hidden_size)
        let mask_sum = mask_sum.clamp(1e-9, f64::MAX)?;

        // 平均
        sum.div(&mask_sum)
    }

    /// 应用投影层
    ///
    /// # Arguments
    /// * `pooled` - 池化后的向量，形状: (batch, hidden_size)
    ///
    /// # Returns
    /// 投影后的向量，形状: (batch, hidden_size)
    pub fn project(&self, pooled: &Tensor) -> Result<Tensor> {
        match &self.projection {
            Some(proj) => proj.forward(pooled),
            None => Ok(pooled.clone()),
        }
    }

    /// L2 归一化
    ///
    /// # Arguments
    /// * `embeddings` - 输入向量，形状: (batch, dim)
    ///
    /// # Returns
    /// 归一化后的向量，形状: (batch, dim)
    pub fn l2_normalize(&self, embeddings: &Tensor) -> Result<Tensor> {
        let norm = embeddings
            .sqr()?
            .sum_keepdim(1)?
            .sqrt()?
            .clamp(1e-9, f64::MAX)?;

        embeddings.broadcast_div(&norm)
    }

    /// 完整的 embedding 生成流程
    ///
    /// # Arguments
    /// * `input_ids` - 输入 token IDs，形状: (batch, seq_len)
    /// * `attention_mask` - 注意力掩码，形状: (batch, seq_len)
    /// * `matryoshka_dim` - 可选的 Matryoshka 截断维度
    ///
    /// # Returns
    /// 归一化的 embedding，形状: (batch, dim)
    pub fn encode(
        &self,
        input_ids: &Tensor,
        attention_mask: &Tensor,
        matryoshka_dim: Option<usize>,
    ) -> Result<Tensor> {
        // 1. 前向传播获取隐藏状态
        let hidden_states = self.forward(input_ids, Some(attention_mask))?;

        // 2. Mean Pooling
        let pooled = self.mean_pooling(&hidden_states, attention_mask)?;

        // 3. 投影
        let projected = self.project(&pooled)?;

        // 4. Matryoshka 截断（可选）
        let output = match matryoshka_dim {
            Some(dim) if dim < self.config.hidden_size => projected.narrow(1, 0, dim)?,
            _ => projected,
        };

        // 5. L2 归一化
        self.l2_normalize(&output)
    }

    /// 获取模型配置
    pub fn config(&self) -> &EmbeddingGemmaConfig {
        &self.config
    }

    /// 获取计算设备
    pub fn device(&self) -> &Device {
        &self.device
    }

    /// 获取隐藏层维度
    pub fn hidden_size(&self) -> usize {
        self.config.hidden_size
    }

    /// 获取层数
    pub fn num_layers(&self) -> usize {
        self.config.num_hidden_layers
    }

    /// 获取最大序列长度
    pub fn max_seq_len(&self) -> usize {
        self.config.max_position_embeddings
    }
}

/// 模型加载选项
#[derive(Debug, Clone)]
pub struct LoadOptions {
    /// 计算设备
    pub device: Device,
    /// 数据类型
    pub dtype: DType,
    /// 是否加载投影层
    pub load_projection: bool,
}

impl Default for LoadOptions {
    fn default() -> Self {
        Self {
            device: Device::Cpu,
            dtype: DType::F32,
            load_projection: true,
        }
    }
}

impl LoadOptions {
    /// 使用 CPU
    pub fn cpu() -> Self {
        Self::default()
    }

    /// 使用 CUDA (如果可用)
    pub fn cuda(ordinal: usize) -> Result<Self> {
        Ok(Self {
            device: Device::new_cuda(ordinal)?,
            ..Default::default()
        })
    }

    /// 使用 Metal (Apple Silicon，如果可用)
    pub fn metal() -> Result<Self> {
        Ok(Self {
            device: Device::new_metal(0)?,
            ..Default::default()
        })
    }

    /// 设置数据类型
    pub fn with_dtype(mut self, dtype: DType) -> Self {
        self.dtype = dtype;
        self
    }

    /// 是否加载投影层
    pub fn with_projection(mut self, load: bool) -> Self {
        self.load_projection = load;
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_nn::VarMap;

    fn create_test_varbuilder(device: &Device) -> VarBuilder<'static> {
        let varmap = VarMap::new();
        VarBuilder::from_varmap(&varmap, DType::F32, device)
    }

    fn create_small_config() -> EmbeddingGemmaConfig {
        EmbeddingGemmaConfig {
            vocab_size: 1000,
            hidden_size: 64,
            intermediate_size: 256,
            num_hidden_layers: 2,
            num_attention_heads: 2,
            num_key_value_heads: 1, // GQA
            head_dim: 32,
            max_position_embeddings: 128,
            rms_norm_eps: 1e-6,
            rope_theta: 10000.0,
            attention_bias: false,
            projection_dim: 256,
            query_pre_attn_scalar: 32,
        }
    }

    #[test]
    fn test_model_creation() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config.clone(), vb);
        assert!(model.is_ok());

        let model = model.unwrap();
        assert_eq!(model.hidden_size(), 64);
        assert_eq!(model.num_layers(), 2);
        assert_eq!(model.max_seq_len(), 128);
    }

    #[test]
    fn test_model_forward() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config, vb).unwrap();

        // 创建测试输入 (batch=2, seq=16)
        let input_ids = Tensor::zeros((2, 16), DType::U32, &device).unwrap();

        let output = model.forward(&input_ids, None).unwrap();

        // 输出形状应该是 (batch, seq, hidden_size)
        assert_eq!(output.dims(), &[2, 16, 64]);
    }

    #[test]
    fn test_model_with_attention_mask() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config, vb).unwrap();

        // 创建测试输入
        let input_ids = Tensor::zeros((2, 16), DType::U32, &device).unwrap();
        let attention_mask = Tensor::ones((2, 16), DType::F32, &device).unwrap();

        let output = model.forward(&input_ids, Some(&attention_mask)).unwrap();

        assert_eq!(output.dims(), &[2, 16, 64]);
    }

    #[test]
    fn test_mean_pooling() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config, vb).unwrap();

        // 创建测试隐藏状态 (batch=2, seq=16, hidden=64)
        let hidden_states = Tensor::randn(0.0f32, 1.0, (2, 16, 64), &device).unwrap();

        // 创建掩码 (假设后 4 个是 padding)
        let mut mask_data = vec![1.0f32; 2 * 16];
        for i in 0..2 {
            for j in 12..16 {
                mask_data[i * 16 + j] = 0.0;
            }
        }
        let attention_mask = Tensor::from_vec(mask_data, (2, 16), &device).unwrap();

        let pooled = model.mean_pooling(&hidden_states, &attention_mask).unwrap();

        // 输出形状应该是 (batch, hidden_size)
        assert_eq!(pooled.dims(), &[2, 64]);
    }

    #[test]
    fn test_full_encode() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config, vb).unwrap();

        // 创建测试输入 - 使用随机非零 token IDs
        let input_ids_data: Vec<u32> = (0..32).map(|i| (i % 100 + 1) as u32).collect();
        let input_ids = Tensor::from_vec(input_ids_data, (2, 16), &device).unwrap();
        let attention_mask = Tensor::ones((2, 16), DType::F32, &device).unwrap();

        let embeddings = model.encode(&input_ids, &attention_mask, None).unwrap();

        // 输出形状应该是 (batch, hidden_size)
        assert_eq!(embeddings.dims(), &[2, 64]);

        // 检查输出是有效的（finite）
        let flat: Vec<f32> = embeddings.flatten_all().unwrap().to_vec1().unwrap();
        for val in &flat {
            assert!(val.is_finite(), "Embedding values should be finite");
        }

        // 检查 L2 归一化
        let norms: Vec<f32> = embeddings
            .sqr()
            .unwrap()
            .sum(1)
            .unwrap()
            .sqrt()
            .unwrap()
            .to_vec1()
            .unwrap();
        for norm in norms {
            assert!(norm.is_finite(), "Norm should be finite");
            // 如果 norm > 0，检查归一化
            if norm > 1e-6 {
                assert!((norm - 1.0).abs() < 0.1, "Norm should be close to 1.0, got {}", norm);
            }
        }
    }

    #[test]
    fn test_matryoshka_truncation() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config, vb).unwrap();

        // 创建测试输入 - 使用随机非零 token IDs
        let input_ids_data: Vec<u32> = (0..32).map(|i| (i % 100 + 1) as u32).collect();
        let input_ids = Tensor::from_vec(input_ids_data, (2, 16), &device).unwrap();
        let attention_mask = Tensor::ones((2, 16), DType::F32, &device).unwrap();

        // 测试 Matryoshka 截断
        let embeddings = model.encode(&input_ids, &attention_mask, Some(32)).unwrap();

        // 输出维度应该被截断为 32
        assert_eq!(embeddings.dims(), &[2, 32]);

        // 检查输出是有效的
        let flat: Vec<f32> = embeddings.flatten_all().unwrap().to_vec1().unwrap();
        for val in &flat {
            assert!(val.is_finite(), "Embedding values should be finite");
        }

        // 检查 L2 归一化
        let norms: Vec<f32> = embeddings
            .sqr()
            .unwrap()
            .sum(1)
            .unwrap()
            .sqrt()
            .unwrap()
            .to_vec1()
            .unwrap();
        for norm in norms {
            assert!(norm.is_finite(), "Norm should be finite");
            if norm > 1e-6 {
                assert!((norm - 1.0).abs() < 0.1, "Norm should be close to 1.0, got {}", norm);
            }
        }
    }

    #[test]
    fn test_l2_normalize() {
        let device = Device::Cpu;
        let vb = create_test_varbuilder(&device);
        let config = create_small_config();

        let model = EmbeddingGemmaModel::new(config, vb).unwrap();

        // 创建非归一化的向量
        let input = Tensor::new(&[[3.0f32, 4.0]], &device).unwrap();
        let normalized = model.l2_normalize(&input).unwrap();

        // 检查归一化后的范数为 1
        let values: Vec<f32> = normalized.to_vec2().unwrap()[0].clone();
        let norm: f32 = values.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-5);

        // 检查方向正确 (3/5, 4/5)
        assert!((values[0] - 0.6).abs() < 1e-5);
        assert!((values[1] - 0.8).abs() < 1e-5);
    }
}
