#!/usr/bin/env python3
from collections import defaultdict
from itertools import zip_longest
from functools import partial

import lib
from lib import vector
from intcode import parse, emulator

BLACK = 0
WHITE = 1

decode_color = {
        WHITE : lib.Color.WHITE,
        BLACK : lib.Color.BLACK
        }

LEFT = 0
RIGHT = 1

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

    def __str__(self):
        xs, ys = zip(*self.keys())
        return '\n'.join(''.join(decode_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")


def paint(program, canvas):
    pos = (0, 0)
    direction = lib.Cardinals.UP

    comp = emulator(program)
    comp.write_input(pop_input=lambda : canvas[pos])

    zipped = zip_longest(comp, comp, fillvalue=LEFT)
    for color, rotate_dir in zipped: 
        canvas[pos] = color 
        rotate = lib.rotate90pos if rotate_dir == LEFT else lib.rotate90neg
        direction = rotate(direction)
        pos = pos + direction

def preproc(puzzle_input):
    program = parse(puzzle_input) 
    return partial(paint, program)

def partI(paint_canvas):
    canvas = Canvas()
    paint_canvas(canvas)
    return len(canvas)

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