from lib import get_data, str_to_ints
from d9 import Amp


def part_1(data):
    xs = str_to_ints(data)
    a = Amp(xs)
    while not a.done:
        a.go()

    r = sum([a.outputs[i : i + 3][2] == 2 for i in range(0, len(a.outputs), 3)])
    print(r)


def part_2(data):
    xs = str_to_ints(data)
    xs[0] = 2  # play for free
    a = Amp(xs)
    ball_x = 0
    paddle_x = 0
    score = 0
    while not a.done:
        a.go()
        if len(a.outputs) == 3:
            x, y, tile_id = a.pop(), a.pop(), a.pop()
            if x == -1 and y == 0:
                score = tile_id
            elif tile_id == 4:
                ball_x = x
            elif tile_id == 3:
                paddle_x = x
        if a.input_required:
            a.input_required = False
            x = 0
            if ball_x < paddle_x:
                x = -1
            elif ball_x > paddle_x:
                x = 1
            a.feed(x)
    print(score)


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


if __name__ == "__main__":
    main()