data_processing/graph_builder.py

"""Knowledge Graph builder for Myanmar Ghost project.

Represents conversational context as a knowledge graph for better
understanding of complex social interactions.
Example: (Speaker, Role, Customer) --[located_in]--> (Restaurant)
"""

import json
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Set, Tuple

import networkx as nx


class NodeType(str, Enum):
    """Types of nodes in the knowledge graph."""
    SPEAKER = "speaker"
    UTTERANCE = "utterance"
    LOCATION = "location"
    ORGANIZATION = "organization"
    EMOTION = "emotion"
    TOPIC = "topic"
    ACTION = "action"
    TIME = "time"


class RelationType(str, Enum):
    """Types of relations between nodes."""
    SPEAKS = "speaks"
    LOCATED_IN = "located_in"
    WORKS_AT = "works_at"
    VISITS = "visits"
    FEELS = "feels"
    ABOUT = "about"
    BEFORE = "before"
    AFTER = "after"
    IN_RESPONSE_TO = "in_response_to"
    CONTAINS = "contains"
    HAS_ROLE = "has_role"


@dataclass
class Entity:
    """Represents an entity in the knowledge graph."""
    id: str
    type: NodeType
    properties: Dict[str, Any] = field(default_factory=dict)
    aliases: List[str] = field(default_factory=list)
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "id": self.id,
            "type": self.type.value,
            "properties": self.properties,
            "aliases": self.aliases,
        }


@dataclass
class Relation:
    """Represents a relation between entities."""
    source: str  # Entity ID
    target: str  # Entity ID
    type: RelationType
    properties: Dict[str, Any] = field(default_factory=dict)
    confidence: float = 1.0
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            "source": self.source,
            "target": self.target,
            "type": self.type.value,
            "properties": self.properties,
            "confidence": self.confidence,
        }


