import tkinter as tk
from tkinter import messagebox, simpledialog
import math

class Person:
    def __init__(self, name="New Person", x=0, y=0):
        self.name = name
        self.x = x
        self.y = y
        self.spouse = None
        self.children = []
        self.parents = []
        self.rect_id = None
        self.text_id = None
        
    def add_child(self, child):
        if child not in self.children:
            self.children.append(child)
            child.parents.append(self)
            
    def add_parent(self, parent):
        if parent not in self.parents:
            self.parents.append(parent)
            parent.children.append(self)
            
    def set_spouse(self, spouse):
        if self.spouse:
            self.spouse.spouse = None
        self.spouse = spouse
        if spouse:
            spouse.spouse = self
            
    def remove_from_family(self):
        # Remove spouse relationship
        if self.spouse:
            self.spouse.spouse = None
            
        # Remove from parents' children lists
        for parent in self.parents:
            if self in parent.children:
                parent.children.remove(self)
                
        # Remove from children's parents lists
        for child in self.children:
            if self in child.parents:
                child.parents.remove(self)

class FamilyTreeApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Family Tree Diagram")
        self.root.geometry("1000x700")
        
        # Create canvas with scrollbars
        self.create_canvas()
        
        # Data structures
        self.people = []
        self.selected_person = None
        
        # Layout parameters
        self.person_width = 120
        self.person_height = 60
        self.level_height = 120
        self.person_spacing = 150
        
        # Create initial person
        initial_person = Person("Click to Edit", 400, 300)
        self.people.append(initial_person)
        
        self.draw_tree()
        
    def create_canvas(self):
        # Main frame
        main_frame = tk.Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # Canvas with scrollbars
        self.canvas = tk.Canvas(main_frame, bg='white', scrollregion=(0, 0, 2000, 2000))
        
        v_scrollbar = tk.Scrollbar(main_frame, orient=tk.VERTICAL, command=self.canvas.yview)
        h_scrollbar = tk.Scrollbar(main_frame, orient=tk.HORIZONTAL, command=self.canvas.xview)
        
        self.canvas.configure(yscrollcommand=v_scrollbar.set, xscrollcommand=h_scrollbar.set)
        
        v_scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
        h_scrollbar.pack(side=tk.BOTTOM, fill=tk.X)
        self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        # Bind events
        self.canvas.bind("<Button-1>", self.on_click)
        self.canvas.bind("<Button-3>", self.on_right_click)  # Right click for context menu
        
    def calculate_layout(self):
        """Calculate positions for all people in the tree"""
        if not self.people:
            return
            
        # Find root people (those with no parents)
        roots = [p for p in self.people if not p.parents]
        if not roots:
            # If no roots found, pick the first person
            roots = [self.people[0]]
            
        # Calculate levels
        levels = {}
        self.assign_levels(roots, levels, 0)
        
        # Group people by level
        level_groups = {}
        for person, level in levels.items():
            if level not in level_groups:
                level_groups[level] = []
            level_groups[level].append(person)
            
        # Calculate positions
        canvas_width = 2000
        for level, people_in_level in level_groups.items():
            y = 100 + level * self.level_height
            
            # Sort people in level to keep families together
            people_in_level.sort(key=lambda p: self.get_family_sort_key(p))
            
            total_width = len(people_in_level) * self.person_spacing
            start_x = (canvas_width - total_width) // 2
            
            for i, person in enumerate(people_in_level):
                person.x = start_x + i * self.person_spacing
                person.y = y
                
    def assign_levels(self, people, levels, level):
        """Recursively assign levels to people"""
        for person in people:
            if person not in levels:
                levels[person] = level
                # Process children at next level
                if person.children:
                    self.assign_levels(person.children, levels, level + 1)
                    
    def get_family_sort_key(self, person):
        """Generate a sort key to keep families together"""
        # Use parent names as primary sort key
        if person.parents:
            return ''.join(sorted([p.name for p in person.parents]))
        return person.name
        
    def draw_tree(self):
        """Draw the entire family tree"""
        self.canvas.delete("all")
        
        # Calculate layout
        self.calculate_layout()
        
        # Draw connections first (so they appear behind rectangles)
        self.draw_connections()
        
        # Draw people
        for person in self.people:
            self.draw_person(person)
            
    def draw_person(self, person):
        """Draw a single person as a rectangle with text"""
        x1 = person.x - self.person_width // 2
        y1 = person.y - self.person_height // 2
        x2 = person.x + self.person_width // 2
        y2 = person.y + self.person_height // 2
        
        # Choose color based on selection
        color = "lightblue" if person == self.selected_person else "lightgray"
        
        # Draw rectangle
        person.rect_id = self.canvas.create_rectangle(
            x1, y1, x2, y2, 
            fill=color, outline="black", width=2,
            tags=f"person_{id(person)}"
        )
        
        # Draw text
        person.text_id = self.canvas.create_text(
            person.x, person.y, 
            text=person.name, 
            font=("Arial", 10, "bold"),
            width=self.person_width - 10,
            tags=f"person_{id(person)}"
        )
        
    def draw_connections(self):
        """Draw family relationship lines"""
        for person in self.people:
            # Draw spouse connection
            if person.spouse and id(person) < id(person.spouse):  # Draw only once per couple
                self.draw_spouse_line(person, person.spouse)
                
            # Draw parent-child connections
            if person.children:
                self.draw_parent_child_lines(person)
                
    def draw_spouse_line(self, person1, person2):
        """Draw line between spouses"""
        self.canvas.create_line(
            person1.x, person1.y,
            person2.x, person2.y,
            fill="red", width=3, tags="connection"
        )
        
    def draw_parent_child_lines(self, parent):
        """Draw lines from parent to children"""
        if not parent.children:
            return
            
        # Find the spouse to determine the family unit center
        spouse = parent.spouse
        if spouse:
            family_x = (parent.x + spouse.x) // 2
            family_y = min(parent.y, spouse.y)
        else:
            family_x = parent.x
            family_y = parent.y
            
        # Draw vertical line down from family unit
        children_y = min([child.y for child in parent.children])
        mid_y = family_y + (children_y - family_y) // 2
        
        self.canvas.create_line(
            family_x, family_y + self.person_height // 2,
            family_x, mid_y,
            fill="blue", width=2, tags="connection"
        )
        
        # Draw horizontal line across children
        if len(parent.children) > 1:
            left_x = min([child.x for child in parent.children])
            right_x = max([child.x for child in parent.children])
            
            self.canvas.create_line(
                left_x, mid_y,
                right_x, mid_y,
                fill="blue", width=2, tags="connection"
            )
            
        # Draw vertical lines to each child
        for child in parent.children:
            self.canvas.create_line(
                child.x, mid_y,
                child.x, child.y - self.person_height // 2,
                fill="blue", width=2, tags="connection"
            )
            
    def on_click(self, event):
        """Handle left mouse click"""
        clicked_item = self.canvas.find_closest(event.x, event.y)[0]
        
        # Find which person was clicked
        clicked_person = None
        for person in self.people:
            if (person.rect_id == clicked_item or person.text_id == clicked_item):
                clicked_person = person
                break
                
        if clicked_person:
            self.selected_person = clicked_person
            self.show_edit_menu(clicked_person, event.x, event.y)
        else:
            self.selected_person = None
            
        self.draw_tree()
        
    def on_right_click(self, event):
        """Handle right mouse click - same as left click for simplicity"""
        self.on_click(event)
        
    def show_edit_menu(self, person, x, y):
        """Show context menu for editing person"""
        menu = tk.Menu(self.root, tearoff=0)
        menu.add_command(label="Edit Name", command=lambda: self.edit_name(person))
        menu.add_command(label="Add Spouse", command=lambda: self.add_spouse(person))
        menu.add_command(label="Add Child", command=lambda: self.add_child(person))
        menu.add_command(label="Add Parent", command=lambda: self.add_parent(person))
        menu.add_separator()
        menu.add_command(label="Delete Person", command=lambda: self.delete_person(person))
        
        try:
            menu.tk_popup(event.x_root, event.y_root)
        finally:
            menu.grab_release()
            
    def edit_name(self, person):
        """Edit person's name"""
        new_name = simpledialog.askstring("Edit Name", "Enter new name:", initialvalue=person.name)
        if new_name:
            person.name = new_name
            self.draw_tree()
            
    def add_spouse(self, person):
        """Add spouse to person"""
        if person.spouse:
            messagebox.showwarning("Warning", f"{person.name} already has a spouse!")
            return
            
        spouse_name = simpledialog.askstring("Add Spouse", "Enter spouse's name:")
        if spouse_name:
            spouse = Person(spouse_name, person.x + self.person_spacing, person.y)
            person.set_spouse(spouse)
            self.people.append(spouse)
            self.draw_tree()
            
    def add_child(self, person):
        """Add child to person"""
        child_name = simpledialog.askstring("Add Child", "Enter child's name:")
        if child_name:
            child = Person(child_name, person.x, person.y + self.level_height)
            person.add_child(child)
            
            # If person has spouse, add child to spouse too
            if person.spouse:
                person.spouse.add_child(child)
                
            self.people.append(child)
            self.draw_tree()
            
    def add_parent(self, person):
        """Add parent to person"""
        parent_name = simpledialog.askstring("Add Parent", "Enter parent's name:")
        if parent_name:
            parent = Person(parent_name, person.x, person.y - self.level_height)
            person.add_parent(parent)
            self.people.append(parent)
            self.draw_tree()
            
    def delete_person(self, person):
        """Delete person from tree"""
        if len(self.people) == 1:
            messagebox.showwarning("Warning", "Cannot delete the last person!")
            return
            
        result = messagebox.askyesno("Confirm Delete", 
                                   f"Are you sure you want to delete {person.name}?")
        if result:
            person.remove_from_family()
            self.people.remove(person)
            if self.selected_person == person:
                self.selected_person = None
            self.draw_tree()

def main():
    root = tk.Tk()
    app = FamilyTreeApp(root)
    root.mainloop()

if __name__ == "__main__":
    main()