//! Transformer编码器
//!
//! 本模块实现4层8头的Transformer编码器，用于提取50维宏观特征的时序依赖。

use crate::constants::*;
use crate::error::{Error, Result};
use ndarray::{Array1, Array2};
use serde::{Deserialize, Serialize};

// ============================================================================
// 多头注意力层
// ============================================================================

/// 多头注意力层
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MultiHeadAttention {
    /// 注意力头数
    num_heads: usize,
    /// 嵌入维度
    embed_dim: usize,
    /// 每个头的维度
    head_dim: usize,
    /// Query权重矩阵（embed_dim × embed_dim）
    w_q: Array2<f32>,
    /// Key权重矩阵（embed_dim × embed_dim）
    w_k: Array2<f32>,
    /// Value权重矩阵（embed_dim × embed_dim）
    w_v: Array2<f32>,
    /// 输出权重矩阵（embed_dim × embed_dim）
    w_o: Array2<f32>,
}

impl MultiHeadAttention {
    /// 创建新的多头注意力层
    ///
    /// # 参数
    ///
    /// * `num_heads` - 注意力头数
    /// * `embed_dim` - 嵌入维度
    pub fn new(num_heads: usize, embed_dim: usize) -> Result<Self> {
        if embed_dim % num_heads != 0 {
            return Err(Error::ArchitectureError(format!(
                "嵌入维度 {} 必须能被注意力头数 {} 整除",
                embed_dim, num_heads
            )));
        }

        let head_dim = embed_dim / num_heads;

        // 初始化权重矩阵（Xavier初始化）
        let scale = (embed_dim as f32).sqrt();
        let w_q = Array2::zeros((embed_dim, embed_dim)) / scale;
        let w_k = Array2::zeros((embed_dim, embed_dim)) / scale;
        let w_v = Array2::zeros((embed_dim, embed_dim)) / scale;
        let w_o = Array2::zeros((embed_dim, embed_dim)) / scale;

        Ok(Self {
            num_heads,
            embed_dim,
            head_dim,
            w_q,
            w_k,
            w_v,
            w_o,
        })
    }

    /// 前向传播
    ///
    /// # 参数
    ///
    /// * `x` - 输入张量（seq_len × embed_dim）
    ///
    /// # 返回
    ///
    /// 输出张量（seq_len × embed_dim）
    pub fn forward(&self, x: &Array2<f32>) -> Result<Array2<f32>> {
        let (seq_len, embed_dim) = x.dim();

        if embed_dim != self.embed_dim {
            return Err(Error::DimensionMismatch {
                expected: self.embed_dim,
                actual: embed_dim,
            });
        }

        // Q = X @ W_q, K = X @ W_k, V = X @ W_v
        let q = x.dot(&self.w_q);
        let k = x.dot(&self.w_k);
        let v = x.dot(&self.w_v);

        // 分割为多头：(seq_len, embed_dim) -> (num_heads, seq_len, head_dim)
        // 简化实现：直接在embed_dim维度上计算注意力

        // 注意力分数：Q @ K^T / sqrt(head_dim)
        let scale = (self.head_dim as f32).sqrt();
        let scores = q.dot(&k.t()) / scale;

        // Softmax（简化实现：沿最后一维）
        let attention = Self::softmax(&scores);

        // 加权求和：Attention @ V
        let output = attention.dot(&v);

        // 输出投影：Output @ W_o
        let result = output.dot(&self.w_o);

        Ok(result)
    }

    /// Softmax激活函数（简化实现）
    fn softmax(x: &Array2<f32>) -> Array2<f32> {
        let mut result = x.clone();
        for mut row in result.rows_mut() {
            let max = row.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
            for val in row.iter_mut() {
                *val = (*val - max).exp();
            }
            let sum: f32 = row.iter().sum();
            for val in row.iter_mut() {
                *val /= sum;
            }
        }
        result
    }
}

// ============================================================================
// 前馈神经网络
// ============================================================================

/// 前馈神经网络（FFN）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeedForward {
    /// 输入维度
    input_dim: usize,
    /// 隐藏层维度（通常为4倍输入维度）
    hidden_dim: usize,
    /// 第一层权重（input_dim × hidden_dim）
    w1: Array2<f32>,
    /// 第一层偏置
    b1: Array1<f32>,
    /// 第二层权重（hidden_dim × input_dim）
    w2: Array2<f32>,
    /// 第二层偏置
    b2: Array1<f32>,
}

