import aocd
from enum import Enum, nonmember

class Direction(Enum):
    NORTH = (0, -1)
    WEST = (-1, 0)
    SOUTH = (0, 1)
    EAST = (1, 0)

    @nonmember
    def clockwise_from(d):
        return list(Direction)[list(Direction).index(d) - 1]

def facing(coords, direction):
    return tuple(map(sum, zip(coords, direction.value)))

def parse(data):
    """
    Given maze data, return
    (Tuple of guard coords, set of obstacle coords, grid width, grid height)
    """
    guard = None
    obstacles = set()

    for y, row in enumerate(data.split('\n')):
        for x, char in enumerate(row):
            if char == '#':
                obstacles.add((x, y))
            elif char == '^':
                guard = (x, y)
    
    return (guard, obstacles, x, y)

def solve(data):
    guard, obstacles, width, height = parse(data)
    guard_direction = Direction.NORTH

    visited = set()

    while (0 <= guard[0] <= width) and (0 <= guard[1] <= height):
        visited.add(guard)
        while (facing_tile := facing(guard, guard_direction)) in obstacles:
            guard_direction = Direction.clockwise_from(guard_direction)
        guard = facing_tile

    return sum(1 for _ in visited)


if __name__ == '__main__':
    puz = aocd.models.Puzzle(year=2024, day=6)

    # data = puz.examples[0].input_data
    data = puz.input_data

    solution = solve(data)

    print(solution)
    aocd.submit(solution, year=2024, day=6, part='a')