Implement 2019 day 13 part 2

2019/aoc2019/day13.py

@@ -1,15 +1,10 @@
-from typing import TextIO
+import statistics
+from typing import TextIO, Tuple, Dict
 
 from aoc2019.intcode import Computer, read_program
 
 
-def part1(data: TextIO) -> int:
+def render_screen(computer: Computer, screen: Dict[Tuple[int, int], int]):
-    computer = Computer(read_program(data))
-
-    computer.run()
-
-    screen = {}
-
     while computer.output:
         x = computer.output.popleft()
         y = computer.output.popleft()
@@ -17,4 +12,46 @@ def part1(data: TextIO) -> int:
 
         screen[x, y] = val
 
+
+def part1(data: TextIO) -> int:
+    computer = Computer(read_program(data))
+    computer.run()
+
+    screen = {}
+
+    render_screen(computer, screen)
+
     return sum(1 for val in screen.values() if val == 2)
+
+
+def part2(data: TextIO) -> int:
+    computer = Computer(read_program(data))
+
+    computer.program[0] = 2
+
+    screen = {}
+
+    finished = False
+
+    while not finished:
+        try:
+            computer.run()
+            finished = True
+        except IndexError:
+            # Waiting for input
+            pass
+
+        render_screen(computer, screen)
+
+        ball_x = next(x for x, y in screen if screen[x, y] == 4)
+
+        paddle_x = statistics.mean(x for x, y in screen if screen[x, y] == 3)
+
+        if ball_x < paddle_x:
+            computer.input.append(-1)
+        elif ball_x > paddle_x:
+            computer.input.append(1)
+        else:
+            computer.input.append(0)
+
+    return screen[-1, 0]