//! Graph-Aware Attention — 图增强 GPT (Phase-71)
//!
//! 将 SocialGraphV2 的关系信息注入 GPT 的 Attention 机制，
//! 使模型能够在推理和生成时感知图结构。
//!
//! ## 核心创新
//!
//! ```text
//! Standard GPT:   attn = softmax(Q·K^T / √d + causal_mask)
//! Graph-Aware:    attn = softmax(Q·K^T / √d + causal_mask + graph_bias + sifu_bias)
//!                                                            ↑             ↑
//!                                                 图连通性权重    SiFu 信号传播
//! ```
//!
//! ## 功能模块
//!
//! - **T71.1** `GraphAttentionMask`: 图连通性 → 注意力偏置矩阵
//! - **T71.2** `RelationTypeEmbedding`: 关系类型 → 额外 embedding 维度
//! - **T71.3** `KHopContextBuilder`: k 跳邻居 → 上下文 token 扩展
//! - **T71.4** `SiFuAttentionBias`: SiFu edge_weight → attention bias
//! - **T71.5** `GraphAwareGpt`: 集成以上所有组件的增强版 GPT

use std::collections::HashMap;

use candle_core::{DType, Device, Result as CandleResult, Tensor};
use candle_nn::{embedding, linear, Embedding, Linear, Module, VarBuilder};
use serde::{Deserialize, Serialize};

// ============================================================================
// T71.1 — Graph-Aware Attention Mask
// ============================================================================

/// 图节点关系对
#[derive(Debug, Clone)]
pub struct GraphEdge {
    /// 源 token 位置（在序列中的索引）
    pub src_pos: usize,
    /// 目标 token 位置
    pub tgt_pos: usize,
    /// 关系权重 (0.0 ~ 1.0)
    pub weight: f64,
    /// 关系类型（可选，用于 relation embedding）
    pub relation_type: Option<String>,
}

/// 图结构信息，用于构建注意力偏置
#[derive(Debug, Clone)]
pub struct GraphContext {
    /// 序列长度
    pub seq_len: usize,
    /// 图中的边（token 位置对 + 权重）
    pub edges: Vec<GraphEdge>,
    /// token 位置 → 实体 ID 映射
    pub token_entity_map: HashMap<usize, String>,
    /// k-hop 额外上下文 token（追加在序列后面）
    pub context_tokens: Vec<u32>,
    /// context token 位置 → 实体 ID 映射
    pub context_entity_map: HashMap<usize, String>,
}

impl GraphContext {
    pub fn new(seq_len: usize) -> Self {
        Self {
            seq_len,
            edges: Vec::new(),
            token_entity_map: HashMap::new(),
            context_tokens: Vec::new(),
            context_entity_map: HashMap::new(),
        }
    }

    /// 添加一条图边
    pub fn add_edge(&mut self, src_pos: usize, tgt_pos: usize, weight: f64, rel_type: Option<String>) {
        self.edges.push(GraphEdge {
            src_pos,
            tgt_pos,
            weight,
            relation_type: rel_type,
        });
    }

    /// 总序列长度（原始 + 上下文）
    pub fn total_len(&self) -> usize {
        self.seq_len + self.context_tokens.len()
    }
}

/// 图注意力偏置配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphAttentionConfig {
    /// 图连通性偏置强度 (默认 1.0)
    pub graph_bias_scale: f64,
    /// SiFu 信号传播偏置强度 (默认 0.5)
    pub sifu_bias_scale: f64,
    /// 关系类型数量上限
    pub max_relation_types: usize,
    /// 关系类型 embedding 维度
    pub relation_embd_dim: usize,
    /// k-hop 邻居搜索深度
    pub k_hop: usize,
    /// 每跳最大邻居数
    pub max_neighbors_per_hop: usize,
    /// 是否启用双向注意力（有图关系时双向可见）
    pub bidirectional_attention: bool,
}

impl Default for GraphAttentionConfig {
    fn default() -> Self {
        Self {
            graph_bias_scale: 1.0,
            sifu_bias_scale: 0.5,
            max_relation_types: 32,
            relation_embd_dim: 16,
            k_hop: 2,
            max_neighbors_per_hop: 10,
            bidirectional_attention: true,
        }
    }
}

