Untitled

from advent_of_code.lib import parse as aoc_parse
from advent_of_code.lib import aoc
from dataclasses import dataclass, field
from collections import defaultdict, deque
from abc import ABC


@dataclass
class Message:
    sender: str
    user: str
    signal: int


@dataclass
class Module(ABC):
    name: str
    links: list[str] = field(default_factory=list)

    def receive(self, sender, signal, q):
        pass

    def send(self, signal, q):
        for link in self.links:
            q.append(Message(self.name, link, signal))


@dataclass
class FlipFlop(Module):
    value: int = 0

    def receive(self, sender, signal, q):
        if signal == 0:
            self.value = 1 - self.value
            self.send(self.value, q)


class Conjunction(Module):
    def set_pre(self, pre):
        self.pre = {p: 0 for p in pre}

    def receive(self, sender, signal, q):
        self.pre[sender] = signal
        self.send(0 if all(self.pre.values()) else 1, q)


class Broadcaster(Module):
    def receive(self, sender, signal, q):
        pass


class Output(Module):
    def send(self, signal, q):
        pass


def get_module(config):
    name, links = config.split(" -> ")
    links = list(links.split(", "))
    if name == "broadcaster":
        return Broadcaster(name, links)
    if name[0] == "%":
        return FlipFlop(name[1:], links)
    if name[0] == "&":
        return Conjunction(name[1:], links)


def get_input(file):
    raw = aoc.read_input(2023, 20, file)
    modules = {}
    for line in aoc_parse.as_lines(raw):
        module = get_module(line)
        modules[module.name] = module
    # compute predecessors of conjuctions
    pre = defaultdict(list)
    for name, module in modules.items():
        for link in module.links:
            pre[link].append(name)
    for name, module in modules.items():
        if isinstance(module, Conjunction):
            module.set_pre(pre[name])
    # add output nodes
    for module in list(modules.values()):
        for link in module.links:
            if link not in modules:
                modules[link] = Output(link)
    return modules


@aoc.pretty_solution(1)
def part1(modules):
    low_pulse, high_pulse = 0, 0
    for _ in range(1000):
        # manually handle the queue of messages
        q = deque()
        modules["broadcaster"].send(0, q)
        while q:
            msg = q.popleft()
            low_pulse += msg.signal == 0
            high_pulse += msg.signal == 1
            # process this message
            # the module that receives it will update the q
            modules[msg.user].receive(msg.sender, msg.signal, q)
    return (low_pulse + 1000) * high_pulse


if __name__ == "__main__":
    data = get_input("input.txt")
    part1(data)