import nbformat as nbf

# Create a new notebook
nb = nbf.v4.new_notebook()

# Add cells
nb.cells = [
    nbf.v4.new_markdown_cell("# Demo: AI-Powered Scientific Research Companion\nThis notebook demonstrates how to use the `Dispatcher` to search for papers, retrieve reproducible notebook cells, and fetch a knowledge graph."),
    nbf.v4.new_code_cell("""\
from orchestrator.dispatcher import Dispatcher

# Initialize dispatcher
dispatcher = Dispatcher()

# Example query
query = "CRISPR delivery"

# 1. Search for papers
papers = dispatcher.search_papers(query, limit=3)
print("Papers found:")
for p in papers:
    print(f"- {p['title']} (ID: {p['id']})")
"""),
    nbf.v4.new_code_cell("""\
# 2. Retrieve notebook cells for the first paper
if papers:
    first_id = papers[0]['id']
    cells = dispatcher.get_notebook_cells(first_id)
    print(f"Notebook cells for paper {first_id}:")
    for i, cell in enumerate(cells, 1):
        print(f"Cell {i}:")
        print(cell)
        print("------")
"""),
    nbf.v4.new_code_cell("""\
# 3. Fetch knowledge graph for the first paper
if papers:
    graph = dispatcher.get_graph(first_id)
    print("Graph nodes:")
    for node in graph.get("nodes", []):
        print(node)
    print("Graph edges:")
    for edge in graph.get("edges", []):
        print(edge)
""")
]

# Save the notebook
nb_path = "/mnt/data/demo.ipynb"
with open(nb_path, "w") as f:
    nbf.write(nb, f)

nb_path