from tools.aoc import AOCDay
from typing import Any


class Programm:
    def __init__(self, pid: int):
        self.pid = pid
        self.children = set()


class Day(AOCDay):
    inputs = [
        [
            (6, "input12_test"),
            (145, "input12")
        ],
        [
            (207, "input12")
        ]
    ]

    def build_programm_tree(self) -> dict:
        ret = {}
        for line in self.getInput():
            pid, peer_str = line.split(" <-> ")
            pid = int(pid)
            if pid not in ret:
                ret[pid] = Programm(pid)
            peers = list(map(int, peer_str.split(", ")))
            for x in peers:
                if x not in ret:
                    ret[x] = Programm(x)
                ret[pid].children.add(ret[x])
                ret[x].children.add(ret[pid])

        return ret

    def count_group_zero(self, programs: dict, start: int = 0, counted: set = None) -> (int, set):
        count = 0
        if counted is None:
            counted = set()
        counted.add(start)
        for p in  programs[start].children:
            if p.pid not in counted:
                count += self.count_group_zero(programs, p.pid, counted)[0]

        return count + 1, counted

    def part1(self) -> Any:
        return self.count_group_zero(self.build_programm_tree(), 0)[0]

    def part2(self) -> Any:
        p = self.build_programm_tree()
        g = 0
        c = set()
        for x in p:
            if x in c:
                continue
            c = c.union(self.count_group_zero(p, x)[1])
            g += 1
        return g


if __name__ == '__main__':
    day = Day(2017, 12)
    day.run(verbose=True)