memory.py

"""
Hypergraph Memory System

This module implements a knowledge graph using NetworkX MultiDiGraph
to represent complex relationships between agents, shards, and assets.
"""

import random
from typing import List, Dict, Set, Optional, Tuple
import networkx as nx
import plotly.graph_objects as go


class HypergraphMemory:
    """
    Represents a hypergraph memory system for the ASI Chain.
    
    Stores nodes (agents, shards, assets) and labeled relationships
    between them. Supports queries and visualization.
    
    Node Types:
        - agent: Autonomous agents
        - shard: Processing shards
        - asset: Digital assets or resources
        
    Relation Types:
        - trusts: Agent trusts another agent
        - owns: Agent owns an asset
        - collaborates: Agents work together
        - trades: Trade relationship
        - processes: Shard processes transactions
        - stores: Shard stores data
    """
    
    def __init__(self):
        """Initialize an empty hypergraph memory."""
        self.graph = nx.MultiDiGraph()
        self.node_types: Dict[str, str] = {}
        self.relation_count: Dict[str, int] = {}
        
    def add_node(self, node_id: str, node_type: str, **attributes):
        """
        Add a node to the hypergraph.
        
        Args:
            node_id: Unique identifier for the node
            node_type: Type of node (agent, shard, asset)
            **attributes: Additional node attributes
        """
        self.graph.add_node(node_id, node_type=node_type, **attributes)
        self.node_types[node_id] = node_type
        
    def add_relation(self, source: str, target: str, relation_type: str, 
                    **attributes):
        """
        Add a labeled relation between two nodes.
        
        Args:
            source: Source node ID
            target: Target node ID
            relation_type: Type of relationship
            **attributes: Additional edge attributes
        """
        if source not in self.graph:
            raise ValueError(f"Source node {source} not in graph")
        if target not in self.graph:
            raise ValueError(f"Target node {target} not in graph")
        
        self.graph.add_edge(source, target, relation=relation_type, **attributes)
        
        # Update relation count
        self.relation_count[relation_type] = self.relation_count.get(relation_type, 0) + 1
        
    def query_relations(self, relation_type: Optional[str] = None) -> List[Tuple[str, str, Dict]]:
        """
        Query relations by type.
        
        Args:
            relation_type: Filter by relation type (None = all relations)
            
        Returns:
            List of (source, target, attributes) tuples
        """
        results = []
        
        for source, target, key, data in self.graph.edges(keys=True, data=True):
            if relation_type is None or data.get('relation') == relation_type:
                results.append((source, target, data))
        
        return results
    
    def get_neighbors(self, node_id: str, relation_type: Optional[str] = None) -> List[str]:
        """
        Get all neighbors of a node, optionally filtered by relation type.
        
        Args:
            node_id: Node to get neighbors for
            relation_type: Optional relation type filter
            
        Returns:
            List of neighbor node IDs
        """
        if node_id not in self.graph:
            return []
        
        neighbors = []
        
        # Outgoing edges
        for _, target, data in self.graph.out_edges(node_id, data=True):
            if relation_type is None or data.get('relation') == relation_type:
                neighbors.append(target)
        
        # Incoming edges
        for source, _, data in self.graph.in_edges(node_id, data=True):
            if relation_type is None or data.get('relation') == relation_type:
                if source not in neighbors:
                    neighbors.append(source)
        
        return neighbors
    
    def get_node_info(self, node_id: str) -> Optional[Dict]:
        """
        Get information about a node.
        
        Args:
            node_id: Node to get info for
            
        Returns:
            Dictionary with node data or None if not found
        """
        if node_id not in self.graph:
            return None
        
        data = dict(self.graph.nodes[node_id])
        
        # Add connection counts
        data['in_degree'] = self.graph.in_degree(node_id)
        data['out_degree'] = self.graph.out_degree(node_id)
        data['total_connections'] = data['in_degree'] + data['out_degree']
        
        return data
    
    def get_subgraph(self, node_ids: List[str]) -> nx.MultiDiGraph:
        """
        Get a subgraph containing only specified nodes.
        
        Args:
            node_ids: List of node IDs to include
            
        Returns:
            NetworkX MultiDiGraph subgraph
        """
        return self.graph.subgraph(node_ids).copy()
    
    def get_nodes_by_type(self, node_type: str) -> List[str]:
        """
        Get all nodes of a specific type.
        
        Args:
            node_type: Type to filter by
            
        Returns:
            List of node IDs
        """
        return [node for node, data in self.graph.nodes(data=True)
                if data.get('node_type') == node_type]
    
    def visualize(self, width: int = 900, height: int = 700, 
                  relation_filter: Optional[str] = None) -> go.Figure:
        """
        Create an interactive Plotly visualization of the hypergraph.
        
        Args:
            width: Figure width in pixels
            height: Figure height in pixels
            relation_filter: Optional relation type to filter edges
            
        Returns:
            Plotly Figure object
        """
        if len(self.graph.nodes) == 0:
            fig = go.Figure()
            fig.add_annotation(text="No nodes in memory graph",
                             xref="paper", yref="paper",
                             x=0.5, y=0.5, showarrow=False)
            return fig
        
        # Create layout using spring layout
        pos = nx.spring_layout(self.graph, seed=42, k=1.5, iterations=50)
        
        # Define colors for node types
        node_colors = {
            'agent': '#FF6B6B',     # Red
            'shard': '#4ECDC4',     # Teal
            'asset': '#FFE66D',     # Yellow
            'default': '#95A5A6'    # Gray
        }
        
        # Create edge traces (grouped by relation type)
        edge_traces = []
        relation_types = set()
        
        for source, target, key, data in self.graph.edges(keys=True, data=True):
            relation = data.get('relation', 'unknown')
            
            # Apply filter if specified
            if relation_filter and relation != relation_filter:
                continue
            
            relation_types.add(relation)
            
            x0, y0 = pos[source]
            x1, y1 = pos[target]
            
            edge_trace = go.Scatter(
                x=[x0, x1, None],
                y=[y0, y1, None],
                mode='lines',
                line=dict(width=1, color='#888'),
                hoverinfo='text',
                hovertext=f"{relation}: {source[:8]}... → {target[:8]}...",
                showlegend=False,
                name=relation
            )
            edge_traces.append(edge_trace)
        
        # Create node trace
        node_x = []
        node_y = []
        node_text = []
        node_color = []
        node_size = []
        
        for node_id in self.graph.nodes():
            x, y = pos[node_id]
            node_x.append(x)
            node_y.append(y)
            
            node_data = self.graph.nodes[node_id]
            node_type = node_data.get('node_type', 'default')
            
            # Node color based on type
            color = node_colors.get(node_type, node_colors['default'])
            node_color.append(color)
            
            # Node size based on connections
            degree = self.graph.degree(node_id)
            node_size.append(max(10, min(30, 10 + degree * 2)))
            
            # Hover text
            name = node_data.get('name', node_id[:8])
            hover_text = (f"<b>{name}</b><br>"
                         f"Type: {node_type}<br>"
                         f"ID: {node_id[:12]}...<br>"
                         f"Connections: {degree}")
            node_text.append(hover_text)
        
        node_trace = go.Scatter(
            x=node_x, y=node_y,
            mode='markers',
            hoverinfo='text',
            hovertext=node_text,
            marker=dict(
                size=node_size,
                color=node_color,
                line=dict(width=2, color='white')
            ),
            showlegend=False
        )
        
        # Create figure
        fig = go.Figure(data=edge_traces + [node_trace])
        
        # Add legend for node types
        for node_type, color in node_colors.items():
            if node_type != 'default':
                fig.add_trace(go.Scatter(
                    x=[None], y=[None],
                    mode='markers',
                    marker=dict(size=10, color=color),
                    showlegend=True,
                    name=node_type.capitalize()
                ))
        
        filter_text = f" (filtered: {relation_filter})" if relation_filter else ""
        
        fig.update_layout(
            title=dict(text=f"Hypergraph Memory Network{filter_text}", font=dict(size=16)),
            showlegend=True,
            hovermode='closest',
            margin=dict(b=20, l=5, r=5, t=40),
            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            width=width,
            height=height,
            plot_bgcolor='rgba(240,240,240,0.9)',
            legend=dict(x=0.02, y=0.98)
        )
        
        return fig
    
    def get_statistics(self) -> Dict:
        """
        Get statistics about the hypergraph.
        
        Returns:
            Dictionary with graph statistics
        """
        node_type_counts = {}
        for node_type in set(self.node_types.values()):
            node_type_counts[node_type] = sum(1 for nt in self.node_types.values() if nt == node_type)
        
        return {
            'total_nodes': len(self.graph.nodes),
            'total_edges': len(self.graph.edges),
            'node_types': node_type_counts,
            'relation_types': dict(self.relation_count),
            'avg_degree': sum(dict(self.graph.degree()).values()) / len(self.graph.nodes) if self.graph.nodes else 0
        }
    
    def __repr__(self) -> str:
        return f"HypergraphMemory(nodes={len(self.graph.nodes)}, edges={len(self.graph.edges)})"


