import tkinter as tk
from tkinter import simpledialog
import matplotlib.pyplot as plt
import matplotlib.patches as patches

class UMLDiagramDrawer:
    def __init__(self, root):
        self.root = root
        self.root.title("UML Diagram Drawer")
        
        self.canvas = tk.Canvas(root, width=800, height=600, bg="white")
        self.canvas.pack()
        
        self.classes = {}  # Stores class names and their properties
        self.relationships = []  # Stores relationships between classes
        self.selected_class = None  # Tracks the currently selected class
        self.drag_data = {"x": 0, "y": 0}  # Tracks mouse drag data
        
        self.create_menu()
        self.create_buttons()
        self.bind_events()
        
    def create_menu(self):
        menubar = tk.Menu(self.root)
        self.root.config(menu=menubar)
        
        diagram_menu = tk.Menu(menubar, tearoff=0)
        diagram_menu.add_command(label="Add Class", command=self.add_class)
        diagram_menu.add_command(label="Add Relationship", command=self.add_relationship)
        menubar.add_cascade(label="Diagram", menu=diagram_menu)
        
    def create_buttons(self):
        # Add a button to render the diagram with matplotlib
        render_button = tk.Button(self.root, text="Render with Matplotlib", command=self.render_with_matplotlib)
        render_button.pack(side=tk.BOTTOM, pady=10)
        
    def bind_events(self):
        # Bind mouse events for selecting and dragging classes
        self.canvas.bind("<ButtonPress-1>", self.on_mouse_press)
        self.canvas.bind("<B1-Motion>", self.on_mouse_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_mouse_release)
        
    def add_class(self):
        class_name = simpledialog.askstring("Input", "Enter class name:")
        if class_name:
            self.classes[class_name] = {
                "x": 100,  # Default x position
                "y": 100,  # Default y position
                "width": 100,  # Default width
                "height": 60,  # Default height
                "rect": None,  # Canvas rectangle object
                "text": None,  # Canvas text object
            }
            self.draw_class(class_name)
            
    def draw_class(self, class_name):
        props = self.classes[class_name]
        x, y, width, height = props["x"], props["y"], props["width"], props["height"]
        
        # Draw rectangle and text on the canvas
        props["rect"] = self.canvas.create_rectangle(x, y, x + width, y + height, outline="black", fill="lightblue")
        props["text"] = self.canvas.create_text(x + width / 2, y + height / 2, text=class_name, fill="black")
        
    def add_relationship(self):
        class1 = simpledialog.askstring("Input", "Enter first class name:")
        class2 = simpledialog.askstring("Input", "Enter second class name:")
        if class1 in self.classes and class2 in self.classes:
            self.relationships.append((class1, class2))
            self.draw_relationship(class1, class2)
            
    def draw_relationship(self, class1, class2):
        props1 = self.classes[class1]
        props2 = self.classes[class2]
        
        x1 = props1["x"] + props1["width"] / 2
        y1 = props1["y"] + props1["height"]
        x2 = props2["x"] + props2["width"] / 2
        y2 = props2["y"]
        
        self.canvas.create_line(x1, y1, x2, y2, arrow=tk.LAST, fill="black", tags="relationship")
        
    def on_mouse_press(self, event):
        # Check if a class rectangle is clicked
        for class_name, props in self.classes.items():
            x, y, width, height = props["x"], props["y"], props["width"], props["height"]
            if x <= event.x <= x + width and y <= event.y <= y + height:
                self.selected_class = class_name
                self.drag_data["x"] = event.x - x
                self.drag_data["y"] = event.y - y
                break
        
    def on_mouse_drag(self, event):
        if self.selected_class:
            # Move the selected class
            props = self.classes[self.selected_class]
            new_x = event.x - self.drag_data["x"]
            new_y = event.y - self.drag_data["y"]
            
            # Update class position
            props["x"], props["y"] = new_x, new_y
            
            # Move the rectangle and text on the canvas
            self.canvas.coords(props["rect"], new_x, new_y, new_x + props["width"], new_y + props["height"])
            self.canvas.coords(props["text"], new_x + props["width"] / 2, new_y + props["height"] / 2)
            
            # Redraw relationships involving the moved class
            self.redraw_relationships()
        
    def on_mouse_release(self, event):
        self.selected_class = None  # Deselect the class
        
    def redraw_relationships(self):
        # Clear all relationships and redraw them
        self.canvas.delete("relationship")  # Delete all relationship lines
        for class1, class2 in self.relationships:
            self.draw_relationship(class1, class2)
        
    def render_with_matplotlib(self):
        # Create a matplotlib figure
        fig, ax = plt.subplots(figsize=(8, 6))
        
        # Draw classes as rectangles
        for class_name, props in self.classes.items():
            x, y, width, height = props["x"], props["y"], props["width"], props["height"]
            rect = patches.Rectangle((x, y), width, height, linewidth=1, edgecolor='black', facecolor='lightblue')
            ax.add_patch(rect)
            ax.text(x + width / 2, y + height / 2, class_name, ha='center', va='center', color='black')
        
        # Draw relationships as arrows
        for class1, class2 in self.relationships:
            props1 = self.classes[class1]
            props2 = self.classes[class2]
            
            x1 = props1["x"] + props1["width"] / 2
            y1 = props1["y"] + props1["height"]
            x2 = props2["x"] + props2["width"] / 2
            y2 = props2["y"]
            
            ax.annotate("", xy=(x2, y2), xytext=(x1, y1), 
                        arrowprops=dict(arrowstyle="->", color="black"))
        
        # Set plot limits and turn off axes
        ax.set_xlim(0, 800)
        ax.set_ylim(600, 0)  # Invert y-axis to match tkinter coordinates
        ax.set_aspect('equal')
        plt.axis('off')
        
        # Show the plot
        plt.show()

if __name__ == "__main__":
    root = tk.Tk()
    app = UMLDiagramDrawer(root)
    root.mainloop()