holographic-memory 0.6.4

// Copyright 2024-2026 WritersLogic Contributors
// SPDX-License-Identifier: Apache-2.0

use crate::core::hash::FastMap;
use crate::core::sync::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::VecDeque;

use super::types::{GraphPath, PathHop, Relation, RelationType};

/// In-memory graph index over explicit relations.
pub struct RelationStore {
    /// All relations, indexed by position.
    relations: RwLock<Vec<StoredRelation>>,
    /// source_id -> indices into relations vec.
    by_source: RwLock<FastMap<String, Vec<usize>>>,
    /// target_id -> indices into relations vec.
    by_target: RwLock<FastMap<String, Vec<usize>>>,
    /// relation_type -> indices into relations vec.
    by_type: RwLock<FastMap<String, Vec<usize>>>,
    /// Declared relation types with inference properties.
    types: RwLock<FastMap<String, RelationType>>,
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct StoredRelation {
    source_id: String,
    relation_type: String,
    target_id: String,
    properties: Option<String>,
    valid_from: f64,
    valid_to: f64,
    deleted: bool,
}

impl StoredRelation {
    fn is_valid_at(&self, timestamp_ms: f64) -> bool {
        if self.deleted {
            return false;
        }
        if timestamp_ms == 0.0 {
            return self.valid_to == 0.0;
        }
        let from_ok = self.valid_from == 0.0 || timestamp_ms >= self.valid_from;
        let to_ok = self.valid_to == 0.0 || timestamp_ms <= self.valid_to;
        from_ok && to_ok
    }

    fn to_relation(&self) -> Relation {
        Relation {
            source_id: self.source_id.clone(),
            relation_type: self.relation_type.clone(),
            target_id: self.target_id.clone(),
            properties: self.properties.clone(),
            valid_from: self.valid_from,
            valid_to: self.valid_to,
        }
    }
}

impl RelationStore {
    pub fn new() -> Self {
        Self {
            relations: RwLock::new(Vec::new()),
            by_source: RwLock::new(FastMap::default()),
            by_target: RwLock::new(FastMap::default()),
            by_type: RwLock::new(FastMap::default()),
            types: RwLock::new(FastMap::default()),
        }
    }

    /// Declare a relation type with inference semantics.
    pub fn type_snapshot(&self) -> Vec<RelationType> {
        self.types.read().values().cloned().collect()
    }

    pub fn declare_type(&self, rel_type: RelationType) {
        self.types.write().insert(rel_type.name.clone(), rel_type);
    }

    /// Add a relation to the store. Returns the index.
    pub fn add(&self, rel: &Relation) -> usize {
        let stored = StoredRelation {
            source_id: rel.source_id.clone(),
            relation_type: rel.relation_type.clone(),
            target_id: rel.target_id.clone(),
            properties: rel.properties.clone(),
            valid_from: rel.valid_from,
            valid_to: rel.valid_to,
            deleted: false,
        };

        let mut rels = self.relations.write();
        let idx = rels.len();
        rels.push(stored);

        self.by_source
            .write()
            .entry(rel.source_id.clone())
            .or_default()
            .push(idx);
        self.by_target
            .write()
            .entry(rel.target_id.clone())
            .or_default()
            .push(idx);
        self.by_type
            .write()
            .entry(rel.relation_type.clone())
            .or_default()
            .push(idx);

        idx
    }

    /// Remove a relation by matching source, type, target. Returns true if found.
    pub fn remove(&self, source_id: &str, relation_type: &str, target_id: &str) -> bool {
        let source_indices = self.by_source.read();
        let indices = match source_indices.get(source_id) {
            Some(v) => v.clone(),
            None => return false,
        };
        drop(source_indices);

        let mut rels = self.relations.write();
        for &idx in &indices {
            let r = &mut rels[idx];
            if !r.deleted && r.relation_type == relation_type && r.target_id == target_id {
                r.deleted = true;
                return true;
            }
        }
        false
    }

    /// Get all outgoing relations from a node, optionally filtered by type and time.
    pub fn outgoing(
        &self,
        source_id: &str,
        relation_type: Option<&str>,
        at_time: f64,
    ) -> Vec<Relation> {
        self.valid_relations(&self.by_source, source_id, relation_type, at_time)
    }

    /// Get all incoming relations to a node.
    pub fn incoming(
        &self,
        target_id: &str,
        relation_type: Option<&str>,
        at_time: f64,
    ) -> Vec<Relation> {
        self.valid_relations(&self.by_target, target_id, relation_type, at_time)
    }