if __name__ == "__main__":
    # Demo usage
    from agents import generate_agents
    from shards import create_default_shards
    
    print("=== Hypergraph Memory Demo ===\n")
    
    # Initialize
    random.seed(42)
    memory = HypergraphMemory()
    
    # Add agents
    agents = generate_agents(5)
    for agent in agents:
        memory.add_node(agent.agent_id, 'agent', name=agent.name, reputation=agent.reputation)
    
    # Add shards
    shards = create_default_shards()
    for shard in shards:
        memory.add_node(shard.name, 'shard', name=shard.name, type=shard.shard_type)
    
    # Add some assets
    for i in range(3):
        asset_id = f"asset_{i}"
        memory.add_node(asset_id, 'asset', name=f"Asset-{i}", value=random.randint(100, 1000))
    
    print(f"Added {len(agents)} agents, {len(shards)} shards, 3 assets\n")
    
    # Add relationships
    print("Creating relationships...")
    
    # Agents trust each other
    for i, agent in enumerate(agents[:-1]):
        target = agents[i + 1]
        memory.add_relation(agent.agent_id, target.agent_id, 'trusts', strength=random.uniform(0.5, 1.0))
    
    # Agents own assets
    all_assets = memory.get_nodes_by_type('asset')
    for agent in agents[:3]:
        asset = random.choice(all_assets)
        memory.add_relation(agent.agent_id, asset, 'owns', amount=random.randint(1, 10))
    
    # Agents collaborate
    for _ in range(3):
        a1, a2 = random.sample(agents, 2)
        memory.add_relation(a1.agent_id, a2.agent_id, 'collaborates', project=f"project_{random.randint(1, 5)}")
    
    # Agents use shards
    for agent in agents:
        shard = random.choice(shards)
        memory.add_relation(agent.agent_id, shard.name, 'uses', frequency='high')
    
    print(f"Relationships created\n")
    print("="*60 + "\n")
    
    # Query examples
    print("Query Examples:\n")
    
    # All trust relations
    trust_relations = memory.query_relations('trusts')
    print(f"Trust relationships: {len(trust_relations)}")
    for source, target, data in trust_relations[:3]:
        print(f"  {source[:8]}... trusts {target[:8]}... (strength: {data.get('strength', 0):.2f})")
    
    print()
    
    # Neighbors of first agent
    first_agent = agents[0]
    neighbors = memory.get_neighbors(first_agent.agent_id)
    print(f"Neighbors of {first_agent.name}: {len(neighbors)}")
    for neighbor in neighbors[:3]:
        print(f"  {neighbor[:12]}...")
    
    print("\n" + "="*60 + "\n")
    
    # Statistics
    stats = memory.get_statistics()
    print("Memory Statistics:")
    print(f"  Total nodes: {stats['total_nodes']}")
    print(f"  Total edges: {stats['total_edges']}")
    print(f"  Node types: {stats['node_types']}")
    print(f"  Relation types: {stats['relation_types']}")
    print(f"  Avg degree: {stats['avg_degree']:.2f}")