import networkx as nx
from typing import Any, Dict, List, Optional

# Scientist AI Knowledge Brain

class KnowledgeBrain:
    """Stores and organizes mathematical knowledge as a graph of concepts, relations, equations, and dependencies."""
    def __init__(self):
        self.graph = nx.MultiDiGraph()
        self.concepts = {}

    def add_concept(self, concept: str, data: Dict[str, Any]):
        self.concepts[concept] = data
        self.graph.add_node(concept, **data)

    def add_relation(self, source: str, target: str, relation: str, attributes: Optional[Dict[str, Any]] = None):
        if attributes is None:
            attributes = {}
        self.graph.add_edge(source, target, relation=relation, **attributes)

    def get_concept(self, concept: str) -> Optional[Dict[str, Any]]:
        return self.concepts.get(concept)

    def get_relations(self, concept: str) -> List[Dict[str, Any]]:
        edges = self.graph.out_edges(concept, data=True)
        return [dict(target=target, **data) for _, target, data in edges]

    def load_papers(self, papers: List[Dict[str, Any]]):
        for paper in papers:
            self.add_concept(paper['title'], {'type': 'paper', **paper})

    def load_formulas(self, formulas: List[Dict[str, Any]]):
        for formula in formulas:
            self.add_concept(formula['name'], {'type': 'formula', **formula})

    def load_proofs(self, proofs: List[Dict[str, Any]]):
        for proof in proofs:
            self.add_concept(proof['theorem'], {'type': 'proof', **proof})

    def load_results(self, results: List[Dict[str, Any]]):
        for result in results:
            self.add_concept(result['name'], {'type': 'result', **result})

    def load_failures(self, failures: List[Dict[str, Any]]):
        for failure in failures:
            self.add_concept(failure['name'], {'type': 'failure', **failure})
