Improved UI, added data save function, missing data import function

GamerAI.py

 import pygame
+import os
 import numpy as np
 from random import randint
 from sklearn.neighbors import NearestNeighbors
@@ -17,6 +18,8 @@ clock.tick(60)
 # keyboard pictures and predictor
 leftPic = pygame.image.load('left.png')
 rightPic = pygame.image.load('right.png')
+downPic = pygame.image.load('down.png')
+upPic = pygame.image.load('up.png')
 predPic = pygame.image.load('predictor.png')
 
 # make the pygame window
@@ -30,7 +33,7 @@ smallfont = pygame.font.SysFont('comicsansms', 25)
 mediumfont = pygame.font.SysFont('comicsansms', 50)
 largefont = pygame.font.SysFont('comicsansms', 80)
 
-# matrixes of color, index 0 should be deleted
+# matrixes of color, index 0 should be deleted, used if trained anew
 points = np.array([[0,0,0]])
 white = np.array([[0,0,0]])
 black = np.array([[0,0,0]])
@@ -57,25 +60,31 @@ while True:
     display.fill((123, 123, 123))
         
     # draw color boxes
-    pygame.draw.rect(display, color, [(width - 200) / 4, 300, 200, 200])
-    pygame.draw.rect(display, color, [(width - 200) / 1.25, 300, 200, 200])
+    pygame.draw.rect(display, color, [width / 3 - 200 / 2, 300, 200, 200])
+    pygame.draw.rect(display, color, [width / 1.5 - 200 / 2, 300, 200, 200])
+    pygame.draw.rect(display, (37,37,37), [width / 3 - 220 / 2 - 4, 75, 220, 40])
+    pygame.draw.rect(display, (37,37,37), [width / 1.5 - 245 / 2 - 4, 75, 245, 40])
+
 
     # blit pics of arrows
-    display.blit(leftPic, ((width - 200) / 4 + 70, 512))
-    display.blit(rightPic, ((width - 200) / 1.25 + 70, 512))
+    display.blit(leftPic, (width / 3 - 63 / 2, 512))
+    display.blit(rightPic, (width / 1.5 - 63 / 2, 512))
+    display.blit(downPic, ((width / 3 - 63 / 2), 125))
+    display.blit(upPic, ((width / 1.5 - 63 / 2), 125))
 
     # text
-    text('text', (255,255,255), largefont, (width - 200) / 4 + 20, 330)
-    text('text', (0,0,0), largefont, (width - 200) / 1.25 + 20, 330)
-    text('bazinga', (255,2,255), largefont, 75, 25)
+    text('text', (255,255,255), largefont, width / 3 - 200 / 2 + 20, 330)
+    text('text', (0,0,0), largefont, width / 1.5 - 200 / 2 + 20, 330)
+    text('save training data', (255,255,255), smallfont, width / 3 - 216 / 2, 75)
+    text('import training data', (255,255,255), smallfont, width / 1.5 - 241 / 2, 75)
 
     # make prediction, display prediction above color box
     y = np.reshape(color, (-1, 3))
     if len(points) > 1:
         if nn(points, y) in white:
-            display.blit(predPic, ((width - 200) / 4 + 60, 200))
+            display.blit(predPic, (width / 3 - 80 / 2, 200))
         elif nn(points, y) in black:
-            display.blit(predPic, ((width - 200) / 1.25 + 60, 200))
+            display.blit(predPic, (width / 1.5 - 80 / 2, 200))
 
     # remove [0,0,0]
     if len(points) == 2 and np.array_equal(points[0], np.array([0,0,0])):
@@ -103,3 +112,18 @@ while True:
                     black = np.vstack((black, color))
                     points = np.vstack((points, color))
                     loop = False
+                if event.key == pygame.K_DOWN:
+                    try:
+                        os.remove("training.txt")
+                    except:
+                        pass
+                    dataFile = open("training.txt", "w")
+                    dataFile.write(str(points) + "\n")
+                    dataFile.write(str(white) + "\n")
+                    dataFile.write((str(black) + "\n"))
+                    dataFile.close()
+                if event.key == pygame.K_UP:
+                    try:
+                        pass
+                    except:
+                        pass
\ No newline at end of file