/// 从 GraphContext 构建注意力偏置矩阵
///
/// 返回 (T, T) 的偏置张量，有图关系的位置对获得正值偏置，
/// 使模型在 softmax 时分配更多注意力给图相关的 token。
pub fn build_graph_attention_bias(
    graph_ctx: &GraphContext,
    config: &GraphAttentionConfig,
    device: &Device,
) -> CandleResult<Tensor> {
    let total_len = graph_ctx.total_len();
    // 初始化零偏置
    let mut bias_data = vec![0.0f32; total_len * total_len];

    for edge in &graph_ctx.edges {
        if edge.src_pos < total_len && edge.tgt_pos < total_len {
            let weight = (edge.weight * config.graph_bias_scale) as f32;
            // 从 src → tgt 的注意力增强
            bias_data[edge.src_pos * total_len + edge.tgt_pos] += weight;

            // 双向模式下也增强 tgt → src
            if config.bidirectional_attention {
                bias_data[edge.tgt_pos * total_len + edge.src_pos] += weight;
            }
        }
    }

    Tensor::from_vec(bias_data, (total_len, total_len), device)
}

// ============================================================================
// T71.2 — Relation Type Embedding
// ============================================================================

/// 关系类型编码器
///
/// 将关系类型字符串映射为 embedding 向量，
/// 注入到 attention score 或 token embedding 中。
pub struct RelationTypeEncoder {
    /// 关系类型 → 索引映射
    type_to_idx: HashMap<String, usize>,
    /// 下一个可用索引
    next_idx: usize,
    /// 最大类型数
    max_types: usize,
    /// Embedding 层
    embedding: Embedding,
    /// 投影到注意力维度
    proj: Linear,
}

impl RelationTypeEncoder {
    /// 创建关系类型编码器
    pub fn new(
        config: &GraphAttentionConfig,
        attn_dim: usize,
        vb: VarBuilder,
    ) -> CandleResult<Self> {
        let embd = embedding(
            config.max_relation_types + 1, // +1 for unknown
            config.relation_embd_dim,
            vb.pp("rel_type_embd"),
        )?;
        let proj = linear(config.relation_embd_dim, attn_dim, vb.pp("rel_type_proj"))?;

        Ok(Self {
            type_to_idx: HashMap::new(),
            next_idx: 0,
            max_types: config.max_relation_types,
            embedding: embd,
            proj,
        })
    }

    /// 注册一个关系类型，返回其索引
    pub fn register_type(&mut self, rel_type: &str) -> usize {
        if let Some(&idx) = self.type_to_idx.get(rel_type) {
            return idx;
        }
        if self.next_idx >= self.max_types {
            return self.max_types; // unknown type
        }
        let idx = self.next_idx;
        self.type_to_idx.insert(rel_type.to_string(), idx);
        self.next_idx += 1;
        idx
    }

    /// 获取关系类型的 embedding 向量 (投影到 attn_dim)
    pub fn encode(&self, rel_type: &str, device: &Device) -> CandleResult<Tensor> {
        let idx = self
            .type_to_idx
            .get(rel_type)
            .copied()
            .unwrap_or(self.max_types);
        let idx_tensor = Tensor::new(&[idx as u32], device)?;
        let embd = self.embedding.forward(&idx_tensor)?;
        self.proj.forward(&embd)
    }

    /// 批量编码关系类型，返回 (N, attn_dim) 张量
    pub fn encode_batch(
        &self,
        rel_types: &[&str],
        device: &Device,
    ) -> CandleResult<Tensor> {
        let indices: Vec<u32> = rel_types
            .iter()
            .map(|t| {
                self.type_to_idx
                    .get(*t)
                    .copied()
                    .unwrap_or(self.max_types) as u32
            })
            .collect();
        let idx_tensor = Tensor::new(indices.as_slice(), device)?;
        let embd = self.embedding.forward(&idx_tensor)?;
        self.proj.forward(&embd)
    }
}

// ============================================================================
// T71.3 — k-hop Context Builder
// ============================================================================

/// k-hop 邻居上下文构建器
///
/// 从图中提取目标实体的 k 跳邻居，将邻居的文本
/// 拼接为额外的上下文 token，扩展模型的可见窗口。
#[derive(Debug, Clone)]
pub struct KHopNeighbor {
    /// 实体 ID
    pub entity_id: String,
    /// 实体文本内容
    pub content: String,
    /// 跳数距离
    pub hop: usize,
    /// 关系类型路径
    pub relation_path: Vec<String>,
    /// 信号强度（从种子到此节点的衰减后强度）
    pub signal_strength: f64,
}

/// k-hop 上下文配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KHopConfig {
    /// 最大跳数
    pub max_hops: usize,
    /// 每跳最大邻居数
    pub max_per_hop: usize,
    /// 总最大上下文 token 数
    pub max_context_tokens: usize,
    /// 信号衰减因子 (每跳乘以此值)
    pub decay_factor: f64,
}