impl FeedForward {
    /// 创建新的前馈神经网络
    ///
    /// # 参数
    ///
    /// * `input_dim` - 输入维度
    /// * `hidden_dim` - 隐藏层维度
    pub fn new(input_dim: usize, hidden_dim: usize) -> Self {
        let scale = (input_dim as f32).sqrt();
        let w1 = Array2::zeros((input_dim, hidden_dim)) / scale;
        let b1 = Array1::zeros(hidden_dim);
        let w2 = Array2::zeros((hidden_dim, input_dim)) / scale;
        let b2 = Array1::zeros(input_dim);

        Self {
            input_dim,
            hidden_dim,
            w1,
            b1,
            w2,
            b2,
        }
    }

    /// 前向传播
    ///
    /// # 参数
    ///
    /// * `x` - 输入张量（seq_len × input_dim）
    ///
    /// # 返回
    ///
    /// 输出张量（seq_len × input_dim）
    pub fn forward(&self, x: &Array2<f32>) -> Result<Array2<f32>> {
        let (seq_len, input_dim) = x.dim();

        if input_dim != self.input_dim {
            return Err(Error::DimensionMismatch {
                expected: self.input_dim,
                actual: input_dim,
            });
        }

        // 第一层：X @ W1 + b1
        let mut hidden = x.dot(&self.w1);
        for mut row in hidden.rows_mut() {
            for (val, &bias) in row.iter_mut().zip(self.b1.iter()) {
                *val += bias;
            }
        }

        // GELU激活函数
        Self::gelu_inplace(&mut hidden);

        // 第二层：Hidden @ W2 + b2
        let mut output = hidden.dot(&self.w2);
        for mut row in output.rows_mut() {
            for (val, &bias) in row.iter_mut().zip(self.b2.iter()) {
                *val += bias;
            }
        }

        Ok(output)
    }

    /// GELU激活函数（原地操作）
    fn gelu_inplace(x: &mut Array2<f32>) {
        for val in x.iter_mut() {
            *val = Self::gelu(*val);
        }
    }

    /// GELU激活函数
    fn gelu(x: f32) -> f32 {
        0.5 * x * (1.0 + (0.7978845608 * (x + 0.044715 * x.powi(3))).tanh())
    }
}

// ============================================================================
// Layer Normalization
// ============================================================================

/// Layer Normalization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LayerNorm {
    /// 归一化维度
    dim: usize,
    /// 缩放参数（gamma）
    gamma: Array1<f32>,
    /// 偏移参数（beta）
    beta: Array1<f32>,
    /// 数值稳定性常数
    eps: f32,
}

impl LayerNorm {
    /// 创建新的Layer Normalization
    ///
    /// # 参数
    ///
    /// * `dim` - 归一化维度
    pub fn new(dim: usize) -> Self {
        Self {
            dim,
            gamma: Array1::ones(dim),
            beta: Array1::zeros(dim),
            eps: 1e-5,
        }
    }

    /// 前向传播
    ///
    /// # 参数
    ///
    /// * `x` - 输入张量（seq_len × dim）
    ///
    /// # 返回
    ///
    /// 归一化后的张量
    pub fn forward(&self, x: &Array2<f32>) -> Result<Array2<f32>> {
        let (seq_len, dim) = x.dim();

        if dim != self.dim {
            return Err(Error::DimensionMismatch {
                expected: self.dim,
                actual: dim,
            });
        }

        let mut result = x.clone();

        for mut row in result.rows_mut() {
            // 计算均值和方差
            let mean = row.mean().unwrap_or(0.0);
            let variance = row.iter().map(|&v| (v - mean).powi(2)).sum::<f32>() / dim as f32;
            let std = (variance + self.eps).sqrt();

            // 归一化：(x - mean) / std
            for val in row.iter_mut() {
                *val = (*val - mean) / std;
            }

            // 缩放和偏移：gamma * x + beta
            for (val, (&gamma, &beta)) in row.iter_mut().zip(self.gamma.iter().zip(self.beta.iter())) {
                *val = gamma * *val + beta;
            }
        }

        Ok(result)
    }
}

// ============================================================================
// Transformer编码器层
// ============================================================================

/// Transformer编码器层
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransformerEncoderLayer {
    /// 多头注意力层
    attention: MultiHeadAttention,
    /// 前馈神经网络
    ffn: FeedForward,
    /// 第一个Layer Norm
    norm1: LayerNorm,
    /// 第二个Layer Norm
    norm2: LayerNorm,
}

impl TransformerEncoderLayer {
    /// 创建新的Transformer编码器层
    ///
    /// # 参数
    ///
    /// * `num_heads` - 注意力头数
    /// * `embed_dim` - 嵌入维度
    /// * `ffn_dim` - 前馈网络隐藏层维度
    pub fn new(num_heads: usize, embed_dim: usize, ffn_dim: usize) -> Result<Self> {
        Ok(Self {
            attention: MultiHeadAttention::new(num_heads, embed_dim)?,
            ffn: FeedForward::new(embed_dim, ffn_dim),
            norm1: LayerNorm::new(embed_dim),
            norm2: LayerNorm::new(embed_dim),
        })
    }

