#!/usr/bin/env python3

# import numpy as np
from functools import reduce
from re import findall
from copy import deepcopy
import sys

# filename = "in/day17.ref"
filename = "in/day17.pzl"
data = open(filename).read()
lines = [line for line in data.rstrip('\n').split('\n')]
# print(lines)

res1 = 0
res2 = 0

dxs = list()
for dx in lines[0]:
    if dx == '<':
        dxs.append(-1)
    elif dx == '>':
        dxs.append(1)
    else:
        print('error parsing input')

S = [ { (0,0):1, (1,0):1, (2,0):1, (3,0):1 },
          { (0,1):1, (1,0):1, (1,1):1, (2,1):1, (1,2):1 },
          { (0,0):1, (1,0):1, (2,0):1, (2,1):1, (2,2):1 },
          { (0,0):1, (0,1):1, (0,2):1, (0,3):1 },
          { (0,0):1, (1,0):1, (0,1):1, (1,1):1, }
]

def print_grid(yrange, G):
    for y in yrange:
        for x in range(-1, 7+1):
            if (x,y) in G:
                print('#', end = '')
            else:
                print('.', end = '')
        print()

most_y = 0
si = 0
restart = 1
dxi = 0

G = dict()
for i in range(7):
    G[(i,0)] = 1

fallen_rocks = 0

# repeat finder
rf_start = 0
rf_start_y = 0
rf_start_rocks = 0
rf_next = 0
rf_arr = list()
rf_extra_rocks = 0
rf_extra_y = 0

rocks_part1 = 2022
rocks_part2 = 1000000000000
while True:

    # create
    if restart == 1:
        x = 2
        y = most_y + 4
        restart = 0


    # move left/right
    dx = dxs[dxi% len(dxs)]
    can_go_lr = not any( [ ((x+sx+dx, y+sy) in G) or (x+sx+dx < 0) or (x+sx+dx > 6) for (sx,sy) in S[si].keys() ] )
    if can_go_lr:
        x += dx
    dxi += 1


    # repeat finder first
    if (dxi%len(dxs)) == 0 and fallen_rocks > rocks_part1 and rf_start == 0:
        # print_grid(range(most_y + 5, most_y - 20, -1), G)
        y_dist = [0 for _ in range(7)]
        for (xx,yy) in G.keys():
            y_dist[xx] = max(y_dist[xx], yy)
        y_dist = [i-most_y for i in y_dist]
        rf_arr = y_dist
        rf_start = dxi
        rf_start_y = most_y
        rf_start_rocks = fallen_rocks

    # repeat finder repeat
    if (dxi%len(dxs)) == 0 and dxi > rf_start and rf_extra_rocks == 0:

        y_dist = [0 for _ in range(7)]
        for (xx,yy) in G.keys():
            y_dist[xx] = max(y_dist[xx], yy)
        y_dist = [i-most_y for i in y_dist]

        if y_dist == rf_arr:
            rf_next = dxi
            cycle_len = rf_next - rf_start
            rocks_per_cycle = fallen_rocks - rf_start_rocks
            rocks_needed = rocks_part2 - fallen_rocks
            cycles_needed = rocks_needed // rocks_per_cycle
            rf_extra_rocks = rocks_needed - (rocks_needed % rocks_per_cycle)
            rf_extra_y = cycles_needed * (most_y - rf_start_y)


    # fall down
    can_fall_down = not any( [ (x+sx, y+sy-1) in G for (sx,sy) in S[si].keys() ] )
    if can_fall_down:
        y -= 1
    else:
        # rock stopped
        for (sx,sy) in S[si].keys():
            G[(x+sx, y+sy)] = 1
        most_y = max(most_y, y+sy)
        fallen_rocks += 1
        si = (si + 1) % 5
        restart = 1


    # part1
    if fallen_rocks == (rocks_part1-1) and restart == 1:
        max_y = 0
        for (sx,sy) in S[si].keys():
            max_y = max(max_y, y+sy)
        res1 = max_y + 1
        # break

    # part2
    if fallen_rocks + rf_extra_rocks == rocks_part2:
        max_y = 0
        for (sx,sy) in S[si].keys():
            max_y = max(max_y, y+sy)
        res2 = max_y + 1 + rf_extra_y
        break

print('res1:', res1)
print('res2:', res2)