lesson_graph.py

import networkx as nx
import matplotlib.pyplot as plt
from io import BytesIO
from PIL import Image
import matplotlib.patches as mpatches
import mplcursors
import json
from typing import Dict, Tuple, Any, NamedTuple, Optional
from enum import Enum

class NodeType(Enum):
    USER = "User"
    SUBJECT = "Subject"
    GRADE_LEVEL = "Grade Level"
    LEARNING_OBJECTIVE = "Learning Objective"
    ACTIVITY = "Activity"
    ASSESSMENT = "Assessment"
    RESOURCE = "Resource"
    SCHOOL_BOARD = "School Board"
    COUNTRY_AUTHORITY = "Country Authority"  # New NodeType


class ResetState(NamedTuple):
    teacher_name: str
    subject: str
    grade_level: str
    learning_objective: str
    activity: str
    assessment: str
    resource: str
    school_board: str
    country_authority: str  # New field
    message: str

class LessonGraph:
    INITIAL_STATE: Dict[str, str] = {
        "teacher_name": "",
        "subject": "",
        "grade_level": "",
        "learning_objective": "",
        "activity": "",
        "assessment": "",
        "resource": "",
        "school_board": "",
        "country_authority": ""  # New field
    }

    REQUIRED_FIELDS = ["teacher_name", "subject", "grade_level"]

    COLOR_MAP: Dict[NodeType, str] = {
        NodeType.USER: "#FF9999",
        NodeType.SUBJECT: "#66B2FF",
        NodeType.GRADE_LEVEL: "#99FF99",
        NodeType.LEARNING_OBJECTIVE: "#FFCC99",
        NodeType.ACTIVITY: "#FF99FF",
        NodeType.ASSESSMENT: "#FFFF99",
        NodeType.RESOURCE: "#99FFFF",
        NodeType.SCHOOL_BOARD: "#CCCCCC",
        NodeType.COUNTRY_AUTHORITY: "#FFA07A"  # New color for Country Authority
    }

    def __init__(self):
        self.graph = nx.DiGraph()
        self.inputs = self.INITIAL_STATE.copy()

    def validate_required_fields(self):
        """
        Validate that all required fields are filled.
        Raises a ValueError if any required field is empty.
        """
        missing_fields = [field for field in self.REQUIRED_FIELDS if not self.inputs.get(field)]
        if missing_fields:
            raise ValueError(f"The following required fields are missing: {', '.join(missing_fields)}")

    def add_lesson_plan(self, **kwargs) -> Tuple[str, Image.Image]:
        """
        Add nodes and edges to the lesson plan graph for the given inputs.
        Returns a search string and the graph image.
        """
        self.graph.clear()
        self.inputs.update(kwargs)
        self.validate_required_fields()
    
        # Define required nodes
        nodes = {
            self.inputs["teacher_name"]: {"type": NodeType.USER, "role": "Teacher"},
            self.inputs["subject"]: {"type": NodeType.SUBJECT, "description": "Core subject area"},
            self.inputs["grade_level"]: {"type": NodeType.GRADE_LEVEL, "description": "Target grade for the lesson"},
        }
    
        # Include country authority if provided
        if self.inputs.get("country_authority"):
            nodes[self.inputs["country_authority"]] = {
                "type": NodeType.COUNTRY_AUTHORITY,
                "description": "Sets national curriculum standards"
            }
    
        # Optional nodes
        optional_nodes = {
            "school_board": NodeType.SCHOOL_BOARD,
            "learning_objective": NodeType.LEARNING_OBJECTIVE,
            "activity": NodeType.ACTIVITY,
            "assessment": NodeType.ASSESSMENT,
            "resource": NodeType.RESOURCE
        }
    
        for field, node_type in optional_nodes.items():
            if self.inputs.get(field):
                nodes[self.inputs[field]] = {"type": node_type, "description": f"{node_type.value}"}
    
        # Add nodes to the graph
        for node, attributes in nodes.items():
            self.graph.add_node(node, **attributes)
    
        # Define relationships between nodes
        edges = [
            (self.inputs["teacher_name"], self.inputs["subject"], {"relationship": "TEACHES"}),
            (self.inputs["subject"], self.inputs["grade_level"], {"relationship": "HAS_GRADE"})
        ]
    
        # Relationships involving country authority
        if self.inputs.get("country_authority"):
            if self.inputs.get("learning_objective"):
                edges.append((self.inputs["country_authority"], self.inputs["learning_objective"], {"relationship": "DEFINES"}))
            if self.inputs.get("school_board"):
                edges.append((self.inputs["country_authority"], self.inputs["school_board"], {"relationship": "OVERSEES"}))
    
        # Existing optional edges
        if self.inputs.get("learning_objective"):
            edges.append((self.inputs["subject"], self.inputs["learning_objective"], {"relationship": "COVERS"}))
            if self.inputs.get("school_board"):
                edges.append((self.inputs["learning_objective"], self.inputs["school_board"], {"relationship": "ALIGNS_WITH"}))
    
        if self.inputs.get("activity") and self.inputs.get("learning_objective"):
            edges.append((self.inputs["activity"], self.inputs["learning_objective"], {"relationship": "ACHIEVES"}))
    
        if self.inputs.get("activity") and self.inputs.get("resource"):
            edges.append((self.inputs["activity"], self.inputs["resource"], {"relationship": "REQUIRES"}))
    
        if self.inputs.get("learning_objective") and self.inputs.get("assessment"):
            edges.append((self.inputs["learning_objective"], self.inputs["assessment"], {"relationship": "EVALUATED_BY"}))
    
        if self.inputs.get("school_board"):
            edges.append((self.inputs["teacher_name"], self.inputs["school_board"], {"relationship": "BELONGS_TO"}))
    
        # Remove None entries from edges list
        edges = [edge for edge in edges if edge is not None]
        self.graph.add_edges_from(edges)
    
        # Generate the search string for content discovery
        search_string = f"{self.inputs['subject']} {self.inputs['grade_level']} {self.inputs.get('learning_objective', '')} {self.inputs.get('activity', '')} {self.inputs.get('resource', '')}".strip()
    
        # Get the graph image
        image = self.draw_graph()
    
        return search_string, image

    def draw_graph(self) -> Image.Image:
        """
        Visualize the graph using Matplotlib, handling layout, labels, and interactivity.
        """
        fig, ax = plt.subplots(figsize=(14, 10))
        pos = nx.spring_layout(self.graph, k=1.2, iterations=100)

        self._draw_nodes(ax, pos)
        self._draw_edges(ax, pos)
        self._add_legend(ax)
        
        plt.title("Your Educational Landscape", fontsize=16)
        plt.axis('off')
        plt.tight_layout()
        
        self._add_interactivity()
        
        # Save the plot to a BytesIO object
        buf = BytesIO()
        plt.savefig(buf, format="png", dpi=300, bbox_inches="tight", pad_inches=0.5)
        buf.seek(0)
        plt.close(fig)

        return Image.open(buf)

    def _draw_nodes(self, ax, pos):
        node_colors = [self.COLOR_MAP[self.graph.nodes[node]['type']] for node in self.graph.nodes()]
        nx.draw_networkx_nodes(self.graph, pos, node_color=node_colors, node_size=3000, alpha=0.8, ax=ax)
        nx.draw_networkx_labels(self.graph, pos, font_size=10, font_weight="bold", ax=ax)

    def _draw_edges(self, ax, pos):
        nx.draw_networkx_edges(self.graph, pos, edge_color='gray', arrows=True, arrowsize=20, ax=ax)
        edge_labels = nx.get_edge_attributes(self.graph, 'relationship')
        nx.draw_networkx_edge_labels(self.graph, pos, edge_labels=edge_labels, font_size=8, ax=ax)
    
    def _add_legend(self, ax):
        legend_elements = [mpatches.Patch(color=color, label=node_type.value) for node_type, color in self.COLOR_MAP.items()]
        ax.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1, 1), title="Node Types")

    def _add_interactivity(self):
        cursor = mplcursors.cursor(hover=True)
        @cursor.connect("add")
        def on_add(sel):
            node = list(self.graph.nodes())[sel.target.index]
            node_data = self.graph.nodes[node]
            sel.annotation.set_text(f"Node: {node}\nType: {node_data['type'].value}\n{node_data.get('description', '')}")

    def reset_state(self) -> ResetState:
        """
        Resets all input states to their default values and clears the graph.
        Returns a named tuple of the cleared input values and a status message.
        """
        self.inputs = self.INITIAL_STATE.copy()
        self.graph.clear()
        return ResetState(**self.inputs, message="Landscape cleared. You can start a new lesson plan.")

    def graph_to_json(self) -> str:
        """
        Converts the current lesson plan graph into a JSON string format and returns the result.
        """
        try:
            graph_data = {
                "nodes": [
                    {
                        "id": node,
                        "type": self.graph.nodes[node]["type"].value,
                        "description": self.graph.nodes[node].get("description", "")
                    }
                    for node in self.graph.nodes()
                ],
                "edges": [
                    {
                        "source": u,
                        "target": v,
                        "relationship": self.graph.edges[u, v]["relationship"]
                    }
                    for u, v in self.graph.edges()
                ]
            }
            return json.dumps(graph_data, indent=4)
        except (KeyError, TypeError) as e:
            return f"An error occurred while converting the graph to JSON: {str(e)}"

    def process_inputs(self, *args) -> Tuple[str, Optional[Image.Image]]:
        """
        Process input arguments and create a lesson plan.
        Returns a tuple of search string and graph image, or error message and None.
        """
        try:
            self.inputs.update(dict(zip(self.INITIAL_STATE.keys(), args)))
            return self.add_lesson_plan(**self.inputs)
        except ValueError as e:
            return str(e), None

    @property
    def is_empty(self) -> bool:
        """Check if all inputs are empty."""
        return all(value == "" for value in self.inputs.values())

    def __repr__(self) -> str:
        return f"LessonGraph(inputs={self.inputs}, graph_size={len(self.graph)})"