from typing import Dict, List
import os


def parse(input: str) -> Dict[str, List[str]]:
    lines = input.split("\n")[:-1]
    # Bidrectional adjacency list
    t = {}
    for line in lines:
        # b orbits a
        a, b = line.split(")")

        children = t.get(a, [])
        children.append(b)
        t[a] = children

        children = t.get(b, [])
        children.append(a)
        t[b] = children
    return t


def open_file(filename: str):
    # Relative to script directory
    path = os.path.join(os.path.dirname(__file__), filename)
    return open(path, "r")


def part1():
    with open_file("input1.txt") as f:
        # Tree with root node "com"
        tree = parse(f.read())
        root = "COM"
        # Calculate sum of distances between root and each node
        # BFS traversal
        q = []
        q.extend(tree[root])
        seen = set([root])
        distances = 0
        level = 1
        while len(q) > 0:
            q_next = []
            for node in q:
                if node not in seen:
                    distances += level
                    children = tree.get(node, [])
                    q_next.extend(children)
                    seen.add(node)
            level += 1
            q = q_next
        return distances


def part2():
    with open_file("input2.txt") as f:
        # Tree with root node "com"
        tree = parse(f.read())
        you = "YOU"
        santa = "SAN"
        # Use BFS to find shortest path
        # between nodes "YOU" and "SAN"
        q = []
        q.extend(tree[you])
        seen = set([you])
        dist = 1
        while len(q) > 0:
            q_next = []
            for node in q:
                if node not in seen:
                    if node == santa:
                        # Don't count santa and you
                        return dist - 2
                    children = tree.get(node, [])
                    q_next.extend(children)
                    seen.add(node)
            dist += 1
            q = q_next
        raise Exception("Didn't find Santa")


def main() -> None:
    print(part1())  # 147807
    print(part2())  # 229


if __name__ == '__main__':
    main()