impl Default for KHopConfig {
    fn default() -> Self {
        Self {
            max_hops: 2,
            max_per_hop: 5,
            max_context_tokens: 256,
            decay_factor: 0.5,
        }
    }
}

/// 从邻居列表构建 GraphContext 的 context tokens
///
/// 将每个邻居的内容 tokenize 后追加到 context_tokens，
/// 并记录 token 位置 → 实体 ID 映射。
pub fn build_khop_context(
    neighbors: &[KHopNeighbor],
    tokenizer: &super::TokenizerManager,
    base_seq_len: usize,
    max_tokens: usize,
) -> Result<(Vec<u32>, HashMap<usize, String>), String> {
    let mut context_tokens = Vec::new();
    let mut entity_map = HashMap::new();
    let mut current_pos = base_seq_len;

    // 按信号强度降序排列
    let mut sorted: Vec<_> = neighbors.to_vec();
    sorted.sort_by(|a, b| {
        b.signal_strength
            .partial_cmp(&a.signal_strength)
            .unwrap_or(std::cmp::Ordering::Equal)
    });

    for neighbor in &sorted {
        if context_tokens.len() >= max_tokens {
            break;
        }

        let tokens = tokenizer.encode(&neighbor.content)?;
        let take = tokens.len().min(max_tokens - context_tokens.len());

        for &tok in &tokens[..take] {
            entity_map.insert(current_pos, neighbor.entity_id.clone());
            context_tokens.push(tok);
            current_pos += 1;
        }
    }

    Ok((context_tokens, entity_map))
}

// ============================================================================
// T71.4 — SiFu Attention Bias
// ============================================================================

/// SiFu 信号传播 → 注意力偏置
///
/// 从 SiFu 的 EdgeWeightStore 中提取边权重（signal_strength），
/// 转换为注意力偏置矩阵。使 GPT 在推理时"感知"图中信号的流动方向。
#[derive(Debug, Clone)]
pub struct SiFuEdge {
    /// 源 token 位置
    pub src_pos: usize,
    /// 目标 token 位置
    pub tgt_pos: usize,
    /// SiFu 信号强度
    pub signal_strength: f64,
    /// 共激活次数（赫布学习权重）
    pub co_activation_count: u32,
}

/// 从 SiFu 边权重构建注意力偏置
pub fn build_sifu_attention_bias(
    sifu_edges: &[SiFuEdge],
    total_len: usize,
    scale: f64,
    device: &Device,
) -> CandleResult<Tensor> {
    let mut bias_data = vec![0.0f32; total_len * total_len];

    for edge in sifu_edges {
        if edge.src_pos < total_len && edge.tgt_pos < total_len {
            // 综合信号强度和共激活权重
            let hebbian_boost = (edge.co_activation_count as f64).ln().max(0.0) * 0.1;
            let weight = ((edge.signal_strength + hebbian_boost) * scale) as f32;

            bias_data[edge.src_pos * total_len + edge.tgt_pos] += weight;
            // 反向也施加衰减的偏置
            bias_data[edge.tgt_pos * total_len + edge.src_pos] += weight * 0.5;
        }
    }

    Tensor::from_vec(bias_data, (total_len, total_len), device)
}

// ============================================================================
// T71.5 — Graph-Aware GPT (集成)
// ============================================================================

/// 图增强的注意力偏置（合并后的最终偏置矩阵）
pub struct GraphAwareBias {
    /// 合并后的 (1, 1, T, T) 偏置张量（可直接 broadcast_add 到 attention scores）
    pub combined_bias: Tensor,
}

impl GraphAwareBias {
    /// 构建图增强注意力偏置
    ///
    /// 合并：因果掩码 + 图连通性偏置 + SiFu 信号偏置
    pub fn build(
        graph_ctx: &GraphContext,
        sifu_edges: &[SiFuEdge],
        config: &GraphAttentionConfig,
        device: &Device,
    ) -> CandleResult<Self> {
        let total_len = graph_ctx.total_len();

        // 1. 因果掩码 (上三角 → -inf)
        let causal = super::causal_mask(total_len, device)?;

        // 2. 图连通性偏置
        let graph_bias = build_graph_attention_bias(graph_ctx, config, device)?;

        // 3. SiFu 信号偏置
        let sifu_bias =
            build_sifu_attention_bias(sifu_edges, total_len, config.sifu_bias_scale, device)?;

        // 合并: causal_mask + graph_bias + sifu_bias
        // causal_mask: (1, 1, T, T), graph_bias: (T, T), sifu_bias: (T, T)
        let combined = graph_bias.add(&sifu_bias)?;
        let combined = combined.unsqueeze(0)?.unsqueeze(0)?; // (1, 1, T, T)
        let combined = causal.broadcast_add(&combined)?;

        Ok(Self {
            combined_bias: combined,
        })
    }

