reasonkit-mem 0.1.7

//! Graph-Based Knowledge Representation using Petgraph
//!
//! This module provides graph algorithms and data structures for
//! knowledge graph reasoning, relationship traversal, and semantic
//! network operations in the memory layer.
//!
//! # Features
//! - Directed and undirected knowledge graphs
//! - Relationship traversal and path finding
//! - Semantic similarity through graph structure
//! - Integration with vector embeddings
//! - RAPTOR tree augmentation
//!
//! Enable with: `cargo build --features graph`

#![cfg(feature = "graph")]

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

// Re-export petgraph for advanced usage
#[cfg(feature = "graph")]
pub use graphalgs;
#[cfg(feature = "graph")]
pub use petgraph;

use petgraph::algo::{astar, dijkstra};
use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::visit::EdgeRef;
use petgraph::Direction;

/// A node in the knowledge graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeNode {
    /// Unique identifier
    pub id: String,
    /// Node label/type
    pub label: String,
    /// Node content/name
    pub content: String,
    /// Embedding vector (optional)
    pub embedding: Option<Vec<f32>>,
    /// Additional metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

/// An edge/relationship in the knowledge graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeEdge {
    /// Relationship type
    pub relation: String,
    /// Edge weight/strength
    pub weight: f32,
    /// Additional properties
    pub properties: HashMap<String, serde_json::Value>,
}

/// A knowledge graph for semantic reasoning
pub struct KnowledgeGraph {
    graph: DiGraph<KnowledgeNode, KnowledgeEdge>,
    node_index: HashMap<String, NodeIndex>,
}

impl KnowledgeGraph {
    /// Create a new empty knowledge graph
    pub fn new() -> Self {
        Self {
            graph: DiGraph::new(),
            node_index: HashMap::new(),
        }
    }

    /// Add a node to the graph
    pub fn add_node(&mut self, node: KnowledgeNode) -> NodeIndex {
        let id = node.id.clone();
        let idx = self.graph.add_node(node);
        self.node_index.insert(id, idx);
        idx
    }

    /// Add an edge between two nodes
    pub fn add_edge(&mut self, from_id: &str, to_id: &str, edge: KnowledgeEdge) -> Result<()> {
        let from_idx = self
            .node_index
            .get(from_id)
            .ok_or_else(|| anyhow::anyhow!("Source node not found: {}", from_id))?;
        let to_idx = self
            .node_index
            .get(to_id)
            .ok_or_else(|| anyhow::anyhow!("Target node not found: {}", to_id))?;

        self.graph.add_edge(*from_idx, *to_idx, edge);
        Ok(())
    }

    /// Get a node by ID
    pub fn get_node(&self, id: &str) -> Option<&KnowledgeNode> {
        self.node_index
            .get(id)
            .and_then(|idx| self.graph.node_weight(*idx))
    }

    /// Get all neighbors of a node
    pub fn neighbors(&self, id: &str, direction: Direction) -> Vec<&KnowledgeNode> {
        let idx = match self.node_index.get(id) {
            Some(idx) => *idx,
            None => return Vec::new(),
        };

        self.graph
            .neighbors_directed(idx, direction)
            .filter_map(|n| self.graph.node_weight(n))
            .collect()
    }

    /// Find shortest path between two nodes
    pub fn shortest_path(&self, from_id: &str, to_id: &str) -> Option<Vec<String>> {
        let from_idx = self.node_index.get(from_id)?;
        let to_idx = self.node_index.get(to_id)?;

        let path = astar(
            &self.graph,
            *from_idx,
            |finish| finish == *to_idx,
            |e| e.weight().weight as i32,
            |_| 0,
        )?;

        Some(
            path.1
                .iter()
                .filter_map(|idx| self.graph.node_weight(*idx))
                .map(|n| n.id.clone())
                .collect(),
        )
    }

    /// Find all nodes within a certain distance
    pub fn nodes_within_distance(
        &self,
        from_id: &str,
        max_distance: f32,
    ) -> Vec<(&KnowledgeNode, f32)> {
        let from_idx = match self.node_index.get(from_id) {
            Some(idx) => *idx,
            None => return Vec::new(),
        };

        let distances = dijkstra(&self.graph, from_idx, None, |e| e.weight().weight);

        distances
            .iter()
            .filter(|(_, &dist)| dist <= max_distance)
            .filter_map(|(idx, &dist)| self.graph.node_weight(*idx).map(|n| (n, dist)))
            .collect()
    }