    /// Relations listed under `node_id` in `index` that are live at `at_time`.
    fn valid_relations(
        &self,
        index: &RwLock<FastMap<String, Vec<usize>>>,
        node_id: &str,
        relation_type: Option<&str>,
        at_time: f64,
    ) -> Vec<Relation> {
        // `relations` before any index, the order `add` takes them in.
        let rels = self.relations.read();
        let index = index.read();
        let Some(indices) = index.get(node_id) else {
            return Vec::new();
        };
        indices
            .iter()
            .map(|&idx| &rels[idx])
            .filter(|r| {
                r.is_valid_at(at_time) && relation_type.is_none_or(|rt| r.relation_type == rt)
            })
            .map(StoredRelation::to_relation)
            .collect()
    }

    /// BFS traversal from start_id following relations of given type (or all if None).
    /// Returns paths up to max_depth hops. Includes inferred edges from symmetric/transitive types.
    pub fn traverse(
        &self,
        start_id: &str,
        relation_type: Option<&str>,
        max_depth: u32,
        at_time: f64,
        similarity_fn: &dyn Fn(&str, &str) -> f64,
    ) -> Vec<GraphPath> {
        let mut results = Vec::new();
        let mut visited = FastMap::<String, u32>::default();
        visited.insert(start_id.to_string(), 0);

        // BFS queue: (node_id, current_path, depth)
        let mut queue: VecDeque<(String, Vec<PathHop>, u32)> = VecDeque::new();
        queue.push_back((start_id.to_string(), Vec::new(), 0));

        let types = self.types.read();

        while let Some((node_id, path, depth)) = queue.pop_front() {
            if depth >= max_depth {
                continue;
            }

            // Collect neighbors: direct outgoing + symmetric incoming
            let mut neighbors: Vec<(String, String)> = Vec::new();

            for rel in self.outgoing(&node_id, relation_type, at_time) {
                neighbors.push((rel.target_id, rel.relation_type));
            }

            // Symmetric: incoming edges of symmetric types also count as outgoing
            if let Some(rt) = relation_type {
                if let Some(type_def) = types.get(rt) {
                    if type_def.symmetric {
                        for rel in self.incoming(&node_id, Some(rt), at_time) {
                            neighbors.push((rel.source_id, rel.relation_type));
                        }
                    }
                }
            } else {
                // No type filter: check each incoming relation's type for symmetry
                let rels = self.relations.read();
                let incoming_indices = self.by_target.read();
                if let Some(indices) = incoming_indices.get(&node_id) {
                    for &idx in indices {
                        let r = &rels[idx];
                        if !r.is_valid_at(at_time) || r.deleted {
                            continue;
                        }
                        if let Some(td) = types.get(&r.relation_type) {
                            if td.symmetric {
                                neighbors.push((r.source_id.clone(), r.relation_type.clone()));
                            }
                        }
                    }
                }
            }

            for (next_id, rel_type) in neighbors {
                let next_depth = depth + 1;
                let prev_depth = visited.get(&next_id).copied();
                if prev_depth.is_some_and(|d| d <= next_depth) {
                    continue;
                }
                visited.insert(next_id.clone(), next_depth);

                let sim = similarity_fn(&node_id, &next_id);
                let mut new_path = path.clone();
                new_path.push(PathHop {
                    node_id: next_id.clone(),
                    relation_type: rel_type,
                    similarity: sim,
                });

                let score = new_path.iter().map(|h| h.similarity).product::<f64>();
                results.push(GraphPath {
                    hops: new_path.clone(),
                    score,
                });

                if next_depth < max_depth {
                    queue.push_back((next_id, new_path, next_depth));
                }
            }
        }

        // Transitive inference: if we have A->B->C for a transitive type, add A->C
        if let Some(rt) = relation_type {
            if let Some(type_def) = types.get(rt) {
                if type_def.transitive {
                    let mut inferred = Vec::new();
                    for path in &results {
                        if path.hops.len() >= 2 {
                            let final_node = &path.hops.last().expect("len >= 2").node_id;
                            let sim = similarity_fn(start_id, final_node);
                            inferred.push(GraphPath {
                                hops: vec![PathHop {
                                    node_id: final_node.clone(),
                                    relation_type: rt.to_string(),
                                    similarity: sim,
                                }],
                                score: sim,
                            });
                        }
                    }
                    results.extend(inferred);
                }
            }
        }

        results.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results
    }

    /// Total number of live (non-deleted) relations.
    pub fn count(&self) -> usize {
        self.relations.read().iter().filter(|r| !r.deleted).count()
    }

    /// Snapshot all live relations for compaction.
    pub fn snapshot(&self) -> Vec<Relation> {
        self.relations
            .read()
            .iter()
            .filter(|r| !r.deleted)
            .map(|r| r.to_relation())
            .collect()
    }
}

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

