#!/usr/bin/env python3

import sys
from copy import deepcopy
import numpy as np


def show(maps):
    st = ''
    for line in maps:
        print(''.join(line))


def find_guard(current):
    move = (0,0)
    for y in range(0, len(current)):
        for x, field in enumerate(current[y]):
            if field == '.' or field == '#' or field == 'X': continue
            elif field == '^': move = np.array((0, -1))
            elif field == '>': move = np.array((1, 0))
            elif field == 'v': move = np.array((0, 1))
            elif field == '<': move = np.array((-1, 0))
            guard = np.array((x, y))
            return(guard, move, field)


def guard(m, guard, move):
    next = {(0,-1): '>', (1,0): 'v', (0,1): '<', (-1,0): '^'}
    nextm = {(0,-1): (1,0), (1,0): (0,1), (0,1): (-1, 0), (-1,0): (0,-1)}
    current = m

    while (0 <= guard[0] + move[0] < len(m[0]) and 0 <= guard[1] + move[1] < len(m)):
        gx, gy = guard
        guard = guard + move
        x, y = guard
        if m[y][x] == '#':
            current[gy][gx] = next[tuple(move)]
            return current, True, np.array((gx, gy)), np.array(nextm[tuple(move)])
        else:
            current[gy][gx] = 'X'
    current[guard[1]][guard[0]] = 'X'
    return current, False, np.array((guard[0], guard[1])), np.array(nextm[tuple(move)])


def test_loop(m, x, y, g, move, o):
    current = deepcopy(m)
    current[y][x] = '#'
    positions = set([])
    positions.add((g[0], g[1], o))
    running = True
    while running:
        current, running, g, move = guard(current, g, move)
        o = current[g[1]][g[0]]
        if (g[0], g[1], o) in positions: 
            return True
        positions.add((g[0], g[1], o))
    return False


if __name__ == '__main__':
    maps = [list(line.strip('\n')) for line in open(sys.argv[1])]
    sguard, smove, start = find_guard(maps)

    # challenge 1
    running = True
    current = deepcopy(maps)
    g, m = sguard, smove
    while running:
        current, running, g, m = guard(current, g, m)
    
    res1 = sum([m.count('X') for m in current])
    print(f"challenge 1:\n{res1}\n")

    # challenge 2
    res2 = 0
    for y in range(0, len(current)):
        for x, field in enumerate(current[y]):
            if field == '#' or field == '.': continue
            res2 += test_loop(maps, x, y, sguard, smove, start)
    print(f"challenge 2:\n{res2}")