    /// Get all edges between two nodes
    pub fn edges_between(&self, from_id: &str, to_id: &str) -> Vec<&KnowledgeEdge> {
        let from_idx = match self.node_index.get(from_id) {
            Some(idx) => *idx,
            None => return Vec::new(),
        };
        let to_idx = match self.node_index.get(to_id) {
            Some(idx) => *idx,
            None => return Vec::new(),
        };

        self.graph
            .edges_connecting(from_idx, to_idx)
            .map(|e| e.weight())
            .collect()
    }

    /// Get node count
    pub fn node_count(&self) -> usize {
        self.graph.node_count()
    }

    /// Get edge count
    pub fn edge_count(&self) -> usize {
        self.graph.edge_count()
    }

    /// Find nodes by label
    pub fn find_by_label(&self, label: &str) -> Vec<&KnowledgeNode> {
        self.graph
            .node_weights()
            .filter(|n| n.label == label)
            .collect()
    }

    /// Find related nodes through a specific relation type
    pub fn find_related(&self, id: &str, relation: &str) -> Vec<&KnowledgeNode> {
        let idx = match self.node_index.get(id) {
            Some(idx) => *idx,
            None => return Vec::new(),
        };

        self.graph
            .edges_directed(idx, Direction::Outgoing)
            .filter(|e| e.weight().relation == relation)
            .filter_map(|e| self.graph.node_weight(e.target()))
            .collect()
    }
}

impl Default for KnowledgeGraph {
    fn default() -> Self {
        Self::new()
    }
}

/// Reasoning chain built from graph traversal
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ReasoningPath {
    /// Nodes in the reasoning chain
    pub nodes: Vec<String>,
    /// Relationships traversed
    pub relations: Vec<String>,
    /// Total path weight
    pub total_weight: f32,
    /// Semantic confidence score
    pub confidence: f32,
}

/// Graph-based reasoner for knowledge exploration
pub struct GraphReasoner {
    graph: Arc<KnowledgeGraph>,
    max_depth: usize,
}

impl GraphReasoner {
    pub fn new(graph: Arc<KnowledgeGraph>, max_depth: usize) -> Self {
        Self { graph, max_depth }
    }

    /// Find reasoning paths between two concepts
    pub fn find_reasoning_paths(
        &self,
        from: &str,
        to: &str,
        max_paths: usize,
    ) -> Vec<ReasoningPath> {
        // Simplified implementation - would use more sophisticated algorithms
        let mut paths = Vec::new();

        if let Some(path) = self.graph.shortest_path(from, to) {
            paths.push(ReasoningPath {
                nodes: path,
                relations: Vec::new(), // Would be populated with edge relations
                total_weight: 1.0,
                confidence: 0.8,
            });
        }

        paths.truncate(max_paths);
        paths
    }

    /// Expand a concept through its relationships
    pub fn expand_concept(&self, concept: &str, depth: usize) -> ConceptExpansion {
        let neighbors = self.graph.neighbors(concept, Direction::Outgoing);

        ConceptExpansion {
            root: concept.to_string(),
            related: neighbors.iter().map(|n| n.id.clone()).collect(),
            depth_reached: depth.min(self.max_depth),
        }
    }
}

/// Result of concept expansion
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConceptExpansion {
    pub root: String,
    pub related: Vec<String>,
    pub depth_reached: usize,
}

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

    fn create_test_graph() -> KnowledgeGraph {
        let mut graph = KnowledgeGraph::new();

        graph.add_node(KnowledgeNode {
            id: "A".to_string(),
            label: "concept".to_string(),
            content: "Concept A".to_string(),
            embedding: None,
            metadata: HashMap::new(),
        });

        graph.add_node(KnowledgeNode {
            id: "B".to_string(),
            label: "concept".to_string(),
            content: "Concept B".to_string(),
            embedding: None,
            metadata: HashMap::new(),
        });

        graph
            .add_edge(
                "A",
                "B",
                KnowledgeEdge {
                    relation: "related_to".to_string(),
                    weight: 1.0,
                    properties: HashMap::new(),
                },
            )
            .unwrap();

        graph
    }

    #[test]
    fn test_add_and_get_node() {
        let graph = create_test_graph();

        let node = graph.get_node("A").unwrap();
        assert_eq!(node.content, "Concept A");
    }

    #[test]
    fn test_neighbors() {
        let graph = create_test_graph();

        let neighbors = graph.neighbors("A", Direction::Outgoing);
        assert_eq!(neighbors.len(), 1);
        assert_eq!(neighbors[0].id, "B");
    }

    #[test]
    fn test_node_count() {
        let graph = create_test_graph();
        assert_eq!(graph.node_count(), 2);
        assert_eq!(graph.edge_count(), 1);
    }
}