    fn dummy_sim(_a: &str, _b: &str) -> f64 {
        0.8
    }

    /// An always-valid relation without properties.
    fn rel(source: &str, relation_type: &str, target: &str) -> Relation {
        Relation {
            source_id: source.into(),
            relation_type: relation_type.into(),
            target_id: target.into(),
            properties: None,
            valid_from: 0.0,
            valid_to: 0.0,
        }
    }

    #[test]
    fn test_add_and_query_outgoing() {
        let store = RelationStore::new();
        store.add(&rel("paris", "capital_of", "france"));
        store.add(&rel("berlin", "capital_of", "germany"));

        let out = store.outgoing("paris", Some("capital_of"), 0.0);
        assert_eq!(out.len(), 1);
        assert_eq!(out[0].target_id, "france");

        let out_all = store.outgoing("paris", None, 0.0);
        assert_eq!(out_all.len(), 1);

        let none = store.outgoing("tokyo", None, 0.0);
        assert!(none.is_empty());
    }

    #[test]
    fn test_remove_relation() {
        let store = RelationStore::new();
        store.add(&rel("a", "knows", "b"));
        assert_eq!(store.count(), 1);
        assert!(store.remove("a", "knows", "b"));
        assert_eq!(store.count(), 0);
        assert!(!store.remove("a", "knows", "b"));
    }

    #[test]
    fn test_symmetric_traversal() {
        let store = RelationStore::new();
        store.declare_type(RelationType {
            name: "friend_of".into(),
            transitive: false,
            symmetric: true,
        });
        store.add(&rel("alice", "friend_of", "bob"));

        // Bob should see Alice via symmetric inference
        let paths = store.traverse("bob", Some("friend_of"), 1, 0.0, &dummy_sim);
        assert!(!paths.is_empty());
        assert!(paths.iter().any(|p| p.hops[0].node_id == "alice"));
    }

    #[test]
    fn test_transitive_inference() {
        let store = RelationStore::new();
        store.declare_type(RelationType {
            name: "is_in".into(),
            transitive: true,
            symmetric: false,
        });
        store.add(&rel("paris", "is_in", "france"));
        store.add(&rel("france", "is_in", "europe"));

        let paths = store.traverse("paris", Some("is_in"), 3, 0.0, &dummy_sim);
        // Should find both direct (paris->france) and transitive (paris->europe)
        let targets: Vec<&str> = paths
            .iter()
            .flat_map(|p| p.hops.iter().map(|h| h.node_id.as_str()))
            .collect();
        assert!(targets.contains(&"france"));
        assert!(targets.contains(&"europe"));

        // Transitive inference should produce a 1-hop inferred path to europe
        let inferred_single_hop = paths
            .iter()
            .find(|p| p.hops.len() == 1 && p.hops[0].node_id == "europe");
        assert!(
            inferred_single_hop.is_some(),
            "Should have inferred single-hop to europe"
        );
    }

    #[test]
    fn test_multi_hop_traversal() {
        let store = RelationStore::new();
        store.add(&rel("a", "links", "b"));
        store.add(&rel("b", "links", "c"));
        store.add(&rel("c", "links", "d"));

        let paths = store.traverse("a", Some("links"), 3, 0.0, &dummy_sim);
        let max_hops = paths.iter().map(|p| p.hops.len()).max().unwrap_or(0);
        assert_eq!(max_hops, 3); // a->b->c->d
    }

    #[test]
    fn test_temporal_filter() {
        let store = RelationStore::new();
        store.add(&Relation {
            valid_from: 1000.0,
            valid_to: 2000.0,
            ..rel("a", "employs", "b")
        });
        store.add(&Relation {
            valid_from: 1500.0,
            valid_to: 0.0,
            ..rel("a", "employs", "c")
        });

        // At time 1200: only b
        let at_1200 = store.outgoing("a", None, 1200.0);
        assert_eq!(at_1200.len(), 1);
        assert_eq!(at_1200[0].target_id, "b");

        // At time 1700: both
        let at_1700 = store.outgoing("a", None, 1700.0);
        assert_eq!(at_1700.len(), 2);

        // At time 2500: only c (b expired)
        let at_2500 = store.outgoing("a", None, 2500.0);
        assert_eq!(at_2500.len(), 1);
        assert_eq!(at_2500[0].target_id, "c");
    }

    #[test]
    fn test_snapshot() {
        let store = RelationStore::new();
        store.add(&rel("a", "r", "b"));
        store.add(&rel("c", "r", "d"));
        store.remove("a", "r", "b");

        let snap = store.snapshot();
        assert_eq!(snap.len(), 1);
        assert_eq!(snap[0].source_id, "c");
    }
}