From c09db757bcb49b1945000402fee0c7e922ca40ed Mon Sep 17 00:00:00 2001
From: Toby Lane <tobylane1@gmail.com>
Date: Sun, 30 Jun 2024 23:24:07 +0200
Subject: [PATCH] added edge selection

---
 src/editor.py      | 13 ++++++++++++-
 src/graph/graph.py | 11 ++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/editor.py b/src/editor.py
index 234d9da..fc75b99 100644
--- a/src/editor.py
+++ b/src/editor.py
@@ -47,6 +47,7 @@ class Editor:
         self.node_radius: int = 10
         self.line_size: int = (int) (self.node_radius / 5)
         self.selected_node: int = -1
+        self.selected_edge: int = -1
 
     def mainloop(self) -> None:
         self.state = State.LOADING
@@ -256,6 +257,9 @@ class Editor:
         for edge in self.graph.edges:
             node_1, node_2 = self.graph.get_edge_nodes(edge)
             pygame.draw.line(self.win, (255, 0, 0), self.world_to_screen(node_1.x, node_1.z), self.world_to_screen(node_2.x, node_2.z), self.line_size)
+        if self.selected_edge != -1:
+            node_1, node_2 = self.graph.get_edge_nodes(self.graph.edges[self.selected_edge])
+            pygame.draw.line(self.win, (0, 255, 255), self.world_to_screen(node_1.x, node_1.z), self.world_to_screen(node_2.x, node_2.z), self.line_size)
     
     def render_nodes(self) -> None:
         for node in self.graph.nodes:
@@ -303,10 +307,17 @@ class Editor:
         return int(screen_x), int(screen_y)
     
     def select_node(self, node: int):
+        self.selected_edge = -1
         if self.selected_node != -1 and node != -1:
-            self.create_edge(self.selected_node, node)
+            self.link_nodes(self.selected_node, node)
         self.selected_node = node
     
+    def link_nodes(self, node_1: int, node_2: int):
+        if self.graph.edge_exists(node_1, node_2):
+            self.selected_edge = self.graph.get_edge(node_1, node_2)
+        else:
+            self.create_edge(node_1, node_2)
+    
     def create_node(self) -> None:
         self.graph.add_node(self.node_candidate_pos[0], self.node_candidate_pos[1], self.typing_text)
         self.typing_text = ""
diff --git a/src/graph/graph.py b/src/graph/graph.py
index 0844f80..cfeb64f 100644
--- a/src/graph/graph.py
+++ b/src/graph/graph.py
@@ -18,11 +18,20 @@ class Graph:
     def number_of_nodes(self) -> int:
         return len(self.nodes)
     
-    def get_edge_nodes(self, edge) -> tuple[Node, Node]:
+    def get_edge(self, node_1: int, node_2: int) -> int:
+        for edge in self.edges:
+            if (edge.start == node_1 and edge.end == node_2) or (edge.start == node_2 and edge.end == node_1):
+                return self.edges.index(edge)
+        return -1
+    
+    def get_edge_nodes(self, edge: Edge) -> tuple[Node, Node]:
         return self.nodes[edge.start], self.nodes[edge.end]
 
     def edges_adjacent_to(self, node_i: int) -> Iterator[Edge]:
         return filter(lambda e: e.start == node_i or e.end == node_i, self.edges)
+    
+    def edge_exists(self, node_1: int, node_2: int) -> bool:
+        return self.get_edge(node_1, node_2) != -1
 
     def dijkstra(self, source_index: int, target_index: int) -> Optional[list[int]]:
         n = len(self.nodes)