#!/bin/python3

import re
from dijkstar import Graph, find_path
import numpy as np
from itertools import combinations
import itertools as it
from tqdm import tqdm

valves = list()

with open('inputs/day16.txt', 'r') as f:
    for line in f:
        m = re.match(r"Valve (?P<name>[A-Z]{2}) has flow rate=(?P<rate>\d+); tunnels? leads? to valves? (?P<tunnels>.*)", line.strip())
        name = m.group("name")
        rate = int(m.group("rate"))
        tunnels = m.group("tunnels").split(", ")

        valves.append({"name": name, "rate": rate, "tunnels": tunnels})

indices = {}
for i, valve in enumerate(valves):
    indices[valve["name"]] = i

graph = Graph()
for v1 in valves:
    for v2 in v1["tunnels"]:
        graph.add_edge(v1["name"], v2, 1)
        graph.add_edge(v2, v1["name"], 1)

adjs = np.zeros((len(valves), len(valves)))
for v1, v2 in combinations(map(lambda v: v["name"], valves), 2):
    cost = find_path(graph, v1, v2).total_cost
    adjs[indices[v1], indices[v2]] = cost
    adjs[indices[v2], indices[v1]] = cost

class Valve:
    def __init__(self, name, rate):
        self.name = name
        self.rate = rate

    def value(self, start, time_left):
        cost = adjs[indices[start], indices[self.name]] + 1
        return (self.rate * (time_left - cost), cost) # one off?

valve_map = dict()
valve_set = set()
for valve in valves:
    if valve["rate"] > 0:
        valve_set.add(valve["name"])
        valve_map[valve["name"]] = Valve(valve["name"], valve["rate"])

def search(valve_set, current, time_left, visited):
    to_visit_set = valve_set.difference(visited)
    best_score = 0
    for to_visit in to_visit_set:
        # print("time left:", time_left)
        # print("current:", current)
        # print("going to visit:", to_visit)
        visited_copy = visited.copy()
        visited_copy.add(to_visit)
        score, dtime = valve_map[to_visit].value(current, time_left)
        # print("+score:", score, "+time", dtime)
        if time_left - dtime >= 0:
            score = score + search(valve_set, to_visit, time_left - dtime, visited_copy)
            if not best_score or score > best_score:
                best_score = score
    return best_score


# print('part1', search(valve_set, "AA", 30, set()))

# part2 idea
# BIGBRAIN
# divide valve set in two disjoint sets
# apply same search algorithm to both sets.
# sum the scores of the sets to get total score
# brute force set distribution

def partition(pred, iterable):
    t1, t2 = it.tee(iterable)
    return it.filterfalse(pred, t1), filter(pred, t2)

def partitions(l):
    return list(filter(lambda x: [] not in x, [[[x[1] for x in f] for f in partition(lambda x: x[0], zip(pattern, l))] for pattern in it.product([True, False], repeat=len(l))]))
    # for [lhs, rhs] in parts:
    #     if (set(lhs), set(rhs)) not in result and (set(rhs), set(lhs)) not in result:
    #         result.append((set(lhs), set(rhs)))
    # return result

best_score = 0

for part in tqdm(partitions(list(valve_set))):
    [part1, part2] = part
    part1 = set(part1)
    score1 = search(part1, 'AA', 26, set())
    part2 = set(part2)
    score2 = search(part2, 'AA', 26, set())
    grand_score = score1 + score2
    if grand_score > best_score:
        best_score = grand_score

print('part2', best_score)