import gradio as gr
import networkx as nx
import matplotlib.pyplot as plt
from io import BytesIO

# Initialize the lesson plan graph
lesson_graph = nx.DiGraph()

def add_to_graph(teacher_name, subject, grade_level, learning_objective, activity, assessment, resource, school_board):
    global lesson_graph
    
    # Add nodes to the graph
    lesson_graph.add_node(teacher_name, type="User")
    lesson_graph.add_node(subject, type="Subject")
    lesson_graph.add_node(grade_level, type="Grade Level")
    lesson_graph.add_node(learning_objective, type="Learning Objective")
    lesson_graph.add_node(activity, type="Activity")
    lesson_graph.add_node(assessment, type="Assessment")
    lesson_graph.add_node(resource, type="Resource")
    lesson_graph.add_node(school_board, type="School Board")
    
    # Add edges to the graph
    lesson_graph.add_edge(teacher_name, subject, relationship="TEACHES")
    lesson_graph.add_edge(subject, learning_objective, relationship="COVERS")
    lesson_graph.add_edge(subject, grade_level, relationship="HAS_GRADE")
    lesson_graph.add_edge(activity, learning_objective, relationship="ACHIEVES")
    lesson_graph.add_edge(activity, resource, relationship="REQUIRES")
    lesson_graph.add_edge(learning_objective, assessment, relationship="EVALUATED_BY")
    lesson_graph.add_edge(teacher_name, school_board, relationship="BELONGS_TO")
    lesson_graph.add_edge(learning_objective, school_board, relationship="ALIGNS_WITH")
    
    # Generate search string
    search_string = f"{subject} {grade_level} {learning_objective} {activity} {resource}".strip()
    
    # Visualize the graph
    plt.figure(figsize=(12, 8))
    pos = nx.spring_layout(lesson_graph)
    nx.draw(lesson_graph, pos, with_labels=True, node_color="lightblue", 
            font_size=8, font_weight="bold", node_size=3000, node_shape="o")
    
    # Add node labels
    labels = nx.get_node_attributes(lesson_graph, 'type')
    nx.draw_networkx_labels(lesson_graph, pos, labels, font_size=6)
    
    # Add edge labels
    edge_labels = nx.get_edge_attributes(lesson_graph, 'relationship')
    nx.draw_networkx_edge_labels(lesson_graph, pos, edge_labels=edge_labels, font_size=6)
    
    # Save the plot to a bytes object
    buf = BytesIO()
    plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
    buf.seek(0)
    plt.close()
    
    return search_string, buf

def clear_graph():
    global lesson_graph
    lesson_graph.clear()
    return "Canvas cleared. You can start a new lesson plan."

# Gradio interface
demo = gr.Blocks()

with demo:
    gr.Markdown("# EduCanvas: Craft Your Lesson Masterpiece")
    gr.Markdown("Welcome to EduCanvas, where lesson planning becomes an art. Design, visualize, and perfect your educational masterpieces with ease.")
    
    with gr.Row():
        teacher_name = gr.Textbox(label="Teacher Name")
        school_board = gr.Textbox(label="School Board/Region")
    
    with gr.Row():
        subject = gr.Textbox(label="Subject")
        grade_level = gr.Textbox(label="Grade Level")
    
    with gr.Row():
        learning_objective = gr.Textbox(label="Learning Objective")
        activity = gr.Textbox(label="Activity")
    
    with gr.Row():
        assessment = gr.Textbox(label="Assessment")
        resource = gr.Textbox(label="Resource/Material")
    
    with gr.Row():
        generate_btn = gr.Button("Paint Your Lesson Plan")
        clear_btn = gr.Button("Clear Canvas")
    
    search_output = gr.Textbox(label="Content Discovery Search String")
    graph_output = gr.Image(label="Your Lesson Masterpiece")
    message_output = gr.Textbox(label="Canvas Status")
    
    generate_btn.click(
        add_to_graph, 
        inputs=[teacher_name, subject, grade_level, learning_objective, activity, assessment, resource, school_board],
        outputs=[search_output, graph_output]
    )
    
    clear_btn.click(clear_graph, outputs=message_output)

# Launch the EduCanvas app
demo.launch()