src/enhancements/knowledge_graph/graph.py

"""
Knowledge Graph Implementation

Graph-based knowledge representation using networkx.
"""

from typing import List, Dict, Optional, Any, Set, Tuple
import networkx as nx
import numpy as np
from datetime import datetime
import json
from pathlib import Path


class KnowledgeGraph:
    """Graph-based knowledge representation with entities and relationships."""

    def __init__(
        self,
        max_nodes: int = 10000,
        max_edges: int = 50000,
    ):
        """
        Initialize the knowledge graph.

        Args:
            max_nodes: Maximum number of nodes
            max_edges: Maximum number of edges
        """
        self.max_nodes = max_nodes
        self.max_edges = max_edges
        self.graph = nx.MultiDiGraph()
        self._node_attributes: Dict[str, Dict[str, Any]] = {}
        self._edge_attributes: Dict[Tuple[str, str], Dict[str, Any]] = {}

    def add_entity(
        self,
        entity_id: str,
        entity_type: str,
        properties: Optional[Dict[str, Any]] = None,
    ) -> bool:
        """
        Add an entity to the knowledge graph.

        Args:
            entity_id: Unique identifier for the entity
            entity_type: Type of entity (e.g., 'person', 'concept', 'code')
            properties: Additional properties

        Returns:
            True if added, False if limit reached
        """
        if self.graph.number_of_nodes() >= self.max_nodes:
            return False

        if entity_id not in self.graph:
            self.graph.add_node(entity_id, type=entity_type)

        self._node_attributes[entity_id] = {
            "type": entity_type,
            "created_at": datetime.now().isoformat(),
            "properties": properties or {},
        }

        return True

    def add_relationship(
        self,
        source_id: str,
        target_id: str,
        relationship_type: str,
        properties: Optional[Dict[str, Any]] = None,
    ) -> bool:
        """
        Add a relationship between entities.

        Args:
            source_id: Source entity ID
            target_id: Target entity ID
            relationship_type: Type of relationship
            properties: Additional properties

        Returns:
            True if added, False if limit reached or entities don't exist
        """
        if self.graph.number_of_edges() >= self.max_edges:
            return False

        # Ensure entities exist
        if source_id not in self.graph:
            self.add_entity(source_id, "unknown")
        if target_id not in self.graph:
            self.add_entity(target_id, "unknown")

        self.graph.add_edge(source_id, target_id, type=relationship_type)

        edge_key = (source_id, target_id)
        self._edge_attributes[edge_key] = {
            "type": relationship_type,
            "created_at": datetime.now().isoformat(),
            "properties": properties or {},
        }

        return True

    def get_entity(self, entity_id: str) -> Optional[Dict[str, Any]]:
        """Get entity information."""
        if entity_id not in self.graph:
            return None

        return {
            "id": entity_id,
            **self._node_attributes.get(entity_id, {}),
        }

    def get_relationships(
        self,
        entity_id: str,
        relationship_type: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """Get relationships for an entity."""
        if entity_id not in self.graph:
            return []

        relationships = []
        for source, target, data in self.graph.edges(data=True):
            if source == entity_id or target == entity_id:
                rel_type = data.get("type", "unknown")
                if relationship_type and rel_type != relationship_type:
                    continue

                relationships.append({
                    "source": source,
                    "target": target,
                    "type": rel_type,
                })

        return relationships

    def find_similar_entities(
        self,
        entity_id: str,
        max_results: int = 5,
    ) -> List[Tuple[str, float]]:
        """
        Find similar entities using graph-based similarity.

        Args:
            entity_id: Entity to find similar
            max_results: Maximum number of results

        Returns:
            List of (entity_id, similarity_score) tuples
        """
        if entity_id not in self.graph:
            return []

        # Use common neighbors as simple similarity
        neighbors = set(self.graph.neighbors(entity_id))
        scores = []

        for node in self.graph.nodes():
            if node == entity_id:
                continue

            node_neighbors = set(self.graph.neighbors(node))
            common = len(neighbors & node_neighbors)

            if common > 0:
                # Jaccard-like similarity
                union = len(neighbors | node_neighbors)
                score = common / union if union > 0 else 0
                scores.append((node, score))

        scores.sort(key=lambda x: -x[1])
        return scores[:max_results]

    def search_entities(
        self,
        entity_type: Optional[str] = None,
        property_filter: Optional[Dict[str, Any]] = None,
    ) -> List[str]:
        """
        Search for entities.

        Args:
            entity_type: Filter by entity type
            property_filter: Filter by properties

        Returns:
            List of matching entity IDs
        """
        results = []

        for node in self.graph.nodes():
            attrs = self._node_attributes.get(node, {})

            # Check type filter
            if entity_type and attrs.get("type") != entity_type:
                continue

            # Check property filter
            if property_filter:
                props = attrs.get("properties", {})
                if not all(props.get(k) == v for k, v in property_filter.items()):
                    continue

            results.append(node)

        return results

    def get_subgraph(
        self,
        entity_ids: List[str],
        depth: int = 1,
    ) -> nx.MultiDiGraph:
        """
        Get a subgraph around specified entities.

        Args:
            entity_ids: Center entities
            depth: How many hops to include

        Returns:
            Subgraph
        """
        nodes = set(entity_ids)

        for _ in range(depth):
            for entity in list(nodes):
                nodes.update(self.graph.neighbors(entity))

        return self.graph.subgraph(nodes).copy()

    def export_json(self, filepath: str) -> None:
        """Export graph to JSON."""
        data = {
            "nodes": [
                {
                    "id": node,
                    **self._node_attributes.get(node, {}),
                }
                for node in self.graph.nodes()
            ],
            "edges": [
                {
                    "source": source,
                    "target": target,
                    "type": data.get("type", "unknown"),
                }
                for source, target, data in self.graph.edges(data=True)
            ],
        }

        Path(filepath).write_text(json.dumps(data, indent=2))

    def import_json(self, filepath: str) -> None:
        """Import graph from JSON."""
        data = json.loads(Path(filepath).read_text())

        for node_data in data.get("nodes", []):
            node_id = node_data.pop("id")
            self.add_entity(node_id, node_data.get("type", "unknown"), node_data.get("properties"))

        for edge_data in data.get("edges", []):
            self.add_relationship(
                edge_data["source"],
                edge_data["target"],
                edge_data.get("type", "unknown"),
            )

    def get_stats(self) -> Dict[str, Any]:
        """Get graph statistics."""
        return {
            "num_nodes": self.graph.number_of_nodes(),
            "num_edges": self.graph.number_of_edges(),
            "num_node_types": len(set(
                attrs.get("type") for attrs in self._node_attributes.values()
            )),
            "num_edge_types": len(set(
                data.get("type") for _, _, data in self.graph.edges(data=True)
            )),
        }

    def __repr__(self) -> str:
        stats = self.get_stats()
        return f"KnowledgeGraph(nodes={stats['num_nodes']}, edges={stats['num_edges']})"