from lib import get_data
from lib import Grid2D

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

data = get_data(__file__)
data, moves = data.split("\n\n")

g = Grid2D(data)


def simulate(grid: Grid2D, moves: str):
    (p,) = g.find("@")
    for c in moves:
        if c == "\n":
            continue
        d = DIRS[c]

        to_move = set([p])
        row = to_move
        blocked = False
        while True:
            nrow = set()
            for (
                r,
                c,
            ) in row:
                np = r + d[0], c + d[1]
                if g[np] == "#":
                    blocked = True
                elif g[np] == "O":
                    nrow.add(np)
                elif g[np] == "[":
                    nrow.add(np)
                    if d == (1, 0) or d == (-1, 0):
                        nrow.add((np[0], np[1] + 1))
                elif g[np] == "]":
                    if d == (1, 0) or d == (-1, 0):
                        nrow.add((np[0], np[1] - 1))
                    nrow.add(np)
            to_move |= row
            row = nrow
            if len(row) == 0 or blocked:
                break

        if not blocked:
            to_place = {}
            p = p[0] + d[0], p[1] + d[1]

            for tm in to_move:
                np = tm[0] + d[0], tm[1] + d[1]
                to_place[np] = g[tm]
                g[tm] = "."

            for pos, c in to_place.items():
                g[pos] = c


def score(g: Grid2D) -> int:
    os = g.find("O")
    if os:
        return sum(r * 100 + c for r, c in os)
    return sum(r * 100 + c for r, c in g.find("["))


simulate(g, moves)
print(score(g))

ndata = ""
for c in data:
    if c == "#":
        ndata += "##"
    elif c == "O":
        ndata += "[]"
    elif c == ".":
        ndata += ".."
    elif c == "@":
        ndata += "@."
    else:
        ndata += c

g = Grid2D(ndata)
(p,) = g.find("@")

simulate(g, moves)
print(score(g))
exit()