import lib

EXAMPLE = """
Valve AA has flow rate=0; tunnels lead to valves DD, II, BB
Valve BB has flow rate=13; tunnels lead to valves CC, AA
Valve CC has flow rate=2; tunnels lead to valves DD, BB
Valve DD has flow rate=20; tunnels lead to valves CC, AA, EE
Valve EE has flow rate=3; tunnels lead to valves FF, DD
Valve FF has flow rate=0; tunnels lead to valves EE, GG
Valve GG has flow rate=0; tunnels lead to valves FF, HH
Valve HH has flow rate=22; tunnel leads to valve GG
Valve II has flow rate=0; tunnels lead to valves AA, JJ
Valve JJ has flow rate=21; tunnel leads to valve II
"""

def solve(lines: list[str]):
    nodes = {}
    for (i, line) in enumerate(lines):
        source = line.split(" ")[1]
        to_valves = list(map(lambda v: v.replace(",", ""), line.split(" ")[9:]))
        flow_rate = lib.str_to_int(line)
        nodes[source] = (flow_rate, to_valves)


    options = [(0, 0, "AA", ())]
    visited = set()
    for _ in range(30):
        new_options = []
        for (total, flow, current, valves) in options:
            valve_flow_rate = nodes[current][0]
            if not current in valves and valve_flow_rate > 0:
                o = (total + flow, flow + valve_flow_rate, current, valves + (current, ))
                new_options.append(o)
            for new_valve in nodes[current][1]:
                o = (total + flow, flow, new_valve, valves)
                new_options.append(o)

        options = sorted(list(set(new_options)))[-100:]
    return options[-1][0]


def solve2(lines: list[str]):
    nodes = {}
    for (i, line) in enumerate(lines):
        source = line.split(" ")[1]
        to_valves = list(map(lambda v: v.replace(",", ""), line.split(" ")[9:]))
        flow_rate = lib.str_to_int(line)
        nodes[source] = (flow_rate, to_valves)

    options = [(0, 0, ("AA", "AA"), ())]
    for _ in range(26):
        new_options = []
        for (total, flow, (a, b), valves) in options:
            flow_a = nodes[a][0]
            flow_b = nodes[b][0]

            # a turns on b moves
            if not a in valves and flow_a > 0:
                for new_valve_b in nodes[b][1]:
                    o = (total + flow, flow + flow_a, (a, new_valve_b), valves + (a, ))
                    new_options.append(o)

            # a moves and b turns on
            if not b in valves and flow_b > 0:
                for new_valve_a in nodes[a][1]:
                    o = (total + flow, flow + flow_b, (new_valve_a, b), valves + (b, ))
                    new_options.append(o)

            # both turn on
            if a != b and flow_a > 0 and flow_b > 0 and not a in valves and not b in valves:
                o = (total + flow, flow + flow_a + flow_b, (a, b), valves + (a, b, ))
                new_options.append(o)

            # both move
            for new_valve_a in nodes[a][1]:
                for new_valve_b in nodes[b][1]:
                    o = (total + flow, flow, (new_valve_a, new_valve_b), valves)
                    new_options.append(o)

        options = sorted(list(set(new_options)))[-1000:]
    # 52:00
    return options[-1][0]


def main():
    lines = lib.str_to_lines_no_empty(EXAMPLE)
    print("Example  1:", solve(lines))

    lines = lib.str_to_lines_no_empty(open("i16.txt").read())
    print("Solution 1:", solve(lines))

    lines = lib.str_to_lines_no_empty(EXAMPLE)
    print("Example  2:", solve2(lines))

    lines = lib.str_to_lines_no_empty(open("i16.txt").read())
    print("Solution 2:", solve2(lines))
    assert solve2(lines) == 2591

if __name__ == "__main__":
    main()