    /// 获取偏置张量（可直接用于 attention 计算）
    pub fn as_tensor(&self) -> &Tensor {
        &self.combined_bias
    }
}

/// 图增强 GPT 推理配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphAwareGenerateConfig {
    /// 基础生成配置
    pub max_tokens: usize,
    pub temperature: f64,
    pub top_k: Option<usize>,
    /// 图注意力配置
    pub graph_config: GraphAttentionConfig,
    /// k-hop 配置
    pub khop_config: KHopConfig,
}

impl Default for GraphAwareGenerateConfig {
    fn default() -> Self {
        Self {
            max_tokens: 100,
            temperature: 0.8,
            top_k: Some(40),
            graph_config: GraphAttentionConfig::default(),
            khop_config: KHopConfig::default(),
        }
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    fn test_device() -> Device {
        Device::Cpu
    }

    #[test]
    fn test_graph_context_creation() {
        let mut ctx = GraphContext::new(10);
        assert_eq!(ctx.seq_len, 10);
        assert_eq!(ctx.total_len(), 10);

        ctx.add_edge(0, 3, 0.8, Some("friend".to_string()));
        ctx.add_edge(3, 7, 0.6, None);
        assert_eq!(ctx.edges.len(), 2);

        ctx.context_tokens = vec![1, 2, 3];
        assert_eq!(ctx.total_len(), 13);
    }

    #[test]
    fn test_build_graph_attention_bias() {
        let device = test_device();
        let mut ctx = GraphContext::new(4);
        ctx.add_edge(0, 2, 1.0, None);
        ctx.add_edge(1, 3, 0.5, None);

        let config = GraphAttentionConfig::default();
        let bias = build_graph_attention_bias(&ctx, &config, &device).unwrap();

        assert_eq!(bias.dims(), &[4, 4]);

        // Check values
        let data: Vec<Vec<f32>> = bias.to_vec2().unwrap();
        // Edge 0→2: should have positive bias
        assert!(data[0][2] > 0.0);
        // Bidirectional: 2→0 should also have positive bias
        assert!(data[2][0] > 0.0);
        // Non-connected: 0→1 should be 0
        assert_eq!(data[0][1], 0.0);
    }

    #[test]
    fn test_build_graph_bias_unidirectional() {
        let device = test_device();
        let mut ctx = GraphContext::new(4);
        ctx.add_edge(0, 2, 1.0, None);

        let mut config = GraphAttentionConfig::default();
        config.bidirectional_attention = false;

        let bias = build_graph_attention_bias(&ctx, &config, &device).unwrap();
        let data: Vec<Vec<f32>> = bias.to_vec2().unwrap();

        assert!(data[0][2] > 0.0);
        assert_eq!(data[2][0], 0.0); // No reverse in unidirectional
    }

    #[test]
    fn test_build_sifu_attention_bias() {
        let device = test_device();
        let edges = vec![
            SiFuEdge {
                src_pos: 0,
                tgt_pos: 2,
                signal_strength: 0.9,
                co_activation_count: 10,
            },
            SiFuEdge {
                src_pos: 1,
                tgt_pos: 3,
                signal_strength: 0.3,
                co_activation_count: 1,
            },
        ];

        let bias = build_sifu_attention_bias(&edges, 4, 0.5, &device).unwrap();
        assert_eq!(bias.dims(), &[4, 4]);

        let data: Vec<Vec<f32>> = bias.to_vec2().unwrap();
        // High signal + high co-activation → higher bias
        assert!(data[0][2] > data[1][3]);
        // Reverse should have attenuated bias
        assert!(data[2][0] > 0.0);
        assert!(data[2][0] < data[0][2]);
    }

    #[test]
    fn test_graph_aware_bias_build() {
        let device = test_device();
        let mut graph_ctx = GraphContext::new(4);
        graph_ctx.add_edge(0, 2, 0.8, None);

        let sifu_edges = vec![SiFuEdge {
            src_pos: 1,
            tgt_pos: 3,
            signal_strength: 0.7,
            co_activation_count: 5,
        }];

        let config = GraphAttentionConfig::default();
        let bias = GraphAwareBias::build(&graph_ctx, &sifu_edges, &config, &device).unwrap();

        // Should be (1, 1, 4, 4)
        assert_eq!(bias.as_tensor().dims(), &[1, 1, 4, 4]);

        // Extract values
        let data: Vec<f32> = bias.as_tensor().flatten_all().unwrap().to_vec1().unwrap();
        // Causal mask: positions above diagonal should be -inf (very large negative)
        // Position (0, 1) = row 0, col 1 → should be masked (future token)
        let val_0_1 = data[0 * 4 + 1];
        assert!(val_0_1 < -1000.0); // Masked position

        // Position (2, 0) → row 2, col 0 → should be visible + graph bias
        let val_2_0 = data[2 * 4 + 0];
        assert!(val_2_0 > -1000.0); // Not masked (past token)
    }

    #[test]
    fn test_build_khop_context() {
        let tokenizer = super::super::TokenizerManager::char_level();
        let neighbors = vec![
            KHopNeighbor {
                entity_id: "alice".to_string(),
                content: "Hello".to_string(),
                hop: 1,
                relation_path: vec!["friend".to_string()],
                signal_strength: 0.9,
            },
            KHopNeighbor {
                entity_id: "bob".to_string(),
                content: "World".to_string(),
                hop: 2,
                relation_path: vec!["friend".to_string(), "colleague".to_string()],
                signal_strength: 0.5,
            },
        ];

        let (tokens, entity_map) = build_khop_context(&neighbors, &tokenizer, 10, 100).unwrap();

        assert!(!tokens.is_empty());
        assert!(!entity_map.is_empty());

        // First neighbor (higher signal) should come first
        let first_entity = entity_map.get(&10).unwrap();
        assert_eq!(first_entity, "alice");
    }

    #[test]
    fn test_relation_type_encoder() {
        let device = test_device();
        let config = GraphAttentionConfig::default();
        let varmap = candle_nn::VarMap::new();
        let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);

        let mut encoder = RelationTypeEncoder::new(&config, 64, vb).unwrap();

        // Register types
        let idx_friend = encoder.register_type("friend");
        let idx_follow = encoder.register_type("follow");
        assert_ne!(idx_friend, idx_follow);

        // Re-register same type → same index
        assert_eq!(encoder.register_type("friend"), idx_friend);

        // Encode
        let embd = encoder.encode("friend", &device).unwrap();
        assert_eq!(embd.dims(), &[1, 64]);

        // Batch encode
        let batch = encoder.encode_batch(&["friend", "follow", "unknown_type"], &device).unwrap();
        assert_eq!(batch.dims(), &[3, 64]);
    }

