Update agent.py, now with arrows!

agent.py

@@ -66,6 +66,7 @@ def checkBoundary(agent, winWidth, winHeight):
     elif y <= 0 or y >= winHeight:
         agent[2] = agent[2] * (-1)
         agent[0].move(agent[1],agent[2])
+    return agent
 
 def main():
     winWidth = 1000
@@ -75,8 +76,8 @@ def main():
 
     maxTime = 4000
     maxV = 8
-    agentNum = 50
-    agents = [[0 for x in range(3)] for y in range(agentNum)]
+    agentNum = 75
+    agents = [[0 for x in range(4)] for y in range(agentNum)]
 
     #generate point
     for agent in agents:
@@ -84,7 +85,11 @@ def main():
 
         agent[1] = random.uniform(-2,2)
         agent[2] = random.uniform(-2,2)
+
         agent[0].draw(window)
+        agent[3] = Line(agent[0], Point(agent[0].getX() + agent[1], agent[0].getY() + agent[2]))
+        agent[3].setArrow("last")
+        agent[3].draw(window)
 
     #update points
     for i in range(maxTime):
@@ -94,8 +99,13 @@ def main():
             agent = updateV(agent, nn, maxV)
             agent = checkBoundary(agent, winWidth, winHeight)
 
+            agent[3].undraw()
+            agent[3] = Line(agent[0], Point(agent[0].getX() + agent[1], agent[0].getY() + agent[2]))
+            agent[3].setArrow("last")
+            agent[3].draw(window)
+
         time.sleep(0.01)
-        
+
     window.getMouse()
     window.close()