from lib import *

EXAMPLE = """#.#####
#.....#
#>....#
#.....#
#...v.#
#.....#
#####.#
"""

EXAMPLE = """#.######
#>>.<^<#
#.<..<<#
#>v.><>#
#<^v^^>#
######.#
"""

B2DIR = {
    ">": (0, 1),
    "v": (1, 0),
    "<": (0, -1),
    "^": (-1, 0),
}

ADJ5 = [
    (-1, 0),
    (0, 1),
    (1, 0),
    (0, -1),
    (0, 0),
]

def solve(input: Input, second=False):
    g = input.grid2()
    start = (0, g.rows()[0].index('.'))
    poss = set([(start, 0)])
    end = (g.n_rows - 1, g.rows()[-1].index('.'))
    bzs = []
    for d in B2DIR.keys():
        for b in g.find(d):
            bzs.append((b, B2DIR[d]))

    for time in range(10**9):
        nbzs = []
        bzsset = set()
        for pos, dir in bzs:
            row, col = add2(pos, dir)
            if row == 0:
                row = g.n_rows - 2
            elif row == g.n_rows - 1:
                row = 1
            if col == 0:
                col = g.n_cols - 2
            elif col == g.n_cols - 1:
                col = 1
            nbzs.append(((row, col), dir))
            bzsset.add((row, col))
        bzs = nbzs

        nposs = set()
        while poss:
            pos, round = poss.pop()
            for adj in ADJ5:
                npos = add2(pos, adj)
                if npos[0] < 0 or npos[0] >= g.n_rows or npos[1] < 0 or npos[1] >= g.n_cols:
                    continue
                if not second:
                    if npos not in bzsset and g[npos] != '#':
                        nposs.add((npos, round))
                    if npos == end:
                        return time + 1
                else:
                    if npos == end and round == 0:
                        nposs.add((npos, round + 1))
                    elif npos == end and round == 2:
                        return time + 1
                    elif npos == start and round == 1:
                        nposs.add((npos, round + 1))
                    elif npos not in bzsset and g[npos] != '#':
                        nposs.add((npos, round))
        poss = nposs

def main():
    DAY_INPUT = "d24.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    print("Solution 1:", solve(Input(DAY_INPUT)))
    print("Example  2:", solve(Input(EXAMPLE), True))
    print("Solution 2:", solve(Input(DAY_INPUT), True))

if __name__ == "__main__":
    main()