diff --git a/malaria.py b/malaria.py
index d773c84..5689753 100644
--- a/malaria.py
+++ b/malaria.py
@@ -9,41 +9,75 @@ class Model:
     def __init__(self, width=32, height=32, humandens=0.15, mosquitodens=0.10,
                  immunepct=0.1, mosqinfpct=0.1, hm_infpct=0.5, mh_infpct=0.5,
                  hinfdiepct=0.01, mhungrypct=0.1, humandiepct=10**-6,
-                 mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=500):
+                 mosqdiepct=10**-3, mosqnetdens=0.05, time_steps=2000):
+        
+        plt.ion()
         self.width = width
         self.height = height
 
+        # The percentage of tiles that start as humans
         self.humandens = humandens
+        # The percentage of tiles that contain mosquitos
         self.mosquitodens = mosquitodens
+        # Percentage of humans that are immune
         self.immunepct = immunepct
+        # Chance for a mosquito to be infected
         self.mosqinfpct = mosqinfpct
+        # Chance for a human to be infected by a mosquito bite
         self.hm_infpct = hm_infpct
+        # Chance for a mosquito to be infected from biting an infected human
         self.mh_infpct = mh_infpct
+        # Chance that an infected human dies
         self.hinfdiepct = hinfdiepct
+        # Chance for a mosquito to be hungry
         self.mhungrypct = mhungrypct
+        # Chance for human to die from random causes
         self.humandiepct = humandiepct
+        # Chance for a mosquito to die
         self.mosqdiepct = mosqdiepct
+        # Percentage of tiles that contain mosquito nets
         self.mosqnetdens = mosqnetdens
-
+        # The number of timesteps to run te simulation for
         self.time_steps = time_steps
 
         self.humans = self.gen_humans()
         self.mosquitos = self.gen_mosquitos()
+        self.init_draw()
+       
+    def init_draw(self):
+        self.colors = matplotlib.colors.ListedColormap(
+            ["black", "green", "red", "yellow"])
+        bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE]
+        self.norm = matplotlib.colors.BoundaryNorm(bounds, self.colors.N)
 
     def recycle_human(self):
         # Get all living humans
-        humans = np.transpose(np.where(self.grid != Human.DEAD))
+        humans = np.transpose(np.where(self.humans != Human.DEAD))
 
         # Get a mask of humans to kill
         deaths = np.random.rand(len(humans)) < self.humandiepct
 
         # Kill them.
-        self.grid[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD
+        self.humans[humans[deaths][:, 0], humans[deaths][:, 1]] = Human.DEAD
 
         # Pick a random, unpopulated spot
-        x, y = random.choice(np.transpose(np.where(self.grid == Human.DEAD)),
+        x, y = random.choice(np.transpose(np.where(self.humans == Human.DEAD)),
                              size=np.count_nonzero(deaths))
 
+
+    def do_malaria(self):
+        """
+        This function determines who of the infected dies from their illness
+        """
+        # Get all infected humans
+        infected = np.transpose(np.where(self.humans == Human.INFECTED))
+
+        # Decide which infected people die
+        deaths = np.random.rand(len(infected)) < self.hinfdiepct
+
+        # Now let's kill them
+        self.humans[infected[deaths][:, 0], infected[deaths][:, 1]] = Human.DEAD
+
     def gen_humans(self):
         # Calculate the probabilities
         p_dead = 1 - self.humandens
@@ -59,22 +93,17 @@ class Model:
         count = self.width * self.height * self.mosquitodens
 
     def run(self):
-        for _ in range(self.time_steps):
+        for t in range(self.time_steps):
             self.step()
-            self.draw()
+            self.draw(t)
 
     def step(self):
-        # dingen
-        pass
+        self.do_malaria()
 
-    def draw(self):
+    def draw(self, t):
         # this function draws the humans
-        colors = matplotlib.colors.ListedColormap(
-            ["black", "green", "red", "yellow"])
-        bounds = [Human.DEAD, Human.HEALTHY, Human.INFECTED, Human.IMMUNE]
-        norm = matplotlib.colors.BoundaryNorm(bounds, colors.N)
-
-        plt.imshow(self.humans, cmap=colors, norm=norm)
+        plt.title("t={}".format(t))
+        plt.imshow(self.humans, cmap=self.colors, norm=self.norm)
 
         plt.pause(0.0001)
         plt.clf()
@@ -97,8 +126,6 @@ class Human(IntEnum):
 
 
 if __name__ == "__main__":
-    plt.ion()
     model = Model()
-    while True:
-        model.draw()
+    model.run()