class MyanmarKnowledgeGraph:
    """Build and manage knowledge graph for Myanmar conversations."""
    
    # Common Myanmar entities
    LOCATIONS = {
        "စားသောက်ဆိုင်": NodeType.LOCATION,
        "ဆေးရုံ": NodeType.LOCATION,
        "ဈေး": NodeType.LOCATION,
        "ရုံး": NodeType.LOCATION,
        "အိမ်": NodeType.LOCATION,
    }
    
    EMOTIONS = {
        "ပျော်": NodeType.EMOTION,
        "စိတ်ဓာတ်ကျ": NodeType.EMOTION,
        "ဒေါသ": NodeType.EMOTION,
        "ဝမ်းနည်း": NodeType.EMOTION,
        "ပိုးပါး": NodeType.EMOTION,
    }
    
    ROLES = {
        "ဖေါ်သည်": "customer",
        "ဝန်ထမ်း": "staff",
        "ဆရာဝန်": "doctor",
        "ပါးရှင်း": "patient",
        "အရာရှိ": "manager",
    }
    
    def __init__(self):
        self.graph = nx.MultiDiGraph()
        self.entity_index: Dict[str, Entity] = {}
        self.session_id = 0
    
    def add_entity(self, entity: Entity) -> None:
        """Add an entity to the graph."""
        self.entity_index[entity.id] = entity
        self.graph.add_node(
            entity.id,
            type=entity.type.value,
            **entity.properties,
        )
    
    def add_relation(self, relation: Relation) -> None:
        """Add a relation between entities."""
        self.graph.add_edge(
            relation.source,
            relation.target,
            type=relation.type.value,
            **relation.properties,
        )
    
    def extract_speaker_entity(
        self,
        speaker_id: str,
        role: Optional[str] = None,
    ) -> Entity:
        """Create a speaker entity from utterance metadata."""
        entity = Entity(
            id=f"speaker_{speaker_id}",
            type=NodeType.SPEAKER,
            properties={
                "role": role or "unknown",
                "session": self.session_id,
            },
        )
        self.add_entity(entity)
        return entity
    
    def extract_utterance_entity(
        self,
        text: str,
        speaker_id: str,
        timestamp: float,
        prosody: Optional[Dict] = None,
    ) -> Tuple[Entity, List[Entity], List[Relation]]:
        """Extract utterance and related entities from text."""
        utterance_id = f"utt_{speaker_id}_{int(timestamp * 1000)}"
        
        utterance = Entity(
            id=utterance_id,
            type=NodeType.UTTERANCE,
            properties={
                "text": text,
                "timestamp": timestamp,
                "prosody": prosody or {},
            },
        )
        self.add_entity(utterance)
        
        # Extract related entities
        related_entities = []
        relations = []
        
        # Extract location mentions
        for loc, _ in self.LOCATIONS.items():
            if loc in text:
                loc_entity = Entity(
                    id=f"loc_{loc}_{self.session_id}",
                    type=NodeType.LOCATION,
                    properties={"name": loc},
                )
                self.add_entity(loc_entity)
                related_entities.append(loc_entity)
                
                relation = Relation(
                    source=utterance_id,
                    target=loc_entity.id,
                    type=RelationType.LOCATED_IN,
                )
                self.add_relation(relation)
                relations.append(relation)
        
        # Extract emotion mentions
        for emotion, _ in self.EMOTIONS.items():
            if emotion in text:
                emotion_entity = Entity(
                    id=f"emotion_{emotion}_{self.session_id}",
                    type=NodeType.EMOTION,
                    properties={"name": emotion},
                )
                self.add_entity(emotion_entity)
                related_entities.append(emotion_entity)
                
                relation = Relation(
                    source=utterance_id,
                    target=emotion_entity.id,
                    type=RelationType.FEELS,
                )
                self.add_relation(relation)
                relations.append(relation)
        
        # Link to speaker
        speaker_entity = self.entity_index.get(f"speaker_{speaker_id}")
        if speaker_entity:
            relation = Relation(
                source=speaker_entity.id,
                target=utterance_id,
                type=RelationType.SPEAKS,
            )
            self.add_relation(relation)
            relations.append(relation)
        
        return utterance, related_entities, relations
    
    def build_from_conversation(
        self,
        utterances: List[Dict],
        context: Optional[Dict] = None,
    ) -> nx.MultiDiGraph:
        """Build knowledge graph from conversation data."""
        self.session_id += 1
        
        # Set context entities
        if context:
            for key, value in context.items():
                if key == "location" and value in self.LOCATIONS:
                    loc_entity = Entity(
                        id=f"context_location",
                        type=NodeType.LOCATION,
                        properties={"name": value},
                    )
                    self.add_entity(loc_entity)
        
        prev_utterance = None
        
        for i, utt_data in enumerate(utterances):
            speaker_id = utt_data.get("speaker_id", f"s_{i}")
            text = utt_data.get("text", "")
            timestamp = utt_data.get("timestamp", i)
            prosody = utt_data.get("prosody")
            role = utt_data.get("role")
            
            # Add speaker
            self.extract_speaker_entity(speaker_id, role)
            
            # Add utterance
            utterance, related, _ = self.extract_utterance_entity(
                text, speaker_id, timestamp, prosody
            )
            
            # Link to previous utterance (temporal relation)
            if prev_utterance:
                relation = Relation(
                    source=prev_utterance.id,
                    target=utterance.id,
                    type=RelationType.BEFORE,
                )
                self.add_relation(relation)
                
                # In response relation
                response_relation = Relation(
                    source=utterance.id,
                    target=prev_utterance.id,
                    type=RelationType.IN_RESPONSE_TO,
                )
                self.add_relation(response_relation)
            
            prev_utterance = utterance
        
        return self.graph
    
    def query_path(
        self,
        source_type: NodeType,
        target_type: NodeType,
        relation_type: Optional[RelationType] = None,
    ) -> List[Tuple[Entity, Entity, Relation]]:
        """Query paths between entity types."""
        results = []
        
        for source_id in self.entity_index:
            source = self.entity_index[source_id]
            if source.type != source_type:
                continue
            
            for target_id in self.entity_index:
                target = self.entity_index[target_id]
                if target.type != target_type:
                    continue
                
                # Find paths
                try:
                    if relation_type:
                        edges = self.graph.get_edge_data(source_id, target_id)
                        if edges:
                            for edge_data in edges.values():
                                if edge_data.get("type") == relation_type.value:
                                    relation = Relation(
                                        source=source_id,
                                        target=target_id,
                                        type=relation_type,
                                        properties=edge_data,
                                    )
                                    results.append((source, target, relation))
                    else:
                        if nx.has_path(self.graph, source_id, target_id):
                            path = nx.shortest_path(
                                self.graph, source_id, target_id
                            )
                            if len(path) == 2:
                                relation = Relation(
                                    source=source_id,
                                    target=target_id,
                                    type=RelationType.CONTAINS,
                                )
                                results.append((source, target, relation))
                except nx.NetworkXError:
                    continue
        
        return results
    
    def get_utterance_context(self, utterance_id: str) -> Dict:
        """Get full context for an utterance."""
        if utterance_id not in self.entity_index:
            return {}
        
        context = {
            "utterance": self.entity_index[utterance_id].to_dict(),
            "speaker": None,
            "previous": None,
            "next": None,
            "locations": [],
            "emotions": [],
        }
        
        # Get speaker
        for edge in self.graph.out_edges(utterance_id, data=True):
            if edge[2].get("type") == RelationType.FEELS.value:
                context["emotions"].append(self.entity_index[edge[1]].to_dict())
            if edge[2].get("type") == RelationType.LOCATED_IN.value:
                context["locations"].append(self.entity_index[edge[1]].to_dict())
        
        # Get predecessor/successor
        predecessors = list(self.graph.predecessors(utterance_id))
        successors = list(self.graph.successors(utterance_id))
        
        for pred_id in predecessors:
            pred = self.entity_index.get(pred_id)
            if pred and pred.type == NodeType.UTTERANCE:
                context["previous"] = pred.to_dict()
                break
        
        for succ_id in successors:
            succ = self.entity_index.get(succ_id)
            if succ and succ.type == NodeType.UTTERANCE:
                context["next"] = succ.to_dict()
                break
        
        return context
    
    def export_to_json(self, path: str) -> None:
        """Export graph to JSON format."""
        entities = [e.to_dict() for e in self.entity_index.values()]
        
        relations = []
        for source, target, data in self.graph.edges(data=True):
            relations.append({
                "source": source,
                "target": target,
                "type": data.get("type"),
                **data,
            })
        
        output = {
            "entities": entities,
            "relations": relations,
            "metadata": {
                "num_entities": len(entities),
                "num_relations": len(relations),
                "session_id": self.session_id,
            },
        }
        
        with open(path, "w", encoding="utf-8") as f:
            json.dump(output, f, indent=2, ensure_ascii=False)
    
    def load_from_json(self, path: str) -> None:
        """Load graph from JSON format."""
        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)
        
        self.entity_index = {}
        self.graph = nx.MultiDiGraph()
        
        for entity_data in data.get("entities", []):
            entity = Entity(
                id=entity_data["id"],
                type=NodeType(entity_data["type"]),
                properties=entity_data.get("properties", {}),
                aliases=entity_data.get("aliases", []),
            )
            self.add_entity(entity)
        
        for rel_data in data.get("relations", []):
            relation = Relation(
                source=rel_data["source"],
                target=rel_data["target"],
                type=RelationType(rel_data["type"]),
                properties=rel_data,
                confidence=rel_data.get("confidence", 1.0),
            )
            self.add_relation(relation)
    
    def visualize(self) -> nx.MultiDiGraph:
        """Return the graph for visualization."""
        return self.graph


def create_knowledge_graph() -> MyanmarKnowledgeGraph:
    """Factory function to create knowledge graph."""
    return MyanmarKnowledgeGraph()


if __name__ == "__main__":
    # Example usage
    kg = create_knowledge_graph()
    
    # Sample conversation
    utterances = [
        {
            "speaker_id": "customer_1",
            "text": "ဆိုင်သို့ ကျွန်ုပ်လာပါပြီ",
            "timestamp": 0,
            "role": "customer",
        },
        {
            "speaker_id": "staff_1",
            "text": "ကြိုဆိုပါတယ်",
            "timestamp": 1,
            "role": "staff",
        },
        {
            "speaker_id": "customer_1",
            "text": "ကျေးဇူးပါ",
            "timestamp": 2,
            "prosody": {"mean_pitch": 150, "speaking_rate": 3},
            "role": "customer",
        },
    ]
    
    context = {"location": "စားသောက်ဆိုင်"}
    
    kg.build_from_conversation(utterances, context)
    
    # Export
    kg.export_to_json("data/graph/conversation_graph.json")
    print(f"Graph exported with {len(kg.entity_index)} entities")