from lib import get_data, str_to_ints, add2
from d9 import Amp

DIRS = [(-1, 0), (0, 1), (1, 0), (0, -1)]
BLACK = 0
WHITE = 1


def part_1(data):
    xs = str_to_ints(data)
    a = Amp(xs)
    current_dir_idx = 0
    pos = (0, 0)
    white_panels = set()
    got_painted = set()

    while not a.done:
        if pos in white_panels:
            a.feed(WHITE)
        else:
            a.feed(BLACK)
        a.go()
        a.go()
        while a.outputs:
            paint_color = a.pop()
            if paint_color == BLACK:
                white_panels.remove(pos)
            elif paint_color == WHITE:
                white_panels.add(pos)
            else:
                assert False
            got_painted.add(pos)
            turn_dir_code = a.pop()
            turn_dir = -1 if turn_dir_code == 0 else 1
            current_dir_idx = (current_dir_idx + turn_dir) % len(DIRS)
            pos = add2(pos, DIRS[current_dir_idx])
    print(len(got_painted))


def part_2(data):
    xs = str_to_ints(data)
    a = Amp(xs)
    current_dir_idx = 0
    pos = (0, 0)
    white_panels = set([pos])
    got_painted = set()

    while not a.done:
        if pos in white_panels:
            a.feed(WHITE)
        else:
            a.feed(BLACK)
        a.go()
        a.go()
        while a.outputs:
            paint_color = a.pop()
            if paint_color == BLACK:
                if pos in white_panels:
                    white_panels.remove(pos)
            elif paint_color == WHITE:
                white_panels.add(pos)
            else:
                assert False
            got_painted.add(pos)
            turn_dir_code = a.pop()
            turn_dir = -1 if turn_dir_code == 0 else 1
            current_dir_idx = (current_dir_idx + turn_dir) % len(DIRS)
            pos = add2(pos, DIRS[current_dir_idx])

    white_panels = [(p[1], p[0]) for p in white_panels]
    xs = [p[0] for p in white_panels]
    ys = [p[1] for p in white_panels]
    white_panels_set = set(white_panels)
    x_min, x_max = min(xs), max(xs)
    y_min, y_max = min(ys), max(ys)
    for y in range(y_min, y_max + 1):
        s = ""
        for x in range(x_min, x_max + 1):
            if (x, y) in white_panels_set:
                s += "X"
            else:
                s += " "
        print(s)


def main():
    data = get_data(__file__)
    part_1(data)
    part_2(data)


if __name__ == "__main__":
    main()