advent-of-code-2019/day11.py

#!/usr/bin/env python3

import intcode
from collections import defaultdict
from itertools import cycle

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 Robot:
    def __init__(self, program):
        self.pos = (0, 0)
        self.dir = UP

        self.comp = intcode.emulator(program)

    def paint(self, canvas):
        self.comp.write_input(pop_input=lambda : canvas[self.pos])
        paint_flags = cycle((True, False))

        for pf, x in zip(paint_flags, self.comp):
            if pf:
                canvas[self.pos] = x
            else:
                self.dir = left(self.dir) if x == LEFT else right(self.dir)
                self.pos = add(self.pos, self.dir)

def preproc(puzzle_input):
    program = intcode.parse(puzzle_input) 
    paint = lambda c : Robot(program).paint(c)
    return paint

def partI(paint):
    canvas = Canvas(color)
    paint(canvas)
    return len(canvas)

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