    #[test]
    fn test_graph_attention_config_default() {
        let config = GraphAttentionConfig::default();
        assert_eq!(config.graph_bias_scale, 1.0);
        assert_eq!(config.sifu_bias_scale, 0.5);
        assert_eq!(config.k_hop, 2);
        assert!(config.bidirectional_attention);
    }

    #[test]
    fn test_khop_config_default() {
        let config = KHopConfig::default();
        assert_eq!(config.max_hops, 2);
        assert_eq!(config.max_per_hop, 5);
        assert_eq!(config.decay_factor, 0.5);
    }

    #[test]
    fn test_graph_aware_generate_config() {
        let config = GraphAwareGenerateConfig::default();
        assert_eq!(config.max_tokens, 100);
        assert_eq!(config.temperature, 0.8);
        assert_eq!(config.graph_config.k_hop, 2);
    }

    #[test]
    fn test_empty_graph_bias() {
        let device = test_device();
        let ctx = GraphContext::new(4);
        let config = GraphAttentionConfig::default();
        let bias = build_graph_attention_bias(&ctx, &config, &device).unwrap();

        // All zeros (no edges)
        let data: Vec<f32> = bias.flatten_all().unwrap().to_vec1().unwrap();
        assert!(data.iter().all(|&v| v == 0.0));
    }

    #[test]
    fn test_out_of_bounds_edge() {
        let device = test_device();
        let mut ctx = GraphContext::new(4);
        ctx.add_edge(0, 10, 1.0, None); // out of bounds

        let config = GraphAttentionConfig::default();
        let bias = build_graph_attention_bias(&ctx, &config, &device).unwrap();

        // Should not panic, out-of-bounds edge is skipped
        let data: Vec<f32> = bias.flatten_all().unwrap().to_vec1().unwrap();
        assert!(data.iter().all(|&v| v == 0.0));
    }
}
