advent-of-code-2019/day11.py

#!/usr/bin/env python3

from intcode import Interpreter
from collections import defaultdict

BLACK = 0
WHITE = 1

color = {WHITE : u"\u2B1C",
         BLACK : u"\u2B1B"
        }

UP    = (0, 1)
LEFT  = 0
RIGHT = 1

def rotate(v, dir):
    return (-dir) * v[1], dir * v[0]

def left(v):
    return rotate(v, 1)

def right(v):
    return rotate(v, -1)

def add(v1, v2):
    return v1[0] + v2[0], v1[1] + v2[1]

class Canvas(defaultdict):
    def __init__(self, color):
        self.color = color
        super(Canvas, self).__init__(int)

    def __str__(self):
        xs, ys = zip(*self.keys())
        return '\n'.join(''.join(self.color[self[(x, y)]]
                         for x in range(min(xs), max(xs)+1))
                            for y in range(max(ys), min(ys)-1, -1))

    def __repr__(self):
        return unicode(self).encode("utf-8")


class robotIO:
    def __init__(self, canvas=Canvas(color)):
        self.canvas = canvas
        self.dir    = UP
        self.pos    = 0, 0
        self.paint  = True

    def pop_input(self):
        return self.canvas[self.pos]

    def append_output(self, x):
        if self.paint:
            self.canvas[self.pos] = x
        else:
            self.dir = left(self.dir) if x == LEFT else right(self.dir)
            self.pos = add(self.pos, self.dir)
        self.paint = not self.paint

    def flush(self):
        return self.canvas

class Robot:
    def __init__(self, program, canvas=Canvas(color)):
        self.canvas = canvas
        self.program = program

    def paint(self):
        comp = Interpreter(self.program, robotIO)
        comp.IO.canvas = self.canvas
        return comp.run()

def preproc(puzzle_input):
    return list(map(int, puzzle_input.split(',')))

def partI(program):
    canvas = Robot(program).paint()
    return len(canvas)

def partII(program):
    canvas = Canvas(color)
    canvas[(0, 0)] = WHITE
    canvas = Robot(program, canvas).paint()
    return '\n' + str(canvas)