import shared
from scanf import scanf
from pprint import pprint
from math import lcm


def setup(rows):
    MAP = {}
    instructions = rows.pop(0)
    rows.pop(0)
    instructions = instructions.replace("L", "0")
    instructions = instructions.replace("R", "1")
    instructions = [int(x) for x in instructions]

    for row in rows:
        entry, left, right = scanf("%s = (%s, %s)", row)
        MAP[entry] = [left, right]
    return MAP, instructions


# @shared.profile
def part1(rows):
    _map, _instructions = setup(rows)
    count, _ = find("AAA", _instructions, _map)
    print(count)


# @shared.profile
def part2(rows):
    _map, _instructions = setup(rows)
    starting_points = [k for k in _map.keys() if k[-1] == "A"]
    ending_points = [k for k in _map.keys() if k[-1] == "Z"]
    ends = []
    for s in starting_points:
        count, end = find(s, _instructions, _map, None)
        ends.append(count)
    print(lcm(*ends))


def find(start, instructions, MAP, target="ZZZ"):
    count = 0
    current = start
    while True:
        idx = count % len(instructions)
        l_r = instructions[idx]
        current = MAP[current][l_r]
        count += 1
        if target is None:
            if current.endswith("Z"):
                return count, current
        elif current == target:
            return count, current


def main():
    rows = [row for row in shared.load_rows(8)]
    with shared.elapsed_timer() as elapsed:
        part1(rows[:])
        print("🕒", elapsed())

    rows = [row for row in shared.load_rows(8, True)]
    with shared.elapsed_timer() as elapsed:
        part2(rows[:])
        print("🕒", elapsed())


if __name__ == "__main__":
    main()