//! EmbeddingGemma 基础层实现
//!
//! 包含 RMSNorm 和 GeGLU Feed-Forward Network 的实现。
//!
//! # 组件
//!
//! - **RMSNorm**: Root Mean Square Layer Normalization
//!   - 比 LayerNorm 更高效，不需要计算均值
//!   - 公式: `output = x / RMS(x) * weight`
//!
//! - **GeGLU FFN**: Gated Linear Unit with GELU activation
//!   - 门控机制提升表达能力
//!   - 公式: `output = down_proj(gelu(gate_proj(x)) * up_proj(x))`

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

/// Root Mean Square Layer Normalization
///
/// RMSNorm 是一种简化的归一化方法，只使用均方根进行归一化，
/// 不减去均值。这使得它比 LayerNorm 更高效。
///
/// # 数学公式
///
/// ```text
/// RMS(x) = sqrt(mean(x^2) + eps)
/// output = x / RMS(x) * weight
/// ```
pub struct RMSNorm {
    /// 可学习的缩放权重
    weight: Tensor,
    /// 数值稳定性 epsilon
    eps: f64,
    /// 隐藏层维度
    hidden_size: usize,
}

impl RMSNorm {
    /// 创建 RMSNorm 层
    ///
    /// # Arguments
    /// * `hidden_size` - 隐藏层维度
    /// * `eps` - 数值稳定性 epsilon
    /// * `vb` - 变量构建器
    pub fn new(hidden_size: usize, eps: f64, vb: VarBuilder) -> Result<Self> {
        let weight = vb.get(hidden_size, "weight")?;
        Ok(Self {
            weight,
            eps,
            hidden_size,
        })
    }

    /// 使用指定权重创建 RMSNorm
    pub fn from_weight(weight: Tensor, eps: f64) -> Result<Self> {
        let hidden_size = weight.dim(0)?;
        Ok(Self {
            weight,
            eps,
            hidden_size,
        })
    }

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

impl Module for RMSNorm {
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        let original_dtype = x.dtype();

        // 转换为 F32 以获得更好的数值精度
        let x = x.to_dtype(DType::F32)?;

        // 计算均方值: mean(x^2)
        let variance = x.sqr()?.mean_keepdim(D::Minus1)?;

        // 计算 RMS: sqrt(variance + eps)
        let rms = (variance + self.eps)?.sqrt()?;

        // 归一化: x / rms
        let normalized = x.broadcast_div(&rms)?;

        // 应用可学习的缩放权重
        let weight = self.weight.to_dtype(DType::F32)?;
        let output = normalized.broadcast_mul(&weight)?;

        // 转换回原始数据类型
        output.to_dtype(original_dtype)
    }
}

/// GeGLU Feed-Forward Network
///
/// Gated Linear Unit 使用门控机制来控制信息流，
/// 结合 GELU 激活函数提升模型表达能力。
///
/// # 结构
///
/// ```text
/// input (hidden_size)
///     ├── gate_proj (hidden_size -> intermediate_size) -> GELU
///     └── up_proj   (hidden_size -> intermediate_size)
///              ↓
///         gate * up (element-wise multiply)
///              ↓
///         down_proj (intermediate_size -> hidden_size)
///              ↓
/// output (hidden_size)
/// ```
pub struct GeGLUFeedForward {
    /// 门控投影层
    gate_proj: candle_nn::Linear,
    /// 上投影层
    up_proj: candle_nn::Linear,
    /// 下投影层
    down_proj: candle_nn::Linear,
    /// 隐藏层维度
    hidden_size: usize,
    /// 中间层维度
    intermediate_size: usize,
}

impl GeGLUFeedForward {
    /// 创建 GeGLU FFN 层
    ///
    /// # Arguments
    /// * `hidden_size` - 输入/输出维度
    /// * `intermediate_size` - 中间层维度
    /// * `vb` - 变量构建器
    pub fn new(hidden_size: usize, intermediate_size: usize, vb: VarBuilder) -> Result<Self> {
        // Gemma 使用无偏置的线性层
        let gate_proj =
            candle_nn::linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
        let up_proj =
            candle_nn::linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?;
        let down_proj =
            candle_nn::linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?;

        Ok(Self {
            gate_proj,
            up_proj,
            down_proj,
            hidden_size,
            intermediate_size,
        })
    }

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

    /// 获取中间层维度
    pub fn intermediate_size(&self) -> usize {
        self.intermediate_size
    }
}

impl Module for GeGLUFeedForward {
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // 门控分支: GELU(gate_proj(x))
        let gate = self.gate_proj.forward(x)?;
        let gate = gelu(&gate)?;

        // 值分支: up_proj(x)
        let up = self.up_proj.forward(x)?;

        // 门控: gate * up
        let hidden = (gate * up)?;

        // 下投影: down_proj(hidden)
        self.down_proj.forward(&hidden)
    }
}

/// GELU 激活函数 (Gaussian Error Linear Unit)
///
/// 使用近似公式: GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/π) * (x + 0.044715 * x^3)))
pub fn gelu(x: &Tensor) -> Result<Tensor> {
    // 使用 Candle 内置的 GELU 激活
    x.gelu_erf()
}

