from lib import *

EXAMPLE = """>>><<><>><<<>><>>><<<>>><<<><<<>><>><<>>"""

ROCKS = [
    [(0, 0), (0, 1), (0, 2), (0, 3)],
    [(0, 1), (-1, 0), (-1, 1), (-1, 2), (-2, 1)],
    [(0, 0), (0, 1), (0, 2), (-1, 2), (-2, 2)],
    [(0, 0), (-1, 0), (-2, 0), (-3, 0)],
    [(0, 0), (0, 1), (-1, 0), (-1, 1)],
]

def printfall(coords):
    row_coords = list(map(lambda c: c[0], coords))
    col_coords = list(map(lambda c: c[1], coords))
    min_row, max_row = min(row_coords), max(row_coords)
    min_col, max_col = min(col_coords), max(col_coords)

    s = ""
    for r in range(min_row, max_row + 1):
        for c in range(min_col, max_col + 1):
            s += "."
        s += "\n"
    g = Input(s).grid2()
    for c in coords:
        c = (c[0] + -min_row, c[1])
        g[tuple(c)] = "#"
    g.print()
    print()

def solve(input: Input, second=False):
    BLOWS = list(input.text.strip())
    width, blow_index = 7, 0

    heighest = [0 for _ in range(width)]
    rocks = set([(0, i) for i in range(width)])

    if second:
        target = 1000000000000
        nhs = {}
    else:
        target = 2022

    i = 0
    while i < target:
        rock = ROCKS[i % len(ROCKS)]
        pos = (min(heighest) - 4, 2)

        if second and i > 0 and all(abs(heighest[i] - heighest[i + 1]) <= 2 for i in range(len(heighest) - 1)):
            # long winded code to find repeating sequence
            nh = tuple([h - min(heighest) for h in heighest])
            if nh in nhs:
                nhs[nh].append((i, heighest[0]))
            else:
                nhs[nh] = [(i, heighest[0])]
            xs = nhs[nh]
            if len(xs) > 4 and all(xs[1][0] - xs[0][0] == xs[i + 1][0] - xs[i][0] for i in range(len(xs) - 1)):
                # we found a repeating sequence and can skip forward
                di = xs[-1][0] - xs[-2][0]
                dh = xs[-1][1] - xs[-2][1]
                repeat = (target - i) // di
                i += repeat * di
                heighest = [h + repeat * dh for h in heighest]
                for c, r in enumerate(heighest):
                    rocks.add((r, c))
                nhs = {}
                continue

        falling = True
        while falling:
            blow = BLOWS[blow_index % len(BLOWS)]
            blow_index += 1

            dc = None
            if blow == ">":
                dc = 1
            elif blow == "<":
                dc = -1

            for (ro, co) in rock:
                r = pos[0] + ro
                c = pos[1] + co + dc
                if c < 0 or c >= width or (r, c) in rocks:
                    break
            else:
                pos = (pos[0], pos[1] + dc)

            for (ro, co) in rock:
                r = pos[0] + ro + 1
                c = pos[1] + co
                if (r, c) in rocks:
                    falling = False
                    break
            else:
                pos = (pos[0] + 1, pos[1])

        for (ro, co) in rock:
            r = pos[0] + ro
            c = pos[1] + co
            rocks.add((r, c))
            heighest[c] = min(heighest[c], r)

        i += 1

    return -min(heighest)

def main():
    _, day = extract_year_and_date(__file__)
    DAY_INPUT = f"i{day}.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    s1 = solve(Input(DAY_INPUT))
    assert s1 == 3102
    print("Solution 1:", s1)
    print("Example  2:", solve(Input(EXAMPLE), True))
    print("Solution 2:", solve(Input(DAY_INPUT), True))
    assert solve(Input(DAY_INPUT), True) == 1539823008825

if __name__ == "__main__":
    main()