//! EmbeddingGemma 模块
//!
//! 实现了 Google EmbeddingGemma-300M 模型，这是一个高效的文本 embedding 模型，
//! 专为设备端 AI 应用设计。
//!
//! # 特性
//!
//! - **高质量**: 在 500M 参数以下模型中 MTEB 排名最高
//! - **多语言**: 支持 100+ 语言
//! - **灵活维度**: 支持 Matryoshka 表示学习 (768/512/256/128)
//! - **长上下文**: 最大 2048 tokens
//!
//! # 架构
//!
//! EmbeddingGemma 基于 Gemma3 架构，但修改为双向注意力用于 embedding 任务：
//!
//! - 24 层 Transformer Encoder
//! - RoPE 旋转位置编码
//! - RMSNorm 归一化
//! - GeGLU 激活函数
//! - Mean Pooling + Dense 投影
//!
//! # 使用示例
//!
//! ```ignore
//! use aifastdb_core::perception::embedding_gemma::{
//!     EmbeddingGemmaConfig,
//!     EmbeddingGemmaModel,
//!     TaskType,
//!     MatryoshkaDim,
//! };
//!
//! // 加载模型
//! let config = EmbeddingGemmaConfig::default();
//! let model = EmbeddingGemmaModel::new(config, vb)?;
//!
//! // 生成 embedding
//! let embeddings = model.encode(&input_ids, &attention_mask, None)?;
//!
//! // 使用 Matryoshka 截断
//! let embeddings_256 = model.encode(&input_ids, &attention_mask, Some(256))?;
//! ```

// 子模块
mod attention;
mod config;
mod engine;
mod layers;
mod model;
mod rope;

// 测试模块
#[cfg(test)]
mod tests;

// 公开导出
pub use attention::{create_attention_mask, create_bidirectional_mask, BidirectionalAttention};
pub use config::{EmbeddingGemmaConfig, MatryoshkaDim, TaskType};
pub use engine::EmbeddingGemmaEngine;
pub use layers::{gelu, EmbeddingProjection, GeGLUFeedForward, RMSNorm};
pub use model::{EmbeddingGemmaModel, LoadOptions};
pub use rope::RotaryEmbedding;

/// 模块版本
pub const VERSION: &str = "0.1.0";

/// 默认模型 ID
pub const DEFAULT_MODEL_ID: &str = "google/embeddinggemma-300m";

/// 模型信息
#[derive(Debug, Clone)]
pub struct ModelInfo {
    /// 模型 ID
    pub model_id: String,
    /// 隐藏层维度
    pub hidden_size: usize,
    /// 最大序列长度
    pub max_seq_len: usize,
    /// 支持的 Matryoshka 维度
    pub matryoshka_dims: Vec<usize>,
    /// 支持的语言数量
    pub num_languages: usize,
}

impl Default for ModelInfo {
    fn default() -> Self {
        Self {
            model_id: DEFAULT_MODEL_ID.to_string(),
            hidden_size: 768,
            max_seq_len: 2048,
            matryoshka_dims: vec![768, 512, 256, 128],
            num_languages: 100,
        }
    }
}

impl ModelInfo {
    /// 获取 EmbeddingGemma-300M 的模型信息
    pub fn embedding_gemma_300m() -> Self {
        Self::default()
    }

    /// 检查维度是否是有效的 Matryoshka 维度
    pub fn is_valid_matryoshka_dim(&self, dim: usize) -> bool {
        self.matryoshka_dims.contains(&dim)
    }
}

