from tools.aoc import AOCDay
from typing import Any


class Program:
    def __init__(self, name: str, weight: int):
        self.name = name
        self.weight = weight
        self.children = set()

    def add(self, other: 'Program'):
        if other not in self.children:
            self.children.add(other)

    def get_weight_total(self):
        return self.weight + sum([x.get_weight_total() for x in self.children])


def get_weight(root: Program, diff: int = 0) -> int:
    weights = {}
    for child in root.children:
        child_weight_total = child.get_weight_total()
        if child_weight_total not in weights:
            weights[child_weight_total] = []
        weights[child_weight_total].append(child)

    present_weights = list(weights.keys())
    if len(present_weights) == 1:
        return root.weight + diff
    else:
        if len(weights[present_weights[0]]) == 1:  # 0 is the problem
            return get_weight(weights[present_weights[0]][0], present_weights[1] - present_weights[0])
        else:  # 1 is the problem
            return get_weight(weights[present_weights[1]][0], present_weights[0] - present_weights[1])


class Day(AOCDay):
    inputs = [
        [
            ("tknk", "input7_test"),
            ("aapssr", "input7")
        ],
        [
            (60, "input7_test"),
            (1458, "input7")
        ]
    ]

    def parse_input(self) -> Program:
        child_cache = {}
        progs = {}
        prog_set = set()
        for line in self.getInput():
            if "->" not in line:
                prog_name, weight_string = line.split(" ")
            else:
                prog_string, children = line.split(" -> ")
                prog_name, weight_string = prog_string.split(" ")
                child_cache[prog_name] = children.split(", ")

            prog_set.add(prog_name)
            progs[prog_name] = Program(prog_name, int(weight_string[1:-1]))

        for prog, childs in child_cache.items():
            for child in childs:
                prog_set.remove(child)
                progs[prog].add(progs[child])

        return progs[prog_set.pop()]

    def part1(self) -> Any:
        root = self.parse_input()
        return root.name

    def part2(self) -> Any:
        root = self.parse_input()
        return get_weight(root)


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