from openai import OpenAI

def extract_entities_and_relations(text, model="gpt-4"):
    prompt = f"""Extract all entities and relationships from this text.
    
    Text: {text}
    
    Return JSON format:
    {{
        "entities": [
            {{"name": "Entity Name", "type": "PERSON|ORG|LOCATION|CONCEPT", "description": "Brief description"}}
        ],
        "relations": [
            {{"source": "Entity1", "target": "Entity2", "relation": "relationship type", "weight": 0.0-1.0}}
        ]
    }}"""
    
    response = OpenAI().chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        response_format={"type": "json_object"}
    )
    return json.loads(response.choices[0].message.content)

import networkx as nx

class KnowledgeGraph:
    def __init__(self):
        self.graph = nx.DiGraph()
    
    def add_document(self, doc_id, text):
        """Extract and add entities from a document."""
        extraction = extract_entities_and_relations(text)
        
        for entity in extraction["entities"]:
            self.graph.add_node(
                entity["name"],
                type=entity["type"],
                description=entity["description"],
                documents=[doc_id]
            )
        
        for relation in extraction["relations"]:
            if self.graph.has_edge(relation["source"], relation["target"]):
                # Update weight
                self.graph[relation["source"]][relation["target"]]["weight"] += relation["weight"]
            else:
                self.graph.add_edge(
                    relation["source"],
                    relation["target"],
                    relation=relation["relation"],
                    weight=relation["weight"]
                )
    
    def find_path(self, entity1, entity2, max_hops=3):
        """Find shortest path between two entities."""
        try:
            path = nx.shortest_path(self.graph, entity1, entity2)
            return path[:max_hops + 1]
        except nx.NetworkXNoPath:
            return None
    
    def get_entity_context(self, entity_name, max_hops=2):
        """Get all related entities and their descriptions."""
        if entity_name not in self.graph:
            return ""
        
        context = []
        for neighbor in nx.neighbors(self.graph, entity_name):
            edge_data = self.graph[entity_name][neighbor]
            node_data = self.graph.nodes[neighbor]
            context.append(
                f"{entity_name} --[{edge_data['relation']}]--> {neighbor} "
                f"({node_data.get('description', '')})"
            )
        
        return "\n".join(context)

import community as community_louvain

def detect_communities(knowledge_graph):
    """Detect communities in the knowledge graph."""
    partition = community_louvain.best_partition(knowledge_graph.graph.to_undirected())
    
    communities = {}
    for node, comm_id in partition.items():
        if comm_id not in communities:
            communities[comm_id] = []
        communities[comm_id].append(node)
    
    return communities

def summarize_community(community_nodes, knowledge_graph, llm):
    """Generate a summary for a community."""
    node_descriptions = []
    for node in community_nodes:
        desc = knowledge_graph.graph.nodes[node].get("description", "")
        node_descriptions.append(f"{node}: {desc}")
    
    prompt = f"""Summarize the following group of related entities and their relationships:

{chr(10).join(node_descriptions)}

Provide a concise summary of the key themes and relationships:"""
    
    response = llm.generate(prompt)
    return response

def local_retrieval(query, knowledge_graph, llm):
    # Extract entities from query
    query_entities = extract_entities_from_query(query, llm)
    
    context_parts = []
    for entity in query_entities:
        # Get entity neighborhood
        context = knowledge_graph.get_entity_context(entity, max_hops=2)
        if context:
            context_parts.append(context)
    
    return "\n\n".join(context_parts)

def global_retrieval(query, communities, community_summaries, llm):
    # Identify relevant communities
    relevant_comms = []
    for comm_id, summary in community_summaries.items():
        relevance = compute_relevance(query, summary, llm)
        if relevance > 0.5:
            relevant_comms.append((comm_id, summary))
    
    # Combine summaries
    context = "\n\n".join([s for _, s in relevant_comms])
    return context

def hybrid_graph_retrieval(query, knowledge_graph, text_index, llm):
    """Combine graph and text retrieval."""
    # Graph retrieval for entities
    graph_context = local_retrieval(query, knowledge_graph, llm)
    
    # Text retrieval for supporting details
    text_context = text_index.search(query, top_k=5)
    
    # Combine and rank
    combined_context = f"Graph Knowledge:\n{graph_context}\n\nSupporting Text:\n{text_context}"
    return combined_context