from lib import *

EXAMPLE = """        0..#
        .#..
        #...
        ....
1..#2...3..#
........#...
..#....#....
..........#.
        4..#5...
        .....#..
        .#......
        ......#.

10R5L5R10L4R5L5"""


E = (0, 1)
S = (1, 0)
W = (0, -1)
N = (-1, 0)

DIRS = [
    E,
    S,
    W,
    N,
]

# XXX: not hands free - i created this manually probably there is some way to
# figure it out from the input, but id don't know how right now.
trans_example = {
    (0, N): (1, S),
    (0, E): (5, W),
    (0, S): (3, S),
    (0, W): (2, S),
    (1, N): (0, S),
    (1, E): (2, E),
    (1, S): (4, N),
    (1, W): (5, N),
    (2, N): (0, E),
    (2, E): (3, E),
    (2, S): (4, E),
    (2, W): (1, W),
    (3, N): (0, N),
    (3, E): (5, S),
    (3, S): (4, S),
    (3, W): (2, W),
    (4, N): (3, N),
    (4, E): (5, E),
    (4, S): (1, N),
    (4, W): (2, N),
    (5, N): (3, W),
    (5, E): (0, W),
    (5, S): (1, E),
    (5, W): (4, W),
}

trans_input = {
    (0, N): (5, E),
    (0, E): (1, E),
    (0, S): (2, S),
    (0, W): (3, E),
    (1, N): (5, N),
    (1, E): (4, W),
    (1, S): (2, W),
    (1, W): (0, W),
    (2, N): (0, N),
    (2, E): (1, N),
    (2, S): (4, S),
    (2, W): (3, S),
    (3, N): (2, E),
    (3, E): (4, E),
    (3, S): (5, S),
    (3, W): (0, E),
    (4, N): (2, N),
    (4, E): (1, W),
    (4, S): (5, W),
    (4, W): (3, W),
    (5, N): (3, N),
    (5, E): (4, N),
    (5, S): (1, S),
    (5, W): (0, S),
}

def parse_steps(steps):
    insts = re.compile(r"(\d+)([LR])").findall(steps)
    lastn = re.compile(r"\d+").findall(steps)[-1]
    insts.append((lastn, None))
    return insts


def solve(input: Input):
    grid, steps = input.text.split("\n\n")
    insts = parse_steps(steps)

    grid = list(map(list, grid.splitlines()))
    rows = len(grid)

    startcol = 0
    while grid[0][startcol] == ' ':
        startcol += 1
    pos = (0, startcol)
    dir = (0, 1)
    for n, turn in insts:
        nrow, ncol = pos
        for _ in range(int(n)):
            nrow = (nrow + dir[0]) % rows
            ncol = (ncol + dir[1]) % len(grid[nrow])
            while grid[nrow][ncol] == ' ':
                nrow = (nrow + dir[0]) % rows
                ncol = (ncol + dir[1]) % len(grid[nrow])
            if grid[nrow][ncol] == '#':
                break
            pos = (nrow, ncol)
        if turn is not None:
            dir = DIRS[(DIRS.index(dir) + 1 if turn == 'R' else DIRS.index(dir) - 1) % 4]

    return 1000 * (pos[0] + 1) + 4 * (pos[1] + 1) + DIRS.index(dir)


def solve2(input: Input, square_size=4):
    squares = []
    grid, steps = input.text.split("\n\n")
    insts = parse_steps(steps)
    grid_rows = grid.splitlines()
    for r in range(0, len(grid_rows), square_size):
        for c in range(0, len(grid_rows[r]), square_size):
            square = []
            for sr in range(r, r + square_size):
                square.append(grid_rows[sr][c:c+square_size])
            if square[0][0] != ' ':
                squares.append(square)

    def flip(n):
        return square_size - 1 - n

    def teleport(dir, sqi, row, col):
        if square_size == 4:
            nsqi, ndir = trans_example[(sqi, dir)]
        elif square_size == 50:
            nsqi, ndir = trans_input[(sqi, dir)]
        else:
            assert False
        nrow, ncol = None, None

        if ndir == N:
            nrow = square_size - 1
        elif ndir == S:
            nrow = 0
        elif ndir == E:
            ncol = 0
        elif ndir == W:
            ncol = square_size - 1

        if dir == N:
            if ndir == N:
                ncol = col
            elif ndir == S:
                ncol = flip(col)
            elif ndir == E:
                nrow = col
            elif ndir == W:
                nrow = square_size - 1 - col
        elif dir == S:
            if ndir == N:
                ncol = flip(col)
            elif ndir == S:
                ncol = col
            elif ndir == E:
                nrow = square_size - 1 - col
            elif ndir == W:
                nrow = col
        elif dir == E:
            if ndir == N:
                ncol = row
            elif ndir == S:
                ncol = flip(row)
            elif ndir == E:
                nrow = row
            elif ndir == W:
                nrow = square_size - 1 - row
        elif dir == W:
            if ndir == N:
                ncol = flip(row)
            elif ndir == S:
                ncol = row
            elif ndir == E:
                nrow = flip(row)
            elif ndir == W:
                nrow = row

        assert nrow is not None
        assert ncol is not None
        return ndir, nsqi, nrow, ncol

    # for square in squares:
    #     print("\n".join(square), "\n")

    dir, sqi, row, col = E, 0, 0, 0
    for n, turn in insts:
        # print(f"{sqi=} {row=} {col=} {dir=}")
        for _ in range(int(n)):
            nrow = row + dir[0]
            ncol = col + dir[1]
            if nrow < 0 or nrow >= square_size or ncol < 0 or ncol >= square_size:
                ndir, nsqi, nrow, ncol = teleport(dir, sqi, nrow, ncol)
            else:
                ndir, nsqi = dir, sqi
            if squares[nsqi][nrow][ncol] != '#':
                dir, sqi, row, col = ndir, nsqi, nrow, ncol
        if turn is not None:
            dir = DIRS[(DIRS.index(dir) + 1 if turn == 'R' else DIRS.index(dir) - 1) % 4]
    print(f"{sqi=} {row=} {col=} {dir=}")

    if square_size == 4:
        # sqi=2 row=0 col=2 dir=(-1, 0)
        return 1000 * (4 + 0 + 1) + 4 * (4 + 2 + 1) + DIRS.index((-1, 0))
    elif square_size == 50:
        # sqi=1 row=11 col=14 dir=(0, -1)
        return 1000 * (0 + 11 + 1) + 4 * (100 + 14 + 1) + DIRS.index((0, -1))
    else:
        assert False

def main():
    DAY_INPUT = "d22.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    print("Solution 1:", solve(Input(DAY_INPUT)))
    assert solve(Input(DAY_INPUT)) == 149250

    print("Example  2:", solve2(Input(EXAMPLE)))
    print("Solution 2:", solve2(Input("d22.txt"), 50))
    return

if __name__ == "__main__":
    main()