from lib import *
from math import lcm
from collections import deque

EXAMPLE = """
broadcaster -> a, b, c
%a -> b
%b -> c
%c -> inv
&inv -> a
"""

EXAMPLE2 = """
broadcaster -> a
%a -> inv, con
&inv -> b
%b -> con
&con -> output
"""

def visualize(modules):
    with open("g20.dot", 'w') as f:
        f.write("digraph G {\n")
        for m in modules.values():
            for cm in m[3]:
                f.write(" " + m[1] + ' -> ' + cm + "\n")
        f.write("}")

def solve(input: Input, second=False):
    modules = {}
    for line in input.lines():
        if not line: continue
        src, dsts = line.split(" -> ")
        dsts = dsts.split(", ")
        if src.startswith("%"):  # flip-flop
            modules[src[1:]] = [src[0], src[1:], 0, dsts]
        elif src.startswith("&"):  # conjunction
            modules[src[1:]] = [src[0], src[1:], {}, dsts]
        elif src.startswith("broadcaster"):
            modules[src] = ["b", src, 0, dsts]
        else:
            raise Exception()

    for m in modules.values():
        for d in m[3]:
            if d in modules and modules[d][0] == "&":
                modules[d][2][m[1]] = 0

    if second:
        # visualize(modules)
        (feed,) = [m[1] for m in modules.values() if "rx" in m[3]]
        periods = {d: [] for d in modules[feed][2].keys()}
        BUTTON_PUSHES = 10000
    else:
        BUTTON_PUSHES = 1000

    lo, hi = 0, 0
    for i in range(BUTTON_PUSHES):
        qs = deque([("button module", "broadcaster", 0)])
        while qs:
            src, dst, sig = qs.popleft()

            if sig == 0:
                lo += 1
            else:
                hi += 1

            if not dst in modules:
                continue

            m = modules[dst]
            new_pulse = None
            if m[0] == "b":
                new_pulse = 0
            elif m[0] == "%":
                if sig == 0:
                    m[2] = (m[2] + 1) % 2
                    new_pulse = m[2]
            elif m[0] == "&":
                m[2][src] = sig
                if all(s == 1 for s in m[2].values()):
                    new_pulse = 0
                else:
                    new_pulse = 1
            else:
                raise Exception()

            if new_pulse is not None:
                for nxtdst in m[3]:
                    if second and nxtdst == feed and new_pulse == 1:
                        # print(f"{i:>4}: {dst[:4]:>4} -{new_pulse}> {nxtdst}")
                        periods[dst].append(i)
                    qs.append((dst, nxtdst, new_pulse))
    if second:
        # print(periods)
        # Yeah, we got a bit lucky that this works, I guess, but it does.
        return lcm(*[v[1] - v[0] for v in periods.values()])
    return lo * hi

def main():
    DAY_INPUT = "i20.txt"
    print("Example  1:", solve(Input(EXAMPLE)))
    print("Example  2:", solve(Input(EXAMPLE2)))
    print("Solution 1:", solve(Input(DAY_INPUT)))
    print("Solution 2:", solve(Input(DAY_INPUT), True))
    assert solve(Input(DAY_INPUT), True) == 244178746156661

if __name__ == "__main__":
    main()