diff --git a/network_attack_propagation.ipynb b/network_attack_propagation.ipynb
index 9aede97..6782c2a 100644
--- a/network_attack_propagation.ipynb
+++ b/network_attack_propagation.ipynb
@@ -95,4 +95,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}
\ No newline at end of file
diff --git a/network_attack_propagation.py b/network_attack_propagation.py
index d49d59f..13bf9fe 100644
--- a/network_attack_propagation.py
+++ b/network_attack_propagation.py
@@ -35,6 +35,7 @@ class Graph:
     def __init__(self, use_weights=False):
         self.edges = []
         self.use_weights = use_weights
+        self.rounds_survived = 0
 
     def add_edge(self, edge: Edge):
         self.edges.append(edge)
@@ -64,6 +65,16 @@ class Graph:
 
         return nodes
 
+    def is_alive(self):
+        nodes_alive = list(filter(lambda x: not x.is_infected, self.get_nodes()))
+        return len(nodes_alive) > 0
+
+    def update_survived(self):
+        if not self.is_alive():
+            return
+
+        self.rounds_survived += 1
+
     def infect_step(self):
         infected_nodes = list(filter(lambda n: n.is_infected, self.get_nodes()))
         for node in infected_nodes:
@@ -73,16 +84,21 @@ class Graph:
             else:
                 to_be_infected = random.choice([n[0] for n in adjacent_nodes])
             to_be_infected.is_infected = True
+        self.update_survived()
 
 
 def update(num, layout, g_repr, ax, our_graph: Graph):
     """
     This function is called every 'step', so if you wish to update the graph, do it here
     """
+    if not our_graph.is_alive():
+        return
+
+    our_graph.infect_step()
     ax.clear()
 
     colors = ['red' if n.is_infected else 'blue' for n in g_repr]
-    sizes = [50 if n.is_infected else 1 for n in g_repr]
+    sizes = [200 if n.is_infected else 50 for n in g_repr]
     nx.draw(
         g_repr,
         ax=ax,
@@ -90,13 +106,9 @@ def update(num, layout, g_repr, ax, our_graph: Graph):
         node_color=colors,
         with_labels=False,
         node_size=sizes,
-        node_shape="s",
-        alpha=0.5,
-        linewidths=40,
+        alpha=0.7,
     )
 
-    our_graph.infect_step()
-
 
 def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int, layout):
     g_repr = nx.Graph()
@@ -106,8 +118,8 @@ def do_graph_animation(output_file_name: str, in_graph: Graph, frame_count: int,
     layout = layout(g_repr)
     fig, ax = plt.subplots()
 
-    fig.set_figwidth(15)
-    fig.set_figheight(15)
+    fig.set_figwidth(8)
+    fig.set_figheight(8)
 
     anim = animation.FuncAnimation(
         fig, update, frames=frame_count, interval=500, fargs=(layout, g_repr, ax, in_graph)
@@ -165,31 +177,19 @@ def ring_network(n=30) -> Graph:
 
 
 def main():
-    network = Graph()
-    nodes = [Node(True), Node(), Node(), Node(True), Node()]
+    bus = bus_network(36)
+    do_graph_animation('bus.gif', bus, 90, nx.spiral_layout)
 
-    network.add_edges(
-        [
-            Edge(nodes[1], nodes[0], 0.02),
-            Edge(nodes[1], nodes[2], 0.2),
-            Edge(nodes[2], nodes[0], 0.7),
-            Edge(nodes[3], nodes[2], 0.2),
-            Edge(nodes[3], nodes[1], 0.2),
-            Edge(nodes[4], nodes[3], 0.2),
-        ]
-    )
-
-    do_graph_animation('test.gif', network, 5, nx.spring_layout)
-
-    bus = bus_network()
-    do_graph_animation('bus.gif', bus, 20, nx.spiral_layout)
-
-    ring = ring_network()
-    do_graph_animation('ring.gif', ring, 20, nx.circular_layout)
+    ring = ring_network(36)
+    do_graph_animation('ring.gif', ring, 90, nx.circular_layout)
 
     star, star_avg_rank = star_network()
     list(star.get_nodes())[1].is_infected = True
-    do_graph_animation('star.gif', star, 20, nx.kamada_kawai_layout)
+    do_graph_animation('star.gif', star, 90, nx.kamada_kawai_layout)
+
+    print(f'bus survived {bus.rounds_survived} rounds')
+    print(f'ring survived {ring.rounds_survived} rounds')
+    print(f'star survived {star.rounds_survived} rounds')
 
 
 if __name__ == "__main__":