from utils import open_day, Point2D as P
from itertools import repeat, cycle

pieces = [
    { P(0, 0), P(1, 0), P(2, 0), P(3, 0) },
    { P(1, 0), P(0, 1), P(1, 1), P(2, 1), P(1, 2) },
    { P(0, 0), P(1, 0), P(2, 0), P(2, 1), P(2, 2) },
    { P(0, 0), P(0, 1), P(0, 2), P(0, 3) },
    { P(0, 0), P(1, 0), P(0, 1), P(1, 1) },
]

class Tower:
    def __init__(self, width=7):
        self.width = width
        self.tower = []
    @property
    def height(self):
        return len(self.tower) // self.width
    def __str__(self):
        rows = []
        for y in range(self.height - 1, -1, -1):
            rows.append('|' + ''.join(
                    '#' if self.tower[y * self.width + x] else '.'
                    for x in range(self.width)) + '|')
        return '\n'.join(rows + ['+' + '-' * self.width + '+'])
    def __setitem__(self, key, value):
        if self.height <= key.y:
            needed = self.width * (key.y - self.height + 1)
            self.tower.extend(repeat(False, needed))
        self.tower[key.y * self.width + key.x] = value
    def __getitem__(self, key):
        if key.x < 0 or key.x >= self.width or key.y < 0: raise IndexError
        try: return self.tower[key.y * self.width + key.x]
        except IndexError: return False
    def get(self, key, default=None):
        try: return self[key]
        except IndexError: return default

def place_piece(tower, origin, piece):
    for offset in piece:
        tower[origin + offset] = True

def piece_fits(tower, origin, piece):
    for offset in piece:
        if tower.get(origin + offset, True): return False
    return True

with open_day(17) as f:
    conv = { '<': -1, '>': 1 }
    gusts = [conv[c] for c in f.read().rstrip()]

def solve(gusts, cycles=2022):
    tower = Tower()
    it_gusts_i = cycle(range(len(gusts)))
    gust_i = len(gusts) - 1
    states = {}
    heights = []
    for piece_i, i in zip(cycle(range(len(pieces))), range(cycles)):
        piece = pieces[piece_i]
        pos = P(2, tower.height + 3)
        for gust_i in it_gusts_i:
            gust = gusts[gust_i]
            nextpos = pos + P(gust, 0)
            if piece_fits(tower, nextpos, piece):
                pos = nextpos
            nextpos = pos + P(0, -1)
            if not piece_fits(tower, nextpos, piece):
                break
            pos = nextpos
        place_piece(tower, pos, piece)
        top_row = sum(1 << x if tower.get(P(x, tower.height-1), True) else 0 for x in range(tower.width))
        state_key = (top_row, piece_i, gust_i)
        if state_key in states: break
        states[state_key] = i, tower.height
        heights.append(tower.height)
    else: return tower.height
    base_i, base_height = states[state_key]
    cycle_len = i - base_i
    cycle_height = tower.height - base_height
    height = (cycles - base_i) // cycle_len * cycle_height
    height += heights[(cycles - 1) % cycle_len]
    return height

print(solve(gusts))
print(solve(gusts, 1000000000000))