import re
from itertools import cycle

def parse_input(infile):
    with open(infile) as f:
        content = f.read().rstrip()
        directions, nodes = content.split("\n\n")
        directions = directions.strip()
        nodes = nodes.split("\n")
        nodes = [n.split(" = ") for n in nodes]
        nodes = {k: re.findall(r"\w{3}", v) for k, v in nodes}
        return directions, nodes


def part1(directions, nodes):
    iterations = 0
    current_node = "AAA"
    for d in cycle(directions):
        if current_node == "ZZZ":
            break
        iterations += 1
        if d == "L":
            current_node = nodes[current_node][0]
        else:
            current_node = nodes[current_node][1]
    print(f"Reached 'ZZZ' in {iterations} iterations")


def part2(directions, nodes):
    current_nodes = [k for k in nodes.keys() if k.endswith("A")]
    iterations = [0] * len(current_nodes)
    print(current_nodes, iterations)

    for d in cycle(directions):
        if all(c.endswith("Z") for c in current_nodes):
            break
        #iterations += len(current_nodes)
        if d == "L":
            new_nodes = []
            for i, n in enumerate(current_nodes):
                if not n.endswith("Z"):
                    new_nodes.append(nodes[n][0])
                    iterations[i] += 1
                else:
                    new_nodes.append(n)
            current_nodes = new_nodes
            #current_nodes = [c if c.endswith("Z") else nodes[c][0] for c in current_nodes]
        else:
            new_nodes = []
            for i, n in enumerate(current_nodes):
                if not n.endswith("Z"):
                    new_nodes.append(nodes[n][1])
                    iterations[i] += 1
                else:
                    new_nodes.append(n)
            current_nodes = new_nodes
            #current_nodes = [c if c.endswith("Z") else nodes[c][1] for c in current_nodes]
        print(current_nodes, iterations)
        
    
    print(f"Reached all nodes such that 'xxZ' in {iterations} iterations")


if __name__ == "__main__":
    import sys
    import os
    SCRIPTPATH = os.path.dirname(os.path.realpath(__file__))
    infile = sys.argv[1] if len(sys.argv) == 2 else "example.txt"

    directions, nodes = parse_input(os.path.join(SCRIPTPATH, infile))
    part1(directions, nodes)

    directions, nodes = parse_input(os.path.join(SCRIPTPATH, infile))
    part2(directions, nodes)