import gradio as gr
import networkx as nx
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image
import matplotlib.patches as mpatches

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

# Define color map for node types
color_map = {
    "User": "#FF9999",  # Light Red
    "Subject": "#66B2FF",  # Light Blue
    "Grade Level": "#99FF99",  # Light Green
    "Learning Objective": "#FFCC99",  # Light Orange
    "Activity": "#FF99FF",  # Light Purple
    "Assessment": "#FFFF99",  # Light Yellow
    "Resource": "#99FFFF",  # Light Cyan
    "School Board": "#CCCCCC"  # Light Gray
}

def add_to_graph(teacher_name, subject, grade_level, learning_objective, activity, assessment, resource, school_board):
    global lesson_graph
    
    # Clear previous graph
    lesson_graph.clear()
    
    # 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=(14, 10))
    pos = nx.spring_layout(lesson_graph, k=0.9, iterations=50)
    
    # Draw nodes with color coding
    for node, node_type in nx.get_node_attributes(lesson_graph, 'type').items():
        nx.draw_networkx_nodes(lesson_graph, pos, nodelist=[node], node_color=color_map[node_type], 
                               node_size=3000, alpha=0.8)
    
    nx.draw_networkx_edges(lesson_graph, pos, edge_color='gray', arrows=True, arrowsize=20)
    nx.draw_networkx_labels(lesson_graph, pos, font_size=8, font_weight="bold")
    
    # 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=7)
    
    # Create legend
    legend_elements = [mpatches.Patch(color=color, label=node_type) for node_type, color in color_map.items()]
    plt.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1, 1), title="Node Types")
    
    plt.title("Your Educational Landscape", fontsize=16)
    plt.axis('off')
    plt.tight_layout()
    
    # Save the plot to a bytes object
    buf = BytesIO()
    plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
    buf.seek(0)
    plt.close()
    
    # Convert BytesIO to PIL Image
    image = Image.open(buf)
    
    return search_string, image

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

# Gradio interface
demo = gr.Blocks()

with demo:
    gr.Markdown("# EduScape: Design Your Educational Landscape")
    gr.Markdown("Welcome to EduScape, where lesson planning becomes an adventure in crafting educational journeys. Design, visualize, and perfect your learning landscapes 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("Map Your Lesson Plan")
        clear_btn = gr.Button("Clear Landscape")
    
    search_output = gr.Textbox(label="Content Discovery Search String")
    graph_output = gr.Image(label="Your Educational Landscape")
    message_output = gr.Textbox(label="Landscape 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 EduScape app
demo.launch()