/// SiLU/Swish 激活函数 (备用)
///
/// SiLU(x) = x * sigmoid(x)
#[allow(dead_code)]
pub fn silu(x: &Tensor) -> Result<Tensor> {
    x.silu()
}

/// 投影层 (用于最终的 embedding 投影)
///
/// EmbeddingGemma 使用两层投影将隐藏状态转换为最终的 embedding：
/// hidden_size -> projection_dim -> hidden_size
pub struct EmbeddingProjection {
    up_proj: candle_nn::Linear,
    down_proj: candle_nn::Linear,
}

impl EmbeddingProjection {
    /// 创建 embedding 投影层
    ///
    /// # Arguments
    /// * `hidden_size` - 输入/输出维度 (768)
    /// * `projection_dim` - 中间投影维度 (3072)
    /// * `vb` - 变量构建器
    pub fn new(hidden_size: usize, projection_dim: usize, vb: VarBuilder) -> Result<Self> {
        let up_proj = candle_nn::linear_no_bias(hidden_size, projection_dim, vb.pp("up_proj"))?;
        let down_proj = candle_nn::linear_no_bias(projection_dim, hidden_size, vb.pp("down_proj"))?;

        Ok(Self { up_proj, down_proj })
    }
}

impl Module for EmbeddingProjection {
    fn forward(&self, x: &Tensor) -> Result<Tensor> {
        // hidden_size -> projection_dim -> hidden_size
        let hidden = self.up_proj.forward(x)?;
        self.down_proj.forward(&hidden)
    }
}

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

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

    #[test]
    fn test_rmsnorm_shape() {
        let device = Device::Cpu;

        // 创建测试权重
        let weight = Tensor::ones((768,), DType::F32, &device).unwrap();
        let rmsnorm = RMSNorm::from_weight(weight, 1e-6).unwrap();

        // 测试输入 (batch=2, seq=16, hidden=768)
        let x = Tensor::randn(0.0f32, 1.0, (2, 16, 768), &device).unwrap();
        let output = rmsnorm.forward(&x).unwrap();

        assert_eq!(output.dims(), x.dims());
    }

    #[test]
    fn test_rmsnorm_normalization() {
        let device = Device::Cpu;

        // 使用单位权重，检查归一化效果
        let weight = Tensor::ones((4,), DType::F32, &device).unwrap();
        let rmsnorm = RMSNorm::from_weight(weight, 1e-6).unwrap();

        // 简单输入
        let x = Tensor::new(&[[1.0f32, 2.0, 3.0, 4.0]], &device).unwrap();
        let output = rmsnorm.forward(&x).unwrap();

        // 检查输出的 RMS 应该接近 1
        let rms: f32 = output
            .sqr()
            .unwrap()
            .mean_all()
            .unwrap()
            .to_scalar()
            .unwrap();
        let rms = rms.sqrt();
        assert!((rms - 1.0).abs() < 0.1);
    }

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

        let geglu = GeGLUFeedForward::new(768, 3072, vb).unwrap();

        // 测试输入 (batch=2, seq=16, hidden=768)
        let x = Tensor::randn(0.0f32, 1.0, (2, 16, 768), &device).unwrap();
        let output = geglu.forward(&x).unwrap();

        // 输出维度应该与输入相同
        assert_eq!(output.dims(), x.dims());
    }

    #[test]
    fn test_gelu_activation() {
        let device = Device::Cpu;

        // GELU(0) ≈ 0
        let x = Tensor::new(&[0.0f32], &device).unwrap();
        let y = gelu(&x).unwrap();
        let val: f32 = y.to_vec1::<f32>().unwrap()[0];
        assert!(val.abs() < 1e-5);

        // GELU(x) > 0 for x > 0
        let x = Tensor::new(&[1.0f32], &device).unwrap();
        let y = gelu(&x).unwrap();
        let val: f32 = y.to_vec1::<f32>().unwrap()[0];
        assert!(val > 0.0);

        // GELU(x) < 0 for some negative x (but close to 0)
        let x = Tensor::new(&[-0.5f32], &device).unwrap();
        let y = gelu(&x).unwrap();
        let val: f32 = y.to_vec1::<f32>().unwrap()[0];
        assert!(val < 0.0);
    }

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

        let proj = EmbeddingProjection::new(768, 3072, vb).unwrap();

        // 测试输入 (batch=2, hidden=768)
        let x = Tensor::randn(0.0f32, 1.0, (2, 768), &device).unwrap();
        let output = proj.forward(&x).unwrap();

        // 输出维度应该与输入相同
        assert_eq!(output.dims(), x.dims());
    }

    #[test]
    fn test_geglu_dimensions() {
        let geglu_ff = GeGLUFeedForward {
            gate_proj: candle_nn::Linear::new(
                Tensor::zeros((3072, 768), DType::F32, &Device::Cpu).unwrap(),
                None,
            ),
            up_proj: candle_nn::Linear::new(
                Tensor::zeros((3072, 768), DType::F32, &Device::Cpu).unwrap(),
                None,
            ),
            down_proj: candle_nn::Linear::new(
                Tensor::zeros((768, 3072), DType::F32, &Device::Cpu).unwrap(),
                None,
            ),
            hidden_size: 768,
            intermediate_size: 3072,
        };

        assert_eq!(geglu_ff.hidden_size(), 768);
        assert_eq!(geglu_ff.intermediate_size(), 3072);
    }
}