    /// 前向传播
    ///
    /// # 参数
    ///
    /// * `x` - 输入张量（seq_len × embed_dim）
    ///
    /// # 返回
    ///
    /// 输出张量（seq_len × embed_dim）
    pub fn forward(&self, x: &Array2<f32>) -> Result<Array2<f32>> {
        // 多头注意力 + 残差连接 + Layer Norm
        let attn_output = self.attention.forward(x)?;
        let x = &(x + &attn_output);
        let x = self.norm1.forward(x)?;

        // 前馈网络 + 残差连接 + Layer Norm
        let ffn_output = self.ffn.forward(&x)?;
        let x = &x + &ffn_output;
        let x = self.norm2.forward(&x)?;

        Ok(x)
    }
}

// ============================================================================
// Transformer编码器
// ============================================================================

/// Transformer编码器（4层）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TransformerEncoder {
    /// 编码器层列表
    layers: Vec<TransformerEncoderLayer>,
    /// 嵌入维度
    embed_dim: usize,
}

impl TransformerEncoder {
    /// 创建新的Transformer编码器
    ///
    /// # 参数
    ///
    /// * `num_layers` - 编码器层数
    /// * `num_heads` - 注意力头数
    /// * `embed_dim` - 嵌入维度
    /// * `ffn_dim` - 前馈网络隐藏层维度
    pub fn new(
        num_layers: usize,
        num_heads: usize,
        embed_dim: usize,
        ffn_dim: usize,
    ) -> Result<Self> {
        let mut layers = Vec::with_capacity(num_layers);
        for _ in 0..num_layers {
            layers.push(TransformerEncoderLayer::new(num_heads, embed_dim, ffn_dim)?);
        }

        Ok(Self { layers, embed_dim })
    }

    /// 前向传播
    ///
    /// # 参数
    ///
    /// * `x` - 输入张量（seq_len × embed_dim）
    ///
    /// # 返回
    ///
    /// 输出张量（seq_len × embed_dim）
    pub fn forward(&self, x: &Array2<f32>) -> Result<Array2<f32>> {
        let mut output = x.clone();
        for layer in &self.layers {
            output = layer.forward(&output)?;
        }
        Ok(output)
    }

    /// 获取嵌入维度
    pub fn embed_dim(&self) -> usize {
        self.embed_dim
    }
}

// ============================================================================
// 测试
// ============================================================================

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

    #[test]
    fn test_multi_head_attention() {
        let mha = MultiHeadAttention::new(ATTENTION_HEADS, EMBED_DIM).unwrap();
        let x = Array2::zeros((10, EMBED_DIM)); // seq_len=10
        let output = mha.forward(&x).unwrap();
        assert_eq!(output.dim(), (10, EMBED_DIM));
    }

    #[test]
    fn test_feed_forward() {
        let ffn = FeedForward::new(EMBED_DIM, EMBED_DIM * 4);
        let x = Array2::zeros((10, EMBED_DIM));
        let output = ffn.forward(&x).unwrap();
        assert_eq!(output.dim(), (10, EMBED_DIM));
    }

    #[test]
    fn test_layer_norm() {
        let ln = LayerNorm::new(EMBED_DIM);
        let x = Array2::ones((10, EMBED_DIM));
        let output = ln.forward(&x).unwrap();
        assert_eq!(output.dim(), (10, EMBED_DIM));
    }

    #[test]
    fn test_transformer_encoder_layer() {
        let layer = TransformerEncoderLayer::new(ATTENTION_HEADS, EMBED_DIM, EMBED_DIM * 4).unwrap();
        let x = Array2::zeros((10, EMBED_DIM));
        let output = layer.forward(&x).unwrap();
        assert_eq!(output.dim(), (10, EMBED_DIM));
    }

    #[test]
    fn test_transformer_encoder() {
        let encoder = TransformerEncoder::new(
            TRANSFORMER_LAYERS,
            ATTENTION_HEADS,
            EMBED_DIM,
            EMBED_DIM * 4,
        )
        .unwrap();
        let x = Array2::zeros((10, EMBED_DIM));
        let output = encoder.forward(&x).unwrap();
        assert_eq!(output.dim(), (10, EMBED_DIM));
    }

    #[test]
    fn test_dimension_mismatch() {
        let mha = MultiHeadAttention::new(ATTENTION_HEADS, EMBED_DIM).unwrap();
        let x = Array2::zeros((10, 128)); // 错误的维度
        assert!(mha.forward(&x).is_err());
    }
}