from utils import open_day
from functools import cache
from dataclasses import dataclass
from itertools import count, chain
import re

@dataclass
class Node:
    flow: int
    neighbours: list[tuple[str, int]]

ids = {'AA': 0}
nextid = 1
def intern(n):
    global nextid
    if n in ids: return ids[n]
    ids[n] = nextid
    nextid += 1
    return nextid - 1

regex = re.compile(r'^Valve (..) has flow rate=([0-9]+); tunnels? leads? to valves? (.*)$')
inp = {}
with open_day(16) as f:
    for line in f:
        m = regex.match(line)
        assert(m)
        valve, flow, neighbours = m.group(1, 2, 3)
        inp[valve] = (int(flow), neighbours.split(', '))

nodes = {}
for valve, (flow, neighbours) in inp.items():
    if valve != 'AA' and flow == 0: continue
    actual_neighbours = []
    for n in neighbours:
        prev = valve
        for cost in count(1):
            if n == 'AA' or inp[n][0] != 0: break
            l, r = inp[n][1]
            nnext = r if l == prev else l
            prev = n
            n = nnext
        actual_neighbours.append((intern(n), cost))
    nodes[intern(valve)] = Node(flow, actual_neighbours)

@cache
def sum_flow(open_valves):
    total = 0
    for i in count():
        if not open_valves: break
        if open_valves & 1: total += nodes[i].flow
        open_valves >>= 1
    return total

recurse_cache = dict()
def recurse(open_valves=1, current=0, flow=0, time_left=30):
    key = (current, flow, time_left)
    if key in recurse_cache: return recurse_cache[key]
    cnode = nodes[current]
    best = flow
    def check(open_valves, current, flow, time_left):
        nonlocal best
        new = recurse(open_valves, current, flow, time_left)
        if new > best: best = new
    for neighbour, cost in cnode.neighbours:
        if cost >= time_left: continue
        check(open_valves, neighbour, flow, time_left - cost)
    if not (open_valves & (1 << current)) and time_left > 0:
        check(open_valves | 1 << current, current, flow + cnode.flow * (time_left - 1), time_left - 1)
    recurse_cache[key] = best
    return